"""Demo 2 — Graph-Augmented Agent with FalkorDB, feedback, and SSE streaming."""

from __future__ import annotations

import json
import re
import sys
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Any

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.types import AgentConfig
from koan.graph import (
    FeedbackType,
    GraphAugmentedAgent,
    GraphConfig,
    create_graph_client,
    record_feedback,
)
from koan.llm.client import LLMClient
from koan.llm.config import ProviderSettings
from koan.tools import ToolRegistry, tool

if TYPE_CHECKING:
    from koan.graph.client import GraphClient

load_dotenv()

# ---------------------------------------------------------------------------
# Stub data
# ---------------------------------------------------------------------------

ACCOUNTS: dict[str, dict[str, Any]] = {
    "acct-1001": {"name": "Alice Johnson", "plan": "Pro", "status": "active", "balance": 142.50},
    "acct-1002": {"name": "Bob Smith", "plan": "Starter", "status": "active", "balance": 27.00},
    "acct-1003": {"name": "Carol Davis", "plan": "Enterprise", "status": "suspended", "balance": 980.00},
}

# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@tool
def lookup_account(account_id: str) -> str:
    """Look up account details by account ID.

    Args:
        account_id: The account identifier, e.g. 'acct-1001'.
    """
    acct = ACCOUNTS.get(account_id)
    if acct is None:
        return f"Account {account_id} not found."
    return json.dumps({"account_id": account_id, **acct})


@tool
def check_balance(account_id: str) -> str:
    """Check the current balance for an account.

    Args:
        account_id: The account identifier, e.g. 'acct-1001'.
    """
    acct = ACCOUNTS.get(account_id)
    if acct is None:
        return f"Account {account_id} not found."
    return f"Account {account_id} balance: ${acct['balance']:.2f} (plan: {acct['plan']})"


@tool
def apply_credit(account_id: str, amount: float) -> str:
    """Apply a credit (discount) to an account balance.

    Args:
        account_id: The account identifier, e.g. 'acct-1001'.
        amount: The dollar amount to credit.
    """
    acct = ACCOUNTS.get(account_id)
    if acct is None:
        return f"Account {account_id} not found."
    old_balance = acct["balance"]
    acct["balance"] = max(0.0, old_balance - amount)
    return (
        f"Applied ${amount:.2f} credit to {account_id}. "
        f"Previous balance: ${old_balance:.2f}, new balance: ${acct['balance']:.2f}"
    )


# ---------------------------------------------------------------------------
# In-memory thread storage
# ---------------------------------------------------------------------------

threads: dict[str, list[dict[str, str]]] = {}

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

app = FastAPI(title="Demo 2: Graph-Augmented Agent")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client = LLMClient(settings)

registry = ToolRegistry()
registry.register(lookup_account)
registry.register(check_balance)
registry.register(apply_credit)

graph_config = GraphConfig(
    provider="falkordb",
    uri="redis://localhost:6379",
    database="koan_demo2",
)

graph_client: GraphClient | None = None

agent_config = AgentConfig(
    provider="openai",
    model="gpt-4o-mini",
    temperature=0.3,
    system_prompt=(
        "You are a billing support agent. Help customers with account lookups, "
        "balance inquiries, and applying credits. Be concise and professional. "
        "Use the provided account ID directly — do not ask the user to confirm it."
    ),
)


@app.on_event("startup")
async def startup() -> None:
    global graph_client  # noqa: PLW0603
    graph_client = create_graph_client(graph_config)
    await graph_client.connect()
    # No policy seeding — policies emerge automatically from observed patterns


@app.on_event("shutdown")
async def shutdown() -> None:
    if graph_client is not None:
        await graph_client.close()


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


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


class FeedbackRequest(BaseModel):
    run_id: str
    satisfied: bool


class CreateThreadResponse(BaseModel):
    thread_id: str


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


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    return HTMLResponse(content=chat_page(
        title="Demo 2: Graph-Augmented Agent",
        badge_text="FalkorDB",
        badge_color="green",
        port=8002,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Tools", "lookup_account, check_balance, apply_credit")
            + info_row("Memory", "In-memory threads")
            + info_row("Graph", "FalkorDB (koan_demo2)")
        ),
    ))


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


@app.post("/api/threads")
async def create_thread() -> CreateThreadResponse:
    thread_id = str(uuid.uuid4())
    threads[thread_id] = []
    return CreateThreadResponse(thread_id=thread_id)


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


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

    # Ensure thread exists
    if thread_id not in threads:
        threads[thread_id] = []

    # Persist user message
    threads[thread_id].append({"role": "user", "content": user_message})

    # Build conversation history
    conversation_parts: list[str] = []
    for m in threads[thread_id]:
        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

    # Extract entity IDs (acct-XXXX) from message
    entity_ids = re.findall(r"acct-\w+", user_message)

    assert graph_client is not None  # noqa: S101

    agent = GraphAugmentedAgent(
        "billing-agent",
        client=llm_client,
        config=agent_config,
        tools=registry,
        graph_client=graph_client,
        entity_ids=entity_ids,
        decision_type="classification",
    )

    async def event_stream():  # noqa: ANN202
        full_response = ""
        run_id = ""
        graph_context_info: dict[str, Any] = {}
        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):
                    run_id = event.metadata.get("run_id", "")
                    graph_context_info = event.metadata.get("graph_context", {})
                    meta: dict[str, Any] = {
                        "iterations": event.iterations,
                        "tool_calls_made": event.tool_calls_made,
                        "graph_context": graph_context_info,
                        "tool_calls": tool_calls,
                        "run_id": run_id,
                    }
                    threads[thread_id].append({
                        "role": "assistant",
                        "content": event.output,
                        "metadata": meta,
                    })
                    data = json.dumps({
                        "type": "complete",
                        "run_id": run_id,
                        "metadata": meta,
                    })
                    yield f"data: {data}\n\n"

                elif isinstance(event, AgentError):
                    threads[thread_id].append({"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")


@app.post("/api/feedback")
async def submit_feedback(req: FeedbackRequest) -> dict[str, str]:
    assert graph_client is not None  # noqa: S101

    score = 1.0 if req.satisfied else -1.0
    ftype = FeedbackType.THUMBS_UP if req.satisfied else FeedbackType.THUMBS_DOWN

    feedback_id = await record_feedback(
        graph_client,
        run_id=req.run_id,
        feedback_type=ftype,
        score=score,
    )
    return {"feedback_id": feedback_id, "status": "recorded"}


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

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