from fastapi import FastAPI, Request, Form, BackgroundTasks, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import jinja2
import os
from typing import Dict, Any
import threading
import time
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Check required environment variables
required_env_vars = [
    'FINNHUB_API_KEY',
    'OPENAI_API_KEY',
    #'REDDIT_CLIENT_ID',
    #'REDDIT_CLIENT_SECRET',
    #'REDDIT_USER_AGENT'
]

missing_vars = [var for var in required_env_vars if not os.getenv(var)]
if missing_vars:
    print(f"Error: Missing required environment variables: {', '.join(missing_vars)}")
    print("Please create a .env file with these variables or set them in your environment.")

from tradingagents.graph.trading_graph import TradingAgentsGraph

app = FastAPI()

# In-memory storage for the process state
# Using a lock for thread-safe access to app_state
app_state_lock = threading.Lock()
app_state: Dict[str, Any] = {
    "process_running": False,
    "company_symbol": None,
    "execution_tree": [],
    "overall_status": "idle", # idle, in_progress, completed, error
    "overall_progress": 0 # 0-100
}

# Mount the static directory to serve CSS, JS, etc.
app.mount("/static", StaticFiles(directory="webapp/static"), name="static")

# Setup Jinja2 for templating
template_dir = os.path.join(os.path.dirname(__file__), "templates")
jinja_env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir))

def update_execution_state(state: Dict[str, Any]):
    """Callback function to update the app_state based on LangGraph's state."""
    print(f"📡 Callback received state keys: {list(state.keys())}")
    
    with app_state_lock:
        # Initialize the complete execution tree structure if not exists
        if not app_state["execution_tree"] or (
            len(app_state["execution_tree"]) == 1 and 
            app_state["execution_tree"][0]["id"] == "initialization"
        ):
            app_state["execution_tree"] = initialize_complete_execution_tree()
        
        # Map LangGraph node names to our tracking system
        agent_state_mapping = {
            "Market Analyst": {
                "phase": "data_collection", 
                "agent_id": "market_analyst",
                "report_key": "market_report",
                "report_name": "Market Analysis Report"
            },
            "Social Analyst": {
                "phase": "data_collection", 
                "agent_id": "social_analyst",
                "report_key": "sentiment_report", 
                "report_name": "Sentiment Analysis Report"
            },
            "News Analyst": {
                "phase": "data_collection", 
                "agent_id": "news_analyst",
                "report_key": "news_report",
                "report_name": "News Analysis Report"
            },
            "Fundamentals Analyst": {
                "phase": "data_collection", 
                "agent_id": "fundamentals_analyst",
                "report_key": "fundamentals_report",
                "report_name": "Fundamentals Report"
            },
            "Bull Researcher": {
                "phase": "research", 
                "agent_id": "bull_researcher",
                "report_key": "investment_debate_state.bull_history",
                "report_name": "Bull Case Analysis"
            },
            "Bear Researcher": {
                "phase": "research", 
                "agent_id": "bear_researcher",
                "report_key": "investment_debate_state.bear_history",
                "report_name": "Bear Case Analysis"
            },
            "Research Manager": {
                "phase": "research", 
                "agent_id": "research_manager",
                "report_key": "investment_debate_state.judge_decision",
                "report_name": "Research Synthesis"
            },
            "Trade Planner": {
                "phase": "planning", 
                "agent_id": "trade_planner",
                "report_key": "trader_investment_plan",
                "report_name": "Trading Plan"
            },
            "Trader": {
                "phase": "execution", 
                "agent_id": "trader",
                "report_key": "investment_plan",
                "report_name": "Execution Report"
            },
            "Risky Analyst": {
                "phase": "risk_analysis", 
                "agent_id": "risky_analyst",
                "report_key": "risk_debate_state.risky_history",
                "report_name": "Risk Assessment (Aggressive)"
            },
            "Neutral Analyst": {
                "phase": "risk_analysis", 
                "agent_id": "neutral_analyst",
                "report_key": "risk_debate_state.neutral_history",
                "report_name": "Risk Assessment (Neutral)"
            },
            "Safe Analyst": {
                "phase": "risk_analysis", 
                "agent_id": "safe_analyst",
                "report_key": "risk_debate_state.safe_history",
                "report_name": "Risk Assessment (Conservative)"
            },
            "Risk Judge": {
                "phase": "risk_analysis", 
                "agent_id": "risk_judge",
                "report_key": "final_trade_decision",
                "report_name": "Final Risk Decision"
            }
        }
        
        # Update agent statuses based on available reports
        for agent_name, agent_info in agent_state_mapping.items():
            # Check if this agent has completed (has report data)
            report_data = get_nested_value(state, agent_info["report_key"])
            if report_data:
                update_agent_status(agent_info, "completed", report_data, state)
        
        # Update overall progress
        root_node = app_state["execution_tree"][0]
        total_agents = len(agent_state_mapping)
        completed_agents = count_completed_agents(root_node)
        app_state["overall_progress"] = min(100, int((completed_agents / max(total_agents, 1)) * 100))
        
        print(f"📊 Progress updated: {app_state['overall_progress']}% ({completed_agents}/{total_agents} agents)")

def initialize_complete_execution_tree():
    """Initialize the complete execution tree with all agents in pending state."""
    return [{
        "id": "root",
        "name": f"📈 Trading Analysis for {app_state['company_symbol']}",
        "status": "in_progress", 
        "content": f"Comprehensive trading analysis for {app_state['company_symbol']}",
        "children": [
            {
                "id": "data_collection_phase",
                "name": "📊 Data Collection Phase",
                "status": "pending",
                "content": "Collecting market data and analysis from various sources",
                "children": [
                    create_agent_node("market_analyst", "📈 Market Analyst"),
                    create_agent_node("social_analyst", "📱 Social Media Analyst"),
                    create_agent_node("news_analyst", "📰 News Analyst"),
                    create_agent_node("fundamentals_analyst", "📊 Fundamentals Analyst")
                ]
            },
            {
                "id": "research_phase",
                "name": "🔍 Research Phase",
                "status": "pending",
                "content": "Research and debate investment perspectives",
                "children": [
                    create_agent_node("bull_researcher", "🐂 Bull Researcher"),
                    create_agent_node("bear_researcher", "🐻 Bear Researcher"),
                    create_agent_node("research_manager", "🔍 Research Manager")
                ]
            },
            {
                "id": "planning_phase",
                "name": "📋 Planning Phase", 
                "status": "pending",
                "content": "Develop trading strategy and execution plan",
                "children": [
                    create_agent_node("trade_planner", "📋 Trade Planner")
                ]
            },
            {
                "id": "execution_phase",
                "name": "⚡ Execution Phase",
                "status": "pending", 
                "content": "Execute trades based on analysis and planning",
                "children": [
                    create_agent_node("trader", "⚡ Trader")
                ]
            },
            {
                "id": "risk_analysis_phase",
                "name": "⚠️ Risk Management Phase",
                "status": "pending",
                "content": "Assess and manage investment risks",
                "children": [
                    create_agent_node("risky_analyst", "🚨 Aggressive Risk Analyst"),
                    create_agent_node("neutral_analyst", "⚖️ Neutral Risk Analyst"),
                    create_agent_node("safe_analyst", "🛡️ Conservative Risk Analyst"),
                    create_agent_node("risk_judge", "⚠️ Risk Judge")
                ]
            }
        ],
        "timestamp": time.time()
    }]

def create_agent_node(agent_id: str, agent_name: str):
    """Create a standardized agent node with report and messages sub-items."""
    return {
        "id": agent_id,
        "name": agent_name,
        "status": "pending",
        "content": f"Agent: {agent_name} - Awaiting execution",
        "children": [
            {
                "id": f"{agent_id}_report",
                "name": "📄 Report",
                "status": "pending",
                "content": "Report not yet generated",
                "children": [],
                "timestamp": time.time()
            },
            {
                "id": f"{agent_id}_messages",
                "name": "💬 Messages",
                "status": "pending", 
                "content": "No messages yet",
                "children": [],
                "timestamp": time.time()
            }
        ],
        "timestamp": time.time()
    }

def get_nested_value(data: dict, key_path: str):
    """Get value from nested dict using dot notation (e.g., 'investment_debate_state.bull_history')."""
    keys = key_path.split('.')
    value = data
    for key in keys:
        if isinstance(value, dict) and key in value:
            value = value[key]
        else:
            return None
    return value

def update_agent_status(agent_info: dict, status: str, report_data: any, full_state: dict):
    """Update an agent's status and content in the execution tree."""
    root_node = app_state["execution_tree"][0]
    
    # Find the agent in the tree
    agent_node = find_agent_in_tree(agent_info["agent_id"], root_node)
    if not agent_node:
        return
        
    # Update agent status
    if agent_node["status"] != "completed":
        agent_node["status"] = status
        agent_node["content"] = f"✅ {agent_node['name']} - Analysis completed"
        
        # Update report sub-item
        report_node = find_item_by_id(f"{agent_info['agent_id']}_report", agent_node["children"])
        if report_node:
            report_node["status"] = "completed"
            report_node["content"] = format_report_content(agent_info["report_name"], report_data)
            
        # Update messages sub-item (extract from state if available)
        messages_node = find_item_by_id(f"{agent_info['agent_id']}_messages", agent_node["children"])
        if messages_node:
            messages_node["status"] = "completed"
            messages_node["content"] = extract_agent_messages(full_state, agent_info["agent_id"])
    
    # Update phase status if all agents in phase are completed
    update_phase_status_if_complete(agent_info["phase"], root_node)

def find_agent_in_tree(agent_id: str, root_node: dict):
    """Find an agent node in the execution tree."""
    for phase in root_node["children"]:
        for agent in phase["children"]:
            if agent["id"] == agent_id:
                return agent
    return None

def find_item_by_id(item_id: str, items: list):
    """Find an item by ID in a list of items."""
    for item in items:
        if item["id"] == item_id:
            return item
    return None

def format_report_content(report_name: str, report_data: any) -> str:
    """Format report data for display."""
    if isinstance(report_data, str):
        return f"📄 {report_name}\n\n{report_data}"
    elif isinstance(report_data, dict):
        return f"📄 {report_name}\n\n{str(report_data)}"
    elif isinstance(report_data, list) and report_data:
        # For debate histories, show the latest message
        latest = report_data[-1] if report_data else "No data"
        return f"📄 {report_name}\n\n{str(latest)}"
    else:
        return f"📄 {report_name}\n\nReport generated successfully"

def extract_agent_messages(state: dict, agent_id: str) -> str:
    """Extract relevant messages for an agent from the state."""
    # This is a simplified version - could be enhanced to extract actual messages
    messages = state.get("messages", [])
    if messages:
        return f"💬 Agent Messages\n\n{len(messages)} messages exchanged during execution"
    else:
        return "💬 Agent Messages\n\nExecution completed without specific message logs"

def update_phase_status_if_complete(phase_id: str, root_node: dict):
    """Update phase status to completed if all its agents are completed."""
    phase_node = find_item_by_id(f"{phase_id}_phase", root_node["children"])
    if not phase_node:
        return
        
    # Check if all agents in this phase are completed
    all_completed = all(agent["status"] == "completed" for agent in phase_node["children"])
    if all_completed and phase_node["status"] != "completed":
        phase_node["status"] = "completed"
        phase_node["content"] = f"✅ {phase_node['name']} - All agents completed successfully"

def count_completed_agents(root_node: dict) -> int:
    """Count the number of completed agents across all phases."""
    count = 0
    for phase in root_node["children"]:
        for agent in phase["children"]:
            if agent["status"] == "completed":
                count += 1
    return count

def run_trading_process(company_symbol: str, config: Dict[str, Any]):
    """Runs the TradingAgentsGraph in a separate thread."""
    with app_state_lock:
        app_state["overall_status"] = "in_progress"
        app_state["overall_progress"] = 0

    try:
        # Import and create custom config
        from tradingagents.default_config import DEFAULT_CONFIG
        
        # Create custom configuration with user selections
        custom_config = DEFAULT_CONFIG.copy()
        custom_config["llm_provider"] = config["llm_provider"]
        custom_config["max_debate_rounds"] = config["max_debate_rounds"]
        custom_config["cost_per_trade"] = config["cost_per_trade"]
        
        # Set the appropriate LLM models based on provider
        if config["llm_provider"] == "google":
            custom_config["gemini_quick_think_llm"] = config["quick_think_llm"]
            custom_config["gemini_deep_think_llm"] = config["deep_think_llm"]
        else:
            custom_config["quick_think_llm"] = config["quick_think_llm"]
            custom_config["deep_think_llm"] = config["deep_think_llm"]
        
        # Set backend URL based on provider
        if config["llm_provider"] == "openrouter":
            custom_config["backend_url"] = "https://openrouter.ai/api/v1"
        elif config["llm_provider"] == "google":
            custom_config["backend_url"] = "https://generativelanguage.googleapis.com/v1"
        elif config["llm_provider"] == "anthropic":
            custom_config["backend_url"] = "https://api.anthropic.com/"
        elif config["llm_provider"] == "ollama":
            custom_config["backend_url"] = f"http://{os.getenv('OLLAMA_HOST', 'localhost')}:11434/v1"
        else:  # openai
            custom_config["backend_url"] = "https://api.openai.com/v1"
        
        print(f"🚀 Initializing TradingAgentsGraph for {company_symbol}")
        graph = TradingAgentsGraph(config=custom_config)
        analysis_date = config["analysis_date"]  # Use user-selected date
        print(f"🔄 Starting propagation for {company_symbol} on {analysis_date}")
        
        # The propagate method now accepts the callback and trade_date
        final_state, processed_signal = graph.propagate(company_symbol, trade_date=analysis_date, on_step_callback=update_execution_state)
        print(f"✅ Propagation completed for {company_symbol}")
        
        with app_state_lock:
            app_state["overall_status"] = "completed"
            app_state["overall_progress"] = 100
            # Update the root node status to completed
            if app_state["execution_tree"]:
                app_state["execution_tree"][0]["status"] = "completed"
                app_state["execution_tree"][0]["content"] = f"✅ Analysis completed successfully!\n\nFinal Decision: {processed_signal}\n\nFull State: {str(final_state)}"

    except Exception as e:
        import traceback
        error_detail = traceback.format_exc()
        with app_state_lock:
            app_state["overall_status"] = "error"
            app_state["overall_progress"] = 100
            if app_state["execution_tree"]:
                app_state["execution_tree"][0]["status"] = "error"
                app_state["execution_tree"][0]["content"] = f"Error during execution: {str(e)}\n\n{error_detail}"
            # Add a specific error item to the tree
            app_state["execution_tree"].append({
                "id": "error",
                "name": "Process Error",
                "status": "error",
                "content": f"Error during execution: {str(e)}\n\n{error_detail}",
                "children": [],
                "timestamp": time.time()
            })
    finally:
        with app_state_lock:
            app_state["process_running"] = False


@app.get("/", response_class=HTMLResponse)
async def read_root():
    template = jinja_env.get_template("index.html")
    return template.render(app_state=app_state)

@app.post("/start", response_class=HTMLResponse)
async def start_process(
    background_tasks: BackgroundTasks, 
    company_symbol: str = Form(...),
    llm_provider: str = Form(...), 
    quick_think_llm: str = Form(...),
    deep_think_llm: str = Form(...),
    max_debate_rounds: int = Form(...),
    cost_per_trade: float = Form(...),
    analysis_date: str = Form(...)
):
    # Check if all required environment variables are set
    missing_vars = [var for var in required_env_vars if not os.getenv(var)]
    if missing_vars:
        app_state["overall_status"] = "error"
        app_state["execution_tree"] = [{
            "id": "error",
            "name": "Configuration Error",
            "status": "error",
            "content": f"Missing required environment variables: {', '.join(missing_vars)}. Please check .env.example file.",
            "children": [],
            "timestamp": time.time()
        }]
        template = jinja_env.get_template("_partials/left_panel.html")
        return template.render(tree=app_state["execution_tree"], app_state=app_state)

    with app_state_lock:
        if app_state["process_running"]:
            # Optionally, return an error or a message that a process is already running
            template = jinja_env.get_template("_partials/left_panel.html")
            return template.render(tree=app_state["execution_tree"], app_state=app_state)

        app_state["process_running"] = True
        app_state["company_symbol"] = company_symbol
        app_state["overall_status"] = "in_progress"
        app_state["overall_progress"] = 5  # Show initial progress
        
        # Store all configuration parameters
        app_state["config"] = {
            "llm_provider": llm_provider,
            "quick_think_llm": quick_think_llm,
            "deep_think_llm": deep_think_llm,
            "max_debate_rounds": max_debate_rounds,
            "cost_per_trade": cost_per_trade,
            "analysis_date": analysis_date
        }
        
        # Initialize execution tree with startup message
        app_state["execution_tree"] = [{
            "id": "initialization",
            "name": f"🚀 Initializing Trading Analysis for {company_symbol}",
            "status": "in_progress",
            "content": f"Starting comprehensive trading analysis for {company_symbol}...\n\nConfiguration:\n• LLM Provider: {llm_provider}\n• Quick Think Model: {quick_think_llm}\n• Deep Think Model: {deep_think_llm}\n• Max Debate Rounds: {max_debate_rounds}\n• Cost Per Trade: ${cost_per_trade}\n• Analysis Date: {analysis_date}\n\nInitializing trading agents and preparing analysis pipeline...",
            "children": [],
            "timestamp": time.time()
        }]

    background_tasks.add_task(run_trading_process, company_symbol, app_state["config"])
    
    template = jinja_env.get_template("_partials/left_panel.html")
    return template.render(tree=app_state["execution_tree"], app_state=app_state)

@app.get("/status", response_class=HTMLResponse)
async def get_status():
    with app_state_lock:
        template = jinja_env.get_template("_partials/left_panel.html")
        return template.render(tree=app_state["execution_tree"], app_state=app_state)

def find_item_in_tree(item_id: str, tree: list) -> Dict[str, Any] | None:
    """Recursively searches the execution tree for an item by its ID."""
    for item in tree:
        if item["id"] == item_id:
            return item
        if item["children"]:
            found_child = find_item_in_tree(item_id, item["children"])
            if found_child:
                return found_child
    return None

@app.get("/content/{item_id}", response_class=HTMLResponse)
async def get_item_content(item_id: str):
    with app_state_lock:
        item = find_item_in_tree(item_id, app_state["execution_tree"])
        if item:
            template = jinja_env.get_template("_partials/right_panel.html")
            return template.render(content=item.get("content", "No content available."))
        else:
            return HTMLResponse(content="<p>Item not found.</p>", status_code=404)

# To run this app:
# uvicorn webapp.main:app --reload