import os
import sys
import gradio as gr
from dotenv import load_dotenv
import subprocess

# Before initializing your model/agent, check and download data
if not all(os.path.exists(f) for f in [
    "embeddings/faiss_index.index", 
    "data/doc_chunks.pkl", 
    "embeddings/embeddings.pkl"
]):
    print("Downloading required data files...")
    try:
        result = subprocess.run(
            [sys.executable, "download_from_hub.py"],
            check=True,
            capture_output=True,
            text=True
        )
        print(result.stdout)
    except subprocess.CalledProcessError as e:
        print(f"Error downloading data: {e}")
        print(e.stdout)
        print(e.stderr)

# Add the current directory to the path so we can import from src
sys.path.append(".")

# Load environment variables
load_dotenv()

# Import after path setup
from src.agents.agent_director import AgentDirector
from src.utils.config import CHAT_MODEL

# Initialize the agent director
director = None

def check_files():
    """Check if required data files exist."""
    required_files = [
        "embeddings/faiss_index.index",
        "data/doc_chunks.pkl",
        "embeddings/embeddings.pkl"
    ]
    
    missing_files = []
    for file_path in required_files:
        if not os.path.exists(file_path):
            missing_files.append(file_path)
    
    return missing_files

def initialize_agent(model=CHAT_MODEL, top_k=50, debug=False):
    """Initialize the agent director."""
    global director
    
    # Check for missing files
    missing_files = check_files()
    if missing_files:
        raise FileNotFoundError(f"Required files not found: {', '.join(missing_files)}")
    
    # Initialize the agent director
    director = AgentDirector(model=model, top_k=top_k, debug=debug)
    return "Agentic Defensor initialized successfully!"

def process_query(query, top_k=50, debug=False):
    """Process a query using the agent director."""
    global director
    
    # Initialize the agent if not already done
    if director is None:
        try:
            initialize_agent(top_k=top_k, debug=debug)
        except Exception as e:
            return f"ERROR: Failed to initialize agent: {str(e)}"
    
    # Process the query
    try:
        result = director.process_query(query)
        
        # Format the response
        answer = result.get("answer", "No answer available.")
        
        # Add debugging information if available
        if debug and "reasoning_steps" in result:
            reasoning = "\n\n## Agent Reasoning\n"
            for step in result["reasoning_steps"]:
                reasoning += f"\n### {step['stage']}\n{step['reasoning']}\n"
            answer += reasoning
        
        return answer
    
    except Exception as e:
        return f"ERROR: {str(e)}"

# Create the Gradio interface
def create_interface():
    with gr.Blocks(title="Agentic Defensor") as demo:
        gr.Markdown("# Agentic Defensor")
        gr.Markdown("An agentic RAG system for legal defense analysis.")
        
        with gr.Tabs():
            with gr.TabItem("Query"):
                query_input = gr.Textbox(
                    label="Legal Query",
                    placeholder="Enter your legal query here...",
                    lines=3
                )
                
                with gr.Row():
                    top_k = gr.Slider(
                        minimum=10, 
                        maximum=200, 
                        value=50, 
                        step=10, 
                        label="Top K Results"
                    )
                    debug = gr.Checkbox(label="Show Agent Reasoning", value=False)
                
                submit_btn = gr.Button("Submit Query")
                output = gr.Markdown(label="Response")
                
                submit_btn.click(
                    fn=process_query,
                    inputs=[query_input, top_k, debug],
                    outputs=output
                )
            
            with gr.TabItem("About"):
                gr.Markdown("""
                ## About Agentic Defensor
                
                Agentic Defensor is a multi-agent legal analysis system that uses a specialized set of agents to:
                
                1. Analyze user queries to understand intent and extract key entities
                2. Retrieve relevant document chunks from a legal knowledge base
                3. Organize and aggregate context for comprehensive analysis
                4. Generate detailed legal defense analyses with source references
                
                This system goes beyond traditional RAG (Retrieval-Augmented Generation) by employing multiple specialized agents working together to produce high-quality legal analyses.
                
                ### Debug Mode
                
                The "Show Agent Reasoning" option allows you to see the agent's thought process during analysis, including query understanding, document retrieval decisions, context organization, and answer formulation.
                
                ### Source Code
                
                [GitHub Repository](https://github.com/vichudo/agentic-defensor)
                """)
    
    return demo

# Launch the app
if __name__ == "__main__":
    # Create and launch the interface
    demo = create_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860)