                                             
import os
import sys
import logging
from typing import Dict, Any, List, Tuple, Optional

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from fortress.data_management.prompt_processor import PromptProcessor
from fortress.core.vector_store_interface import VectorStoreInterface
from fortress.common.data_models import QueryFeatures, DetectionResult, DecisionLabel                                      
from fortress.config import get_config

logger = logging.getLogger(__name__)


class PrimaryDetector:
    """
    Handles the primary detection phase by processing an input query,
    generating its embedding, querying a vector store for similar prompts,
    and applying an ensemble strategy to the results.
    """


    def __init__(self, prompt_processor: PromptProcessor, vector_store: VectorStoreInterface):
        """
        Initializes the PrimaryDetector.

        Args:
            prompt_processor: An instance of PromptProcessor to process queries.
            vector_store: An instance of a VectorStoreInterface for similarity search.
        """
        self.prompt_processor = prompt_processor
        self.vector_store = vector_store
        self.config = get_config()
        self.detection_config = self.config.get('detection_pipeline', {})
        self.top_k = self.detection_config.get('top_k_semantic_search', 10)
        self.ensemble_strategy = self.detection_config.get('ensemble_strategy', 'closest_match')
        self.ensemble_options = self.detection_config.get('ensemble_options', {})

        logger.info(f"PrimaryDetector initialized. Will retrieve top {self.top_k} results. Ensemble strategy: {self.ensemble_strategy}")

    def detect_batch(self, query_texts: List[str]) -> List[Tuple[Optional[QueryFeatures], Optional[List[float]], List[Dict[str, Any]]]]:
        logger.debug(f"PrimaryDetector processing batch of {len(query_texts)} queries...")
        
                                                                                  
                                                                               
                                                                                                        
                                                       
                                                                                    
                                                 
        processed_queries_data = self.prompt_processor.process_for_query_batch(query_texts)

        results_batch = []
        for i, (query_features, query_embedding, error_msg) in enumerate(processed_queries_data):
            if error_msg or not query_embedding or not query_features:
                logger.warning(f"Failed to generate embedding or features for query text: {query_texts[i][:50]}... Error: {error_msg}")
                results_batch.append((None, None, []))                                          
                continue

                                                                                                           
                                                   
            similar_docs = self.vector_store.query_similar(
                embedding=query_embedding,
                top_k=self.top_k
            )
            logger.info(f"PrimaryDetector found {len(similar_docs)} similar documents for query item {i}.")
            results_batch.append((query_features, query_embedding, similar_docs))
            
        return results_batch
    def detect(self, query_text: str) -> Tuple[Optional[QueryFeatures], Optional[List[float]], List[Dict[str, Any]]]:
        """
        Processes an input query to generate features, an embedding, and find similar documents.

        Args:
            query_text: The input query string.

        Returns:
            A tuple containing:
                - Optional[QueryFeatures]: Extracted features for the query (including token log probabilities).
                - Optional[List[float]]: The embedding of the query text.
                - List[Dict[str, Any]]: A list of similar documents from the vector store.
        """
        logger.debug(f"PrimaryDetector processing query: {query_text[:100]}...")
        
                                                                                            
        query_features, query_embedding, _ = self.prompt_processor.process_for_query(query_text)

        if not query_embedding or not query_features:
            logger.warning("Failed to generate embedding or features for query in PrimaryDetector.")
            return None, None, []

        similar_docs = self.vector_store.query_similar(
            embedding=query_embedding,
            top_k=self.top_k                                                                 
                                                               
        )
        
                                                                                          
                                                                                                  
        logger.info(f"PrimaryDetector found {len(similar_docs)} similar documents for query.")
        return query_features, query_embedding, similar_docs


if __name__ == '__main__':
                                                      
                                                                  
    logging.basicConfig(
        level=logging.INFO, 
        format="%(asctime)s - %(name)s - %(levelname)s - %(module)s:%(lineno)d - %(message)s"
    )
    logger = logging.getLogger(__name__)                                          

                                                                         
                                                                                        
    
    from rich.console import Console
    from rich.panel import Panel
    from rich.table import Table
    from unittest.mock import MagicMock

                                              
    from fortress.common.data_models import InputPromptRecord, DatabasePromptRecord, QueryFeatures
    from fortress.core.embedding_model import EmbeddingModel                               
    from fortress.core.nlp_analyzer import NLPAnalyzer                               
                                                                                   
    try:
        from fortress.config import setup_logging as fortress_setup_logging, get_config as fortress_get_config
                                                                        
        fortress_get_config()                                                            
    except ImportError:
        logger.warning("fortress.config.setup_logging or get_config not found. Using basic logging.")

    console = Console()
    console.print(Panel("[bold green]Starting PrimaryDetector Test Suite[/bold green]", expand=False))

                             
    console.print("\n[bold cyan]--- Test Case: Initialization and Basic Detection ---[/bold cyan]")

                          
                                      
    mock_embedding_model = MagicMock(spec=EmbeddingModel)
                                                                                                   

                                   
    mock_nlp_analyzer = MagicMock(spec=NLPAnalyzer)
    mock_nlp_features_dict = {
        "pos_tag_sequence": ["NN", "VB"], "dependency_relations_summary": {"nsubj": 1}, 
        "sentence_complexity": 10.5, "modal_verb_counts": {}, 
        "sentiment_scores": {"compound": 0.0}, "pronoun_type_counts": {},
        "char_level_stats": {"alphanum_ratio": 0.8, "special_char_count": 1},
        "dominant_language": "en"
    }
    mock_nlp_analyzer.extract_all_features.return_value = mock_nlp_features_dict

                                                                     
                                                            
    try:
        prompt_processor_under_test = PromptProcessor(embedding_model=mock_embedding_model, nlp_analyzer=mock_nlp_analyzer)
    except Exception as e:
        console.print(f"[bold red]Error instantiating PromptProcessor: {e}[/bold red]")
        raise
    
                                                                                                                                     
                                                                                                         
    mock_query_embedding = [0.1] * 1024                                          
    mock_embedding_model.get_embedding.return_value = mock_query_embedding                                  

                                                                                                          
    expected_query_features_obj = QueryFeatures(**mock_nlp_features_dict)

    mock_vector_store = MagicMock(spec=VectorStoreInterface)
    mock_similar_docs_return = [
        {"id": "doc1", "distance": 0.1, "metadata": {"original_prompt": "Test prompt 1", "label": 0, "source_file": "test.csv"}},
        {"id": "doc2", "distance": 0.2, "metadata": {"original_prompt": "Another test prompt", "label": 1, "source_file": "test.csv"}},
    ]
    mock_vector_store.query_similar.return_value = mock_similar_docs_return

                                    
    try:
        detector = PrimaryDetector(prompt_processor=prompt_processor_under_test, vector_store=mock_vector_store)
        console.print("PrimaryDetector instantiated successfully.")
    except Exception as e:
        console.print(f"[bold red]Error instantiating PrimaryDetector: {e}[/bold red]")
        raise

                           
    test_query = "This is a test query."
    console.print(f'Testing detect() with query: "{test_query}"')                     
    
                                                    
    query_features, query_embedding, similar_docs = detector.detect(test_query)

                                   
    table = Table(title="Detection Results from PrimaryDetector")
    table.add_column("Aspect", style="dim", width=30)
    table.add_column("Value / Status")

    if query_features:
        table.add_row("Query Features", "Extracted successfully")
                                                                                     
        if hasattr(query_features, 'token_source_log_probabilities') and query_features.token_source_log_probabilities is not None:
            table.add_row("  Token Log Probs", f"Present, count: {len(query_features.token_source_log_probabilities)}")
        else:
            table.add_row("  Token Log Probs", "Not present or None")
    else:
        table.add_row("Query Features", "[red]Failed to extract[/red]")

    if query_embedding:
        table.add_row("Query Embedding", f"Generated, dimensions: {len(query_embedding)}")
    else:
        table.add_row("Query Embedding", "[red]Failed to generate[/red]")
    
    table.add_row("Similar Documents", f"Found: {len(similar_docs)}")
    if similar_docs:
        table.add_row("  Top Match ID", str(similar_docs[0].get('id')))
        table.add_row("  Top Match Distance", f"{similar_docs[0].get('distance'):.4f}")

    console.print(table)
    
                          
    assert query_features is not None, "QueryFeatures should be generated"
    assert query_embedding is not None, "Query embedding should be generated"
                                                                                              
    assert similar_docs is not None, "Similar docs list should be returned"
    if mock_vector_store.query_similar.return_value:                                          
         assert len(similar_docs) == len(mock_similar_docs_return), "Mismatch in number of similar docs"

    console.print("[bold green]--- PrimaryDetector Test Case Completed ---[/bold green]")

                                             
                                    
    console.print("\n[bold cyan]--- Test Case: Empty Query String ---[/bold cyan]")
    empty_query_features, empty_query_embedding, empty_similar_docs = detector.detect("")
    table_empty = Table(title="Empty Query Test Results")
    table_empty.add_column("Aspect", style="dim", width=30)
    table_empty.add_column("Value / Status")
    table_empty.add_row("Query Features", "None" if empty_query_features is None else str(empty_query_features))
    table_empty.add_row("Query Embedding", "None" if empty_query_embedding is None else f"Dim: {len(empty_query_embedding)}")
    table_empty.add_row("Similar Docs", str(empty_similar_docs))
    console.print(table_empty)
    assert empty_query_features is None, "Features should be None for empty query"
    assert empty_query_embedding is None, "Embedding should be None for empty query"
    assert empty_similar_docs == [], "Similar docs should be empty list for empty query"

                                                                   
    console.print("\n[bold cyan]--- Test Case: Prompt Processor fails to return embedding ---[/bold cyan]")
    mock_embedding_model.reset_mock()
    mock_nlp_analyzer.reset_mock()
    mock_vector_store.reset_mock()
    mock_embedding_model.get_embedding.return_value = None                             
                                           
    mock_nlp_analyzer.extract_all_features.return_value = mock_nlp_features_dict 

                                                                                        
                                                                                         
    detector_embed_fail = PrimaryDetector(prompt_processor=prompt_processor_under_test, vector_store=mock_vector_store)
    
    fail_embed_features, fail_embed_embedding, fail_embed_docs = detector_embed_fail.detect("Query that causes embedding failure")
    
    table_fail_embed = Table(title="Embedding Failure Test Results")
    table_fail_embed.add_column("Aspect", style="dim", width=30)
    table_fail_embed.add_column("Value / Status")
    table_fail_embed.add_row("Query Features", "Available" if fail_embed_features else "[red]None[/red]")
    table_fail_embed.add_row("Query Embedding", "[red]None[/red]" if fail_embed_embedding is None else "[yellow]Present unexpectedly[/yellow]")
    table_fail_embed.add_row("Similar Docs", str(fail_embed_docs))
    console.print(table_fail_embed)

    assert fail_embed_features is not None, "Features should still be present even if embedding fails"
    assert isinstance(fail_embed_features, QueryFeatures)
    assert fail_embed_embedding is None, "Embedding should be None when PromptProcessor returns no embedding"
    assert fail_embed_docs == [], "Similar docs should be empty if embedding fails"
    mock_vector_store.query_similar.assert_not_called()                                       

                                                 
    console.print("\n[bold cyan]--- Test Case: Vector Store Query Fails ---[/bold cyan]")
    mock_embedding_model.reset_mock()
    mock_nlp_analyzer.reset_mock()
    mock_vector_store.reset_mock()
    mock_embedding_model.get_embedding.return_value = mock_query_embedding                               
    mock_nlp_analyzer.extract_all_features.return_value = mock_nlp_features_dict
    mock_vector_store.query_similar.side_effect = Exception("DB connection error")                          

    detector_vs_fail = PrimaryDetector(prompt_processor=prompt_processor_under_test, vector_store=mock_vector_store)
    vs_fail_features, vs_fail_embedding, vs_fail_docs = detector_vs_fail.detect("Query for VS failure")

    table_vs_fail = Table(title="Vector Store Failure Test Results")
    table_vs_fail.add_column("Aspect", style="dim", width=30)
    table_vs_fail.add_column("Value / Status")
    table_vs_fail.add_row("Query Features", "Available" if vs_fail_features else "[red]None[/red]")
    table_vs_fail.add_row("Query Embedding", "Available" if vs_fail_embedding else "[red]None[/red]")
    table_vs_fail.add_row("Similar Docs", str(vs_fail_docs))
    console.print(table_vs_fail)

    assert vs_fail_features is not None
    assert vs_fail_embedding is not None
    assert vs_fail_docs == [], "Similar docs should be empty if vector store query fails"
    mock_vector_store.query_similar.assert_called_once()

    console.print("\n[bold green]PrimaryDetector Test Suite Finished.[/bold green]")
