"""Demo 3 — Multi-Agent Router with SSE streaming and SQLite thread memory."""

from __future__ import annotations

import json
import math
import sys
import uuid
from datetime import UTC, datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from ui import chat_page, info_row

from koan.agents import AgentConfig
from koan.agents.events import AgentComplete, AgentError, TextDelta, ToolCallEnd, ToolCallStart
from koan.agents.orchestrator import Orchestrator
from koan.agents.router import RuleRouter
from koan.agents.stream import StreamingReactAgent
from koan.llm.client import LLMClient
from koan.llm.config import ProviderSettings
from koan.memory.thread import ThreadStore
from koan.tools import ToolRegistry, tool

load_dotenv()

# ---------------------------------------------------------------------------
# Weather tools
# ---------------------------------------------------------------------------


@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city.

    Args:
        city: The city name to look up weather for.
    """
    weather_data = {
        "new york": "72°F, partly cloudy",
        "london": "58°F, rainy",
        "tokyo": "80°F, sunny",
        "paris": "65°F, overcast",
        "sydney": "68°F, clear skies",
        "berlin": "55°F, windy",
        "mumbai": "90°F, humid",
    }
    key = city.lower().strip()
    if key in weather_data:
        return f"Weather in {city}: {weather_data[key]}"
    return f"Weather in {city}: 70°F, clear (simulated)"


@tool
def get_forecast(city: str, days: int) -> str:
    """Get a weather forecast for a city for the specified number of days.

    Args:
        city: The city name to get a forecast for.
        days: Number of days to forecast (1-7).
    """
    days = max(1, min(days, 7))
    forecasts = []
    base_temps = {"new york": 72, "london": 58, "tokyo": 80, "paris": 65}
    base = base_temps.get(city.lower().strip(), 70)
    conditions = ["sunny", "partly cloudy", "cloudy", "rainy", "clear"]
    for i in range(days):
        temp = base + (i * 2 - days)
        cond = conditions[i % len(conditions)]
        forecasts.append(f"  Day {i + 1}: {temp}°F, {cond}")
    return f"Forecast for {city} ({days} days):\n" + "\n".join(forecasts)


# ---------------------------------------------------------------------------
# Math tools
# ---------------------------------------------------------------------------


@tool
def calculate(expression: str) -> str:
    """Evaluate a mathematical expression and return the result.

    Args:
        expression: A mathematical expression to evaluate, e.g. '2 + 2' or 'sqrt(16)'.
    """
    allowed = {
        k: getattr(math, k)
        for k in dir(math)
        if not k.startswith("_")
    }
    allowed.update({"abs": abs, "round": round, "min": min, "max": max})
    try:
        # Replace ^ with ** so users can write natural math
        safe_expr = expression.replace("^", "**")
        result = eval(safe_expr, {"__builtins__": {}}, allowed)  # noqa: S307
        return f"{expression} = {result}"
    except Exception as exc:
        return f"Error evaluating '{expression}': {exc}"


@tool
def convert_units(value: float, from_unit: str, to_unit: str) -> str:
    """Convert a value from one unit to another.

    Args:
        value: The numeric value to convert.
        from_unit: The source unit (e.g. 'km', 'miles', 'kg', 'lbs', 'celsius', 'fahrenheit').
        to_unit: The target unit.
    """
    conversions: dict[tuple[str, str], float] = {
        ("km", "miles"): 0.621371,
        ("miles", "km"): 1.60934,
        ("kg", "lbs"): 2.20462,
        ("lbs", "kg"): 0.453592,
        ("m", "ft"): 3.28084,
        ("ft", "m"): 0.3048,
        ("cm", "in"): 0.393701,
        ("in", "cm"): 2.54,
        ("l", "gal"): 0.264172,
        ("gal", "l"): 3.78541,
    }

    f = from_unit.lower().strip()
    t = to_unit.lower().strip()

    # Handle temperature separately
    if f in ("celsius", "c") and t in ("fahrenheit", "f"):
        result = value * 9 / 5 + 32
        return f"{value} {from_unit} = {result:.2f} {to_unit}"
    if f in ("fahrenheit", "f") and t in ("celsius", "c"):
        result = (value - 32) * 5 / 9
        return f"{value} {from_unit} = {result:.2f} {to_unit}"

    factor = conversions.get((f, t))
    if factor is not None:
        result = value * factor
        return f"{value} {from_unit} = {result:.4f} {to_unit}"

    pairs = ", ".join(f"{a}->{b}" for a, b in conversions)
    return f"Cannot convert from '{from_unit}' to '{to_unit}'. Supported: {pairs}, celsius<->fahrenheit"


# ---------------------------------------------------------------------------
# General tools
# ---------------------------------------------------------------------------


@tool
def get_time() -> str:
    """Get the current date and time in UTC."""
    now = datetime.now(UTC)
    return now.strftime("%Y-%m-%d %H:%M:%S UTC")


@tool
def search_web(query: str) -> str:
    """Search the web for information (stub).

    Args:
        query: The search query string.
    """
    return (
        f"[Simulated search results for '{query}']\n"
        f"1. Wikipedia article about {query}\n"
        f"2. Recent news: '{query}' trending topic\n"
        f"3. Stack Overflow discussion on {query}"
    )


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(title="Demo 3: Multi-Agent Router")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client = LLMClient(settings)

# Build per-agent tool registries
weather_tools = ToolRegistry()
weather_tools.register(get_weather)
weather_tools.register(get_forecast)

math_tools = ToolRegistry()
math_tools.register(calculate)
math_tools.register(convert_units)

general_tools = ToolRegistry()
general_tools.register(get_time)
general_tools.register(search_web)

# Create specialized agents
weather_agent = StreamingReactAgent(
    "weather",
    client=llm_client,
    config=AgentConfig(
        provider="openai",
        model="gpt-4o-mini",
        temperature=0.3,
        system_prompt=(
            "You are a weather specialist. Use your tools to answer weather and "
            "forecast questions. Be concise and informative."
        ),
    ),
    tools=weather_tools,
)

math_agent = StreamingReactAgent(
    "math",
    client=llm_client,
    config=AgentConfig(
        provider="openai",
        model="gpt-4o-mini",
        temperature=0.3,
        system_prompt=(
            "You are a math specialist. Use your tools to evaluate expressions "
            "and convert units. Show your work concisely."
        ),
    ),
    tools=math_tools,
)

general_agent = StreamingReactAgent(
    "general",
    client=llm_client,
    config=AgentConfig(
        provider="openai",
        model="gpt-4o-mini",
        temperature=0.3,
        system_prompt=(
            "You are a helpful general assistant. Use your tools when appropriate. "
            "Be concise and friendly."
        ),
    ),
    tools=general_tools,
)

# Router maps keywords to agent names
router = RuleRouter(
    rules={
        "weather": "weather",
        "forecast": "weather",
        "temperature": "weather",
        "calculate": "math",
        "math": "math",
        "compute": "math",
        "convert": "math",
    },
)

# Orchestrator with fallback
orchestrator = Orchestrator(
    agents={
        "weather": weather_agent,
        "math": math_agent,
        "general": general_agent,
    },
    router=router,
    fallback="general",
)

# Thread store
thread_store = ThreadStore("sqlite+aiosqlite:///demos/demo3.db")


@app.on_event("startup")
async def startup() -> None:
    await thread_store.init()


@app.on_event("shutdown")
async def shutdown() -> None:
    await thread_store.close()


# ---------------------------------------------------------------------------
# Request / response models
# ---------------------------------------------------------------------------


class ChatRequest(BaseModel):
    thread_id: str
    message: str


class CreateThreadResponse(BaseModel):
    thread_id: str


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    return HTMLResponse(content=chat_page(
        title="Demo 3: Multi-Agent Router",
        badge_text="3 Agents",
        badge_color="orange",
        port=8003,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Agents", "weather, math, general")
            + info_row("Router", "RuleRouter (keyword)")
            + info_row("Memory", "SQLite thread store")
            + info_row("Graph", "None")
        ),
    ))


@app.get("/api/threads")
async def list_threads() -> list[dict[str, str]]:
    ids = await thread_store.list_threads()
    return [{"thread_id": tid} for tid in ids]


@app.post("/api/threads")
async def create_thread() -> CreateThreadResponse:
    thread_id = str(uuid.uuid4())
    await thread_store.append(
        thread_id,
        role="system",
        content="Thread created.",
    )
    return CreateThreadResponse(thread_id=thread_id)


@app.get("/api/threads/{thread_id}/messages")
async def get_messages(thread_id: str) -> list[dict]:
    msgs = await thread_store.get_messages(thread_id)
    return [
        {
            "role": m.role,
            "content": m.content,
            "metadata": m.metadata or {},
            "created_at": m.created_at.isoformat() if m.created_at else "",
        }
        for m in msgs
        if m.role in ("user", "assistant")
    ]


@app.post("/api/chat")
async def chat(req: ChatRequest) -> StreamingResponse:
    thread_id = req.thread_id
    user_message = req.message

    # Persist user message
    await thread_store.append(thread_id, role="user", content=user_message)

    # Build conversation history from thread
    history_msgs = await thread_store.get_messages(thread_id)
    conversation_parts: list[str] = []
    for m in history_msgs:
        if m.role == "user":
            conversation_parts.append(f"User: {m.content}")
        elif m.role == "assistant":
            conversation_parts.append(f"Assistant: {m.content}")

    input_text = "\n".join(conversation_parts) if len(conversation_parts) > 1 else user_message

    async def event_stream():  # noqa: ANN202
        full_response = ""
        tool_calls: list[dict[str, str]] = []
        # Route on current message only (not full history) to avoid keyword bleed
        chosen = await router.route(user_message, orchestrator.agents)
        if chosen is None:
            chosen = orchestrator.fallback or "general"
        routed_to = chosen
        agent = orchestrator.agents[chosen]

        data = json.dumps({"type": "agent_start", "agent": chosen})
        yield f"data: {data}\n\n"

        try:
            async for event in agent.run_stream(input_text):
                if isinstance(event, TextDelta):
                    full_response += event.text
                    data = json.dumps({"type": "token", "token": event.text})
                    yield f"data: {data}\n\n"

                elif isinstance(event, ToolCallStart):
                    data = json.dumps({
                        "type": "tool_start",
                        "tool": event.tool_name,
                        "arguments": event.arguments,
                    })
                    yield f"data: {data}\n\n"

                elif isinstance(event, ToolCallEnd):
                    tool_calls.append({"tool": event.tool_name, "result": event.result})
                    data = json.dumps({
                        "type": "tool_end",
                        "tool": event.tool_name,
                        "result": event.result,
                    })
                    yield f"data: {data}\n\n"

                elif isinstance(event, AgentComplete):
                    meta = {
                        "routed_to": routed_to,
                        "iterations": event.iterations,
                        "tool_calls_made": event.tool_calls_made,
                        "tool_calls": tool_calls,
                    }
                    await thread_store.append(
                        thread_id, role="assistant", content=event.output,
                        metadata=meta,
                    )
                    data = json.dumps({"type": "complete", "metadata": meta})
                    yield f"data: {data}\n\n"

                elif isinstance(event, AgentError):
                    await thread_store.append(
                        thread_id, role="assistant", content=event.output,
                    )
                    data = json.dumps({"type": "error", "error": event.error})
                    yield f"data: {data}\n\n"

        except Exception as exc:
            data = json.dumps({"type": "error", "error": str(exc)})
            yield f"data: {data}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8003)
