import chainlit as cl
from tig_engine import IntelliMod
from chainlit.input_widget import Select, Switch
import uuid

engine = IntelliMod()

@cl.on_chat_start
async def start():
    settings = await cl.ChatSettings(
        [
            Select(
                id="operation_mode",
                label="Operation Mode",
                values=[
                    "Prompt Builder - Craft Prompts",
                    "Standard - Chat & Tasks",
                    "Auditor - Evaluate Prompts"
                ],
                initial_index=0,
                description="Prompt Builder: interactive prompt crafting. Standard: chat with AI. Auditor: 9-step prompt evaluation."
            ),
            Select(
                id="model_tier",
                label="Model Strategy",
                values=[
                    "Auto-Router (TIG) - Best Match", 
                    "Free/Unlimited - Gemini 2.5 Flash", 
                    "Medium ($0.50/M) - Gemini 3 Flash", 
                    "Heavy ($3.00/M) - Claude Sonnet 4.5"
                ],
                initial_index=0,
                description="Select 'Free/Unlimited' to save credits."
            ),
            Select(
                id="agent_mode",
                label="Agent Persona",
                values=["Collaborator (Default)", "Project Manager", "Brainstorm Buddy", "Strict Coder"],
                initial_index=0,
                description="Adjusts response style."
            ),
            Switch(id="debug_tig", label="Show TIG Routing", initial=False),
        ]
    ).send()
    
    cl.user_session.set("settings", settings)
    cl.user_session.set("session_id", str(uuid.uuid4()))
    
    await cl.Message(
        content="**IntelliMod Online.**\n\n*Select a mode in Settings ⚙️*\n- **Prompt Builder** — I'll ask questions and craft the perfect prompt for you\n- **Standard** — Chat with AI, routed to the best model for your task\n- **Auditor** — Evaluate a prompt with the Modular Prompt Auditor"
    ).send()

@cl.on_settings_update
async def setup_agent(settings):
    cl.user_session.set("settings", settings)
    # Reset prompt builder state on mode change
    cl.user_session.set("session_id", str(uuid.uuid4()))
    mode = settings["operation_mode"]
    tier = settings["model_tier"]
    await cl.Message(content=f"**System Update:**\nMode: `{mode}`\nStrategy: `{tier}`").send()

@cl.on_message
async def main(message: cl.Message):
    user_input = message.content
    settings = cl.user_session.get("settings")
    session_id = cl.user_session.get("session_id")
    
    operation_mode = settings["operation_mode"]
    selected_tier = settings["model_tier"]
    show_tig = settings["debug_tig"]

    # Determine target model
    target_model = None
    if "Free/Unlimited" in selected_tier:
        target_model = "gemini-2.5-flash"
    elif "Medium" in selected_tier:
        target_model = "gemini-3-flash-preview"
    elif "Sonnet 4.5" in selected_tier:
        target_model = "anthropic/claude-sonnet-4"

    # ── PROMPT BUILDER MODE ──
    if "Prompt Builder" in operation_mode:
        msg = cl.Message(content="")
        
        async with cl.Step(name="💡 IntelliMod Prompt Builder", type="process") as step:
            step.input = user_input
            response, ready, compiled = engine.run_prompt_builder(user_input, session_id)
            
            if ready and compiled:
                step.output = "Prompt compiled successfully!"
                # Show the compiled prompt big and clear
                msg.content = f"**✅ Compiled Prompt**\n\n```\n{compiled}\n```\n\n*Send feedback to refine, or start over with a new request.*"
            else:
                step.output = "Asking next question..."
                msg.content = response
        
        await msg.send()

    # ── AUDITOR MODE ──
    elif "Auditor" in operation_mode:
        async with cl.Step(name="📋 MPA - Modular Prompt Auditor", type="process") as step:
            step.input = user_input
            step.output = "Running 9-step evaluation..."

        response_text = engine.run_mpa_pipeline(user_input, force_model=target_model)
        msg = cl.Message(content=response_text)
        await msg.send()

    # ── STANDARD MODE ──
    else:
        async with cl.Step(name="TIG Router", type="process") as step:
            step.input = user_input
            
            if target_model:
                step.output = f"Strategy: Manual Override\nModel: {target_model}"
            else:
                intent = engine.detect_intent(user_input)
                target_model = engine.tool_registry.get(intent, "gemini-2.5-flash")
                if show_tig:
                    step.output = f"Strategy: TIG Auto-Detect\nIntent: {intent.upper()}\nRouting to: {target_model}"
                else:
                    step.output = f"Model: {target_model}"

        response_text = engine.run_tig_pipeline(user_input, force_model=target_model)
        msg = cl.Message(content=response_text)
        await msg.send()