import os
# Force disable tracing to prevent timeout in Docker
os.environ["CREWAI_TRACING_ENABLED"] = "false"

import chainlit as cl
from src.crews.definitions import CrewDefinitions
from src.router import SmartRouter

@cl.on_chat_start
async def on_chat_start():
    # No more menus! Just a welcome.
    await cl.Message(content="🧠 **Antigravity Brain Online.**\n\nI am ready. Just tell me what you need (e.g., *'Check the server health'* or *'Create a new agent named Bob'*).").send()
    cl.user_session.set("selected_crew", None)

@cl.on_message
async def on_message(message: cl.Message):
    user_input = message.content
    
    # 1. Check for commands
    if user_input.strip() == "/reset":
        cl.user_session.set("selected_crew", None)
        await cl.Message(content="🔄 Session reset. I will re-evaluate the best crew for your next request.").send()
        return

    # 2. Determine Crew
    current_crew = cl.user_session.get("selected_crew")
    
    # If no crew selected, OR if the input strongly suggests a switch (naive check, improved by router logic later if we want sticky sessions but smart switching)
    # For now: Sticky session. Once routed, stays routed until /reset.
    # PRO: Better context. CON: User asks 'Fix server' then 'Write poem' -> Infra tries to write poem.
    
    # Let's try "Auto-Detect on Every Turn" IF the context switch is obvious? 
    # No, that's risky. Let's stick to: Route First -> Sticky -> User can /reset.
    
    if not current_crew:
        msg_routing = cl.Message(content="🤔 Analyzing request...")
        await msg_routing.send()
        
        current_crew = SmartRouter.route(user_input)
        cl.user_session.set("selected_crew", current_crew)
        
        await msg_routing.update()
        await cl.Message(content=f"👉 **Routing to:** {current_crew}").send()

    # 3. Execution
    msg = cl.Message(content=f"🚀 **{current_crew}** is working on it...")
    await msg.send()
    
    try:
        # Assemble Crew
        crew = CrewDefinitions.assemble_crew(current_crew, inputs={"topic": user_input})
        
        # Run
        # Note: In a real async production app, we'd offload this to a thread pool properly.
        result = crew.kickoff(inputs={"topic": user_input})
        final_answer = str(result)
        
        # Update message
        msg.content = f"### ✅ Report from {current_crew}\n\n{final_answer}"
        await msg.update()
        
    except Exception as e:
        msg.content = f"❌ **Mission Failed:** {str(e)}"
        await msg.update()
        # Reset crew on failure so user can try again or get re-routed
        cl.user_session.set("selected_crew", None)