#!/usr/bin/env python3
"""
Run script for Agentic Defensor.

This script provides multiple ways to interact with the Agentic Defensor system:
1. API mode: Run the FastAPI server to handle queries over HTTP
2. CLI mode: Run a single query from the command line
3. Agent mode: Use the multi-agent system to process a query
4. Interactive mode: Start an interactive session to ask multiple questions
"""

import os
import sys
import json
import argparse
import uvicorn
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

def run_api(port):
    """Run the API server."""
    print(f"Starting Agentic Defensor API server on port {port}...")
    print(f"The API will be available at http://localhost:{port}")
    print("Press Ctrl+C to stop the server")
    uvicorn.run("src.api.app:app", host="0.0.0.0", port=port, reload=True)

def run_cli(query, top_k, model, output, verbose):
    """Run a query using the standard legal agent."""
    from src.main import process_query, save_result
    from src.utils.config import CHAT_MODEL
    
    # Ensure model is not None
    if model is None:
        model = CHAT_MODEL
    
    print(f"Processing query: {query}")
    print(f"Using model: {model}")
    result = process_query(query, top_k, model)
    
    # Print the answer
    print("\n--- Answer ---")
    print(result["answer"])
    
    # Print additional information if verbose
    if verbose:
        print("\n--- Query Information ---")
        print(f"Model used: {result['model_used']}")
        print(f"Retrieved chunks: {len(result['retrieved_chunks'])}")
    
    # Save the result if output path is provided
    if output:
        save_result(result, output)
        print(f"Results saved to {output}")

def run_agentic(query, top_k, model, output, verbose, debug):
    """Run a query using the multi-agent system."""
    from src.agents.agent_director import AgentDirector
    from src.utils.config import CHAT_MODEL
    
    # Ensure model is not None
    if model is None:
        model = CHAT_MODEL
    
    # Initialize the agent director
    print("Initializing agent director...")
    print(f"Using model: {model}")
    if debug:
        print("Debug mode enabled: Agent reasoning will be shown")
        
    director = AgentDirector(top_k=top_k, model=model, debug=debug)
    
    # Process the query
    print(f"\nProcessing query: {query}")
    result = director.process_query(query)
    
    # Display the result
    print("\n" + "="*80)
    print("QUERY:")
    print(query)
    print("\nANSWER:")
    print(result["answer"])
    print("="*80)
    
    # Display processing steps
    if verbose:
        print("\nPROCESSING STEPS:")
        if "query_analysis" in result:
            print("1. Query Analysis: Completed")
            structured_analysis = result["query_analysis"].get("structured_analysis", "")
            if structured_analysis:
                print(f"   - Extracted structured information from the query")
        
        print(f"2. Retrieved {result.get('num_chunks_retrieved', 0)} document chunks")
        
        if "context_aggregation" in result:
            agg = result["context_aggregation"]
            print("3. Context Aggregation:")
            print(f"   - Processed {agg.get('num_raw_content_items', 0)} content items")
            print(f"   - Organized context: {agg.get('has_organized_content', False)}")
        
        print(f"4. Answer Generation: Completed")
    
    # Save results if requested
    if output:
        print(f"\nSaving results to {output}...")
        with open(output, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"Results saved successfully.")

def run_interactive(model, top_k, agent_mode, verbose, debug):
    """Run an interactive session with the user."""
    from src.agents.agent_director import AgentDirector
    from src.agents.legal_agent import LegalAgent
    from src.utils.config import CHAT_MODEL
    
    # Ensure model is not None
    if model is None:
        model = CHAT_MODEL
    
    print("=== Agentic Defensor Interactive Mode ===")
    print("Type 'exit', 'quit', or 'q' to end the session.")
    print("Type 'help' or '?' for assistance.")
    print()
    
    if agent_mode:
        print("Using multi-agent system for processing queries.")
        if debug:
            print("Debug mode enabled: Agent reasoning will be shown")
        agent = AgentDirector(top_k=top_k, model=model, debug=debug)
    else:
        print("Using standard legal agent for processing queries.")
        agent = LegalAgent(model=model)
    
    print(f"Using model: {model}")
    print(f"Retrieving {top_k} chunks per query")
    
    history = []
    
    while True:
        # Get the query from the user
        try:
            query = input("\nYour query: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nExiting interactive mode.")
            break
        
        # Check for exit commands
        if query.lower() in ['exit', 'quit', 'q']:
            print("Exiting interactive mode.")
            break
        
        # Check for help command
        if query.lower() in ['help', '?']:
            print("\nAgentic Defensor Help:")
            print("- Type your legal query and press Enter to get an answer.")
            print("- Type 'exit', 'quit', or 'q' to end the session.")
            print("- Type 'history' to see your previous queries.")
            print("- Type 'save FILENAME' to save the session history to a file.")
            continue
        
        # Check for history command
        if query.lower() == 'history':
            if not history:
                print("No history available.")
            else:
                print("\nQuery History:")
                for i, item in enumerate(history, start=1):
                    print(f"{i}. {item['query']}")
            continue
        
        # Check for save command
        if query.lower().startswith('save '):
            filename = query[5:].strip()
            if not filename:
                print("Please provide a filename: save FILENAME")
                continue
            
            if not history:
                print("No history to save.")
                continue
            
            try:
                with open(filename, 'w', encoding='utf-8') as f:
                    json.dump(history, f, ensure_ascii=False, indent=2)
                print(f"History saved to {filename}")
            except Exception as e:
                print(f"Error saving history: {e}")
            continue
        
        # Skip empty queries
        if not query:
            continue
        
        # Process the query
        print("Processing query...")
        try:
            if agent_mode:
                result = agent.process_query(query)
            else:
                result = agent.answer_query(query, top_k)
            
            # Store in history
            history.append({
                'query': query,
                'answer': result.get('answer', 'No answer available.')
            })
            
            # Display the answer
            print("\n--- Answer ---")
            print(result.get('answer', 'No answer available.'))
            
            # Print additional information if verbose
            if verbose:
                if agent_mode and 'num_chunks_retrieved' in result:
                    print(f"\nRetrieved {result['num_chunks_retrieved']} document chunks")
                elif not agent_mode and 'retrieved_chunks' in result:
                    print(f"\nRetrieved {len(result['retrieved_chunks'])} document chunks")
                
                print(f"Used model: {result.get('model_used', model)}")
        
        except Exception as e:
            print(f"Error processing query: {e}")

def main():
    """Main function to parse arguments and run the appropriate mode."""
    # Create the top-level parser
    parser = argparse.ArgumentParser(description="Agentic Defensor: Legal RAG System")
    parser.add_argument('--model', type=str, default=None, help='OpenAI model to use')
    parser.add_argument('--verbose', action='store_true', help='Print verbose output')
    parser.add_argument('--debug', action='store_true', help='Show agent reasoning steps')
    
    # Create subparsers for different modes
    subparsers = parser.add_subparsers(dest='mode', help='Operating mode')
    
    # API mode
    api_parser = subparsers.add_parser('api', help='Run the API server')
    api_parser.add_argument('--port', type=int, default=8000, help='Port to run the API server on')
    
    # CLI mode
    cli_parser = subparsers.add_parser('cli', help='Run a query from the command line')
    cli_parser.add_argument('query', type=str, help='The legal query to process')
    cli_parser.add_argument('--top-k', type=int, default=200, help='Number of chunks to retrieve')
    cli_parser.add_argument('--output', type=str, default=None, help='Output file path for saving the response')
    
    # Agent mode
    agent_parser = subparsers.add_parser('agent', help='Run a query using the multi-agent system')
    agent_parser.add_argument('query', type=str, help='The legal query to process')
    agent_parser.add_argument('--top-k', type=int, default=50, help='Number of chunks to retrieve')
    agent_parser.add_argument('--output', type=str, default=None, help='Output file path for saving the response')
    
    # Interactive mode
    interactive_parser = subparsers.add_parser('interactive', help='Start an interactive session')
    interactive_parser.add_argument('--top-k', type=int, default=50, help='Number of chunks to retrieve')
    interactive_parser.add_argument('--agent', action='store_true', help='Use multi-agent system')
    
    # Parse the arguments
    args = parser.parse_args()
    
    # Run the appropriate mode
    if args.mode == 'api':
        run_api(args.port)
    elif args.mode == 'cli':
        run_cli(args.query, args.top_k, args.model, args.output, args.verbose)
    elif args.mode == 'agent':
        run_agentic(args.query, args.top_k, args.model, args.output, args.verbose, args.debug)
    elif args.mode == 'interactive':
        run_interactive(args.model, args.top_k, args.agent, args.verbose, args.debug)
    else:
        # Default to interactive mode if no mode specified
        print("No mode specified, starting interactive mode...\n")
        run_interactive(args.model, 50, False, args.verbose, args.debug)

if __name__ == "__main__":
    main()