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

engine = IntelliMod()

@cl.on_chat_start
async def start():
    # --- UI SETTINGS PANEL ---
    settings = await cl.ChatSettings(
        [
            # 0. Operation Mode
            Select(
                id="operation_mode",
                label="Operation Mode",
                values=[
                    "Standard - Chat & Tasks",
                    "Auditor - Evaluate Prompts"
                ],
                initial_index=0,
                description="Standard: normal chat with TIG routing. Auditor: run the MPA 9-step prompt evaluation."
            ),

            # 1. Budget Tiers
            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",
                    "Heavy ($1.25/M) - GPT-5.1"
                ],
                initial_index=0,
                description="Select 'Free/Unlimited' to save credits."
            ),
            
            # 2. Agent Persona
            Select(
                id="agent_mode",
                label="Agent Persona",
                values=["Collaborator (Default)", "Project Manager", "Brainstorm Buddy", "Strict Coder"],
                initial_index=0,
                description="Adjusts Kael's personality."
            ),

            # 3. Toggles
            Switch(id="debug_tig", label="Show TIG Routing", initial=True),
        ]
    ).send()
    
    cl.user_session.set("settings", settings)
    await cl.Message(content="**Kael Online.**\n\n*System ready. Check Settings ⚙️ to switch between Standard and Auditor modes.*").send()

@cl.on_settings_update
async def setup_agent(settings):
    cl.user_session.set("settings", settings)
    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")
    
    # RETRIEVE SETTINGS
    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"
    elif "GPT-5.1" in selected_tier:
        target_model = "openai/gpt-5.1-chat"

    # BRANCH: Standard vs Auditor
    if "Auditor" in operation_mode:
        # --- AUDITOR MODE ---
        async with cl.Step(name="MPA - Modular Prompt Auditor", type="process") as step:
            step.input = user_input
            step.output = "Loading MPA v3.3 spec... Executing 9-step evaluation."
        
        async with cl.Step(name="⚙️ Loading Auditor", type="process") as loader:
            loader.output = "Auditor loaded. Running full evaluation..."

        response_text = engine.run_mpa_pipeline(user_input, force_model=target_model)

        msg = cl.Message(content=response_text)
        await msg.send()

    else:
        # --- STANDARD MODE (TIG routing) ---
        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")
                step.output = f"Strategy: TIG Auto-Detect\nIntent: {intent.upper()}\nRouting to: {target_model}"

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