"""Demo 1 — Simple Chat Agent with SQLite thread memory and SSE streaming."""

from __future__ import annotations

import json
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.events import AgentComplete, AgentError, TextDelta, ToolCallEnd, ToolCallStart
from koan.agents.stream import StreamingReactAgent
from koan.agents.types import AgentConfig
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()

# ---------------------------------------------------------------------------
# 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.
    """
    # Stub implementation for demo purposes
    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",
    }
    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 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)'.
    """
    import math  # noqa: F811

    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 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")


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

app = FastAPI(title="Demo 1: Simple Chat Agent")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client = LLMClient(settings)

registry = ToolRegistry()
registry.register(get_weather)
registry.register(calculate)
registry.register(get_time)

thread_store = ThreadStore("sqlite+aiosqlite:///demos/demo1.db")

agent_config = AgentConfig(
    provider="openai",
    model="gpt-4o-mini",
    temperature=0.3,
    system_prompt="You are a helpful assistant. Use your tools when needed. Be concise.",
)


@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 1: Simple Chat Agent",
        badge_text="No Graph",
        badge_color="gray",
        port=8001,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Tools", "get_weather, calculate, get_time")
            + 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())
    # Store a system-level marker so the thread shows up in list_threads
    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}")

    # Use the full conversation as context for the agent
    input_text = "\n".join(conversation_parts) if len(conversation_parts) > 1 else user_message

    agent = StreamingReactAgent(
        "chat-agent",
        client=llm_client,
        config=agent_config,
        tools=registry,
    )

    async def event_stream():  # noqa: ANN202
        full_response = ""
        tool_calls: list[dict[str, str]] = []
        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 = {
                        "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=8001)
