                                                        
import os
import sys
import argparse
import logging                       
import time                   
import json                                  

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

                                                                                               
                                                                          
try:
    from fortress.config import setup_logging                                      
    setup_logging()                          
except ImportError:
                                                                   
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(module)s:%(lineno)d - %(message)s")
    logging.warning("Failed to import setup_logging from fortress.config or scripts.utils. Using basicConfig for logging.")

logger = logging.getLogger(__name__)                                                

                  
from fortress.core.embedding_model import EmbeddingModel
from fortress.core.nlp_analyzer import NLPAnalyzer
from fortress.data_management.prompt_processor import PromptProcessor
from fortress.core.vector_store_interface import ChromaVectorStore                                                     
from fortress.detection_pipeline.primary_detector import PrimaryDetector
from fortress.detection_pipeline.secondary_analyzer import SecondaryAnalyzer
from fortress.detection_pipeline.main_pipeline import DetectionPipeline
from fortress.config import get_config                                                              
from fortress.common.data_models import DecisionLabel, FinalDetectionOutput                                   

                 
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.table import Table
from rich.syntax import Syntax
from rich.box import ROUNDED                                        

def main():
    parser = argparse.ArgumentParser(description="Test the FORTRESS detection pipeline on a single prompt string.")
                          
                         
                    
                                            
                                                   
    parser.add_argument(
        "--collection_name",
        type=str,
        default=None,                                                             
        help="Name of the ChromaDB collection to use. Uses default from settings if not provided."
    )   
    parser.add_argument(
        "--log_level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Set the logging level for the script."
    )
                                                                                        

    args = parser.parse_args()

                                                       
    try:
        logging.getLogger().setLevel(args.log_level.upper())
        logger.info(f"Logging level set to {args.log_level.upper()}")
    except Exception as e:
        logger.error(f"Could not set log level to {args.log_level.upper()}: {e}")


    console = Console()
    console.print(Panel(Text("FORTRESS Prompt Detection CLI - Interactive Mode", justify="center", style="bold blue")))
    
    if args.collection_name:
        logger.info(f"Using custom collection name: {args.collection_name}")
    else:
        logger.info(f"Using default collection name from ChromaVectorStore/settings.")

    try:
                                                  
        logger.info("Initializing FORTRESS components...")
        config = get_config()

        embedding_model = EmbeddingModel()
        nlp_analyzer = NLPAnalyzer()
        
        vector_store_params = {}
        if args.collection_name:
            vector_store_params['collection_name'] = args.collection_name
        vector_store = ChromaVectorStore(**vector_store_params)
        
        prompt_processor = PromptProcessor(embedding_model, nlp_analyzer, vector_store)
        
        primary_detector = PrimaryDetector(prompt_processor, vector_store)
        secondary_analyzer = SecondaryAnalyzer()
        
        detection_pipeline = DetectionPipeline(
            prompt_processor=prompt_processor,
            primary_detector=primary_detector,
            secondary_analyzer=secondary_analyzer
        )
        logger.info("FORTRESS components initialized successfully.")
        console.print(Panel(Text("Components Initialized. Ready for prompts.", style="green"), padding=(0,1)))

    except FileNotFoundError as e:
        logger.error(f"Initialization failed: File not found error: {e}. This might be related to model files or config.", exc_info=True)
        console.print(f"[bold red]Error: File Not Found - {e}. Please ensure all necessary files (models, configs) are accessible. Exiting.[/bold red]")
        return
    except ConnectionError as e:
        logger.error(f"Initialization failed: Connection error, possibly with ChromaDB: {e}", exc_info=True)
        console.print(f"[bold red]Error: Connection Failed - {e}. Check if your vector database server is running and accessible. Exiting.[/bold red]")
        return
    except ImportError as e:
        logger.error(f"Initialization failed: Import error: {e}. Some required packages might be missing or not in PYTHONPATH.", exc_info=True)
        console.print(f"[bold red]Error: Import Failed - {e}. Ensure all dependencies are installed correctly. Exiting.[/bold red]")
        return
    except Exception as e:
        logger.critical(f"An unexpected error occurred during initialization: {e}", exc_info=True)
        console.print("[bold red]An unexpected error occurred during initialization:[/bold red]")
        console.print_exception(show_locals=False)
        return

                                             
    while True:
        try:
            prompt_text = console.input(Text("Enter prompt (or 'exit'/'quit' to stop): ", style="bold cyan")).strip()
        except KeyboardInterrupt:
            console.print("\n[bold yellow]Interrupted. Exiting...[/bold yellow]")
            break                 

        if prompt_text.lower() in ['exit', 'quit']:
            break
        
        if not prompt_text:
            logger.warning("Empty prompt received.")
            console.print(Text("Prompt cannot be empty. Please try again.", style="yellow"))
            continue

        console.print(Panel(f"Analyzing prompt: [bold yellow]{Text(prompt_text, overflow='fold')}[/bold yellow]", expand=False, border_style="cyan", title="[dim]Input[/dim]"))
        logger.info(f"Received prompt for detection: '{prompt_text[:100]}...'")                       

        try:
            logger.info(f"Running detection pipeline for: '{prompt_text[:100]}'")
            start_time = time.time()
            final_output: FinalDetectionOutput = detection_pipeline.run(prompt_text)
            end_time = time.time()
            duration_ms = (end_time - start_time) * 1000
            logger.info(f"Detection pipeline finished for current prompt. Time: {duration_ms:.2f} ms")

                                     
            console.print(Panel(Text("Detection Results", justify="center", style="bold magenta"), box=ROUNDED, padding=(0,2)))

            summary_table = Table(box=None, show_header=False, padding=(0, 1, 0, 1))
            summary_table.add_column("Key", style="dim", width=20)
            summary_table.add_column("Value")

            summary_table.add_row("Input Prompt:", Text(final_output.query_text, style="yellow", overflow="fold"))
            
            decision_style = "green"
            if final_output.final_decision == DecisionLabel.UNSAFE:
                decision_style = "red"
            elif final_output.final_decision == DecisionLabel.AMBIGUOUS:
                decision_style = "yellow"
            elif final_output.final_decision == DecisionLabel.ERROR:
                decision_style = "bold red"

            summary_table.add_row("Final Decision:", Text(final_output.final_decision.value, style=decision_style))
            summary_table.add_row("Confidence:", f"{final_output.overall_confidence:.4f}")
            summary_table.add_row("Is Ambiguous:", str(final_output.is_ambiguous))
            summary_table.add_row("Justification:", Text(final_output.justification if final_output.justification else "N/A", overflow="fold"))
            
            if final_output.error_info:
                summary_table.add_row("Error Info:", Text(final_output.error_info, style="red", overflow="fold"))

            console.print(summary_table)

                                                      
            if final_output.primary_detector_top_k_results and len(final_output.primary_detector_top_k_results) > 0:
                console.print(Text("\n--- Top K Similar Documents (Primary Detector) ---", style="bold cyan"))
                docs_table = Table(title="Retrieved Documents Details", box=ROUNDED, show_lines=True, padding=(0,1), expand=False)
                docs_table.add_column("ID", style="dim cyan", width=10, overflow="fold")
                docs_table.add_column("Snippet (Original Prompt)", max_width=60, overflow="fold")
                docs_table.add_column("Similarity\n(Distance)", justify="right", style="magenta")
                docs_table.add_column("True Label", justify="center")
                docs_table.add_column("Perplexity", justify="right", style="blue")                          
                docs_table.add_column("Source File", style="dim green", overflow="fold", max_width=30)

                for doc in final_output.primary_detector_top_k_results:
                    metadata = doc.get("metadata", {})
                    doc_id = str(doc.get("id", "N/A"))
                    snippet = str(metadata.get("original_prompt", "N/A"))
                    distance = f"{doc.get('distance', float('nan')):.4f}"
                    
                    true_label_val = metadata.get("label")                 
                    if true_label_val == 1:         
                        true_label_str = Text("UNSAFE", style="red")
                    elif true_label_val == 0:       
                        true_label_str = Text("SAFE", style="green")
                    else:
                        true_label_str = Text("N/A", style="dim")
                        
                                                                  
                    perplexity_val = metadata.get("perplexity")                                           
                    perplexity_str = f"{perplexity_val:.4f}" if isinstance(perplexity_val, (float, int)) else "N/A"
                    
                    source_file = str(metadata.get("source_file", "N/A"))
                    
                    docs_table.add_row(
                        doc_id,
                        Text(snippet), 
                        distance,
                        true_label_str,
                        perplexity_str,                                    
                        Text(source_file)
                    )
                console.print(docs_table)
            else:
                console.print(Text("No similar documents found by primary detector.", style="yellow"))

                                                           
            if final_output.detection_stages_summary:
                console.print(Text("\n--- Detection Stages Summary ---", style="bold cyan"))
                stages_str = json.dumps(final_output.detection_stages_summary, indent=2)
                console.print(Syntax(stages_str, "json", theme="native", line_numbers=False, word_wrap=True))

            console.print(Text("-" * console.width, style="dim"))                 

        except ConnectionRefusedError as e:
            logger.error(f"Connection refused during prompt processing: {e}", exc_info=True)
            console.print(f"[bold red]Error: Connection refused - {e}[/bold red]")
        except Exception as e_loop:                                                    
            logger.error(f"Error processing prompt '{prompt_text}': {e_loop}", exc_info=True)
            console.print(f"[bold red]Error processing prompt: {e_loop}[/bold red]")
            console.print_exception(show_locals=False)                                         
            console.print("[yellow]Continuing to next prompt...[/yellow]")
                            

    console.print(Panel(Text("Exited FORTRESS Prompt Detection CLI.", justify="center", style="bold blue")))


if __name__ == "__main__":
    main()
