import os
import json
import argparse
from dotenv import load_dotenv

from src.agents.agent_director import AgentDirector

# Load environment variables
load_dotenv()

def main():
    """Example script demonstrating the Agentic Defensor RAG system with multiple agents."""
    
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Agentic Defensor Multi-Agent Example")
    parser.add_argument("--query", type=str, default="En qué tomo se encuentra Contrato Andrea Monsalve",
                      help="The legal query to process")
    parser.add_argument("--top-k", type=int, default=50, help="Number of chunks to retrieve")
    parser.add_argument("--model", type=str, default=None, help="OpenAI model to use")
    parser.add_argument("--save", action="store_true", help="Save results to file")
    parser.add_argument("--output", type=str, default="agentic_results.json", help="Output file")
    args = parser.parse_args()
    
    # Initialize the agent director
    print("Initializing agent director...")
    director = AgentDirector(top_k=args.top_k, model=args.model)
    
    # Process the query
    print(f"\nProcessing query: {args.query}")
    result = director.process_query(args.query)
    
    # Display the result
    print("\n" + "="*80)
    print("QUERY:")
    print(args.query)
    print("\nANSWER:")
    print(result["answer"])
    print("="*80)
    
    # Display processing steps
    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 args.save:
        print(f"\nSaving results to {args.output}...")
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"Results saved successfully.")

if __name__ == "__main__":
    main()