"""Demo 4 — Full Stack KOAN: Multi-agent orchestrator with graph, feedback, evaluation, and SSE."""

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 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.graph import (
    ConfidenceCalibrator,
    DecisionEvaluator,
    EvaluationPolicy,
    FeedbackType,
    GraphAugmentedAgent,
    GraphConfig,
    PolicyManager,
    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},
}

TICKETS: dict[str, dict[str, str]] = {
    "ticket-2001": {"subject": "Login issues", "description": "Cannot log in since update", "status": "open"},
    "ticket-2002": {
        "subject": "Billing discrepancy",
        "description": "Charged twice for March",
        "status": "in_progress",
    },
    "ticket-2003": {"subject": "Feature request", "description": "Add dark mode support", "status": "closed"},
}

FAQ_ENTRIES: list[dict[str, str]] = [
    {"q": "How do I reset my password?", "a": "Go to Settings > Security > Reset Password."},
    {"q": "What are your support hours?", "a": "24/7 via chat, 9am-6pm EST via phone."},
    {"q": "How do I upgrade my plan?", "a": "Navigate to Settings > Subscription > Upgrade."},
    {"q": "How do I cancel my account?", "a": "Contact support or go to Settings > Account > Cancel."},
    {"q": "Where can I find my invoices?", "a": "Go to Billing > Invoice History."},
]

# ---------------------------------------------------------------------------
# Billing 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}"
    )


# ---------------------------------------------------------------------------
# Support tools
# ---------------------------------------------------------------------------


@tool
def search_faq(query: str) -> str:
    """Search the FAQ knowledge base for answers.

    Args:
        query: The search query string.
    """
    query_lower = query.lower()
    matches = [
        entry for entry in FAQ_ENTRIES
        if any(word in entry["q"].lower() for word in query_lower.split())
    ]
    if not matches:
        return f"No FAQ entries found for '{query}'."
    results = "\n".join(f"Q: {m['q']}\nA: {m['a']}" for m in matches[:3])
    return f"FAQ results for '{query}':\n{results}"


@tool
def create_ticket(subject: str, description: str) -> str:
    """Create a new support ticket.

    Args:
        subject: Brief subject line for the ticket.
        description: Detailed description of the issue.
    """
    ticket_id = f"ticket-{3000 + len(TICKETS)}"
    TICKETS[ticket_id] = {"subject": subject, "description": description, "status": "open"}
    return f"Ticket created: {ticket_id} — '{subject}' (status: open)"


@tool
def check_ticket_status(ticket_id: str) -> str:
    """Check the current status of a support ticket.

    Args:
        ticket_id: The ticket identifier, e.g. 'ticket-2001'.
    """
    ticket = TICKETS.get(ticket_id)
    if ticket is None:
        return f"Ticket {ticket_id} not found."
    return json.dumps({"ticket_id": ticket_id, **ticket})


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

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

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

app = FastAPI(title="Demo 4: Full Stack KOAN")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client: LLMClient | None = None
graph_client: GraphClient | None = None
calibrator: ConfidenceCalibrator | None = None
evaluator: DecisionEvaluator | None = None
orchestrator: Orchestrator | None = None
policy_manager: PolicyManager | None = None
demo_router: RuleRouter | None = None

# Agent configs
billing_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."
    ),
)

support_config = AgentConfig(
    provider="openai",
    model="gpt-4o-mini",
    temperature=0.3,
    system_prompt=(
        "You are a customer support agent. Help users find answers in the FAQ, "
        "create support tickets, and check ticket status. Be friendly and thorough."
    ),
)

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


@app.on_event("startup")
async def startup() -> None:
    global llm_client, graph_client, calibrator, evaluator, orchestrator, demo_router, policy_manager  # noqa: PLW0603

    llm_client = LLMClient(settings)
    graph_client = create_graph_client(graph_config)
    await graph_client.connect()

    calibrator = ConfidenceCalibrator(graph_client)
    evaluator = DecisionEvaluator(calibrator, policy=EvaluationPolicy())

    # PolicyManager for API endpoints — no seeding; policies evolve automatically
    policy_manager = PolicyManager(graph_client)

    # Build tool registries
    billing_tools = ToolRegistry()
    billing_tools.register(lookup_account)
    billing_tools.register(check_balance)
    billing_tools.register(apply_credit)

    support_tools = ToolRegistry()
    support_tools.register(search_faq)
    support_tools.register(create_ticket)
    support_tools.register(check_ticket_status)

    # Create graph-augmented agents
    billing_agent = GraphAugmentedAgent(
        "billing",
        client=llm_client,
        config=billing_config,
        tools=billing_tools,
        graph_client=graph_client,
        entity_ids=[],
        decision_type="billing_inquiry",
    )

    support_agent = GraphAugmentedAgent(
        "support",
        client=llm_client,
        config=support_config,
        tools=support_tools,
        graph_client=graph_client,
        entity_ids=[],
        decision_type="support_request",
    )

    # Router: keyword-based routing
    demo_router = RuleRouter(
        rules={
            "billing": "billing",
            "invoice": "billing",
            "account": "billing",
            "credit": "billing",
            "balance": "billing",
            "help": "support",
            "support": "support",
            "ticket": "support",
            "faq": "support",
            "issue": "support",
        },
    )

    orchestrator = Orchestrator(
        agents={"billing": billing_agent, "support": support_agent},
        router=demo_router,
        fallback="support",
    )


@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


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _extract_entity_ids(message: str) -> list[str]:
    """Extract acct-XXXX and ticket-XXXX identifiers from a message."""
    return re.findall(r"(?:acct|ticket)-\w+", message)


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


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    return HTMLResponse(content=chat_page(
        title="Demo 4: Full Stack KOAN",
        badge_text="Multi-Agent + Graph",
        badge_color="purple",
        port=8004,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Agents", "billing, support")
            + info_row("Router", "RuleRouter (keyword)")
            + info_row("Memory", "In-memory threads")
            + info_row("Graph", "FalkorDB (koan_demo4)")
            + info_row("Evaluation", "DecisionEvaluator")
            + info_row("Calibration", "/api/calibration")
        ),
    ))


@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:
    assert orchestrator is not None  # noqa: S101
    assert evaluator is not None  # noqa: S101
    assert calibrator is not None  # noqa: S101
    assert graph_client is not None  # noqa: S101

    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 and update agent entity_ids before routing
    entity_ids = _extract_entity_ids(user_message)

    # Update entity_ids on the graph-augmented agents so context is fetched for the right entities
    for agent in orchestrator.agents.values():
        if isinstance(agent, GraphAugmentedAgent):
            agent._entity_ids = entity_ids  # noqa: SLF001

    async def event_stream():  # noqa: ANN202
        full_response = ""
        run_id = ""
        graph_context_info: dict[str, Any] = {}
        tool_calls: list[dict[str, str]] = []

        # Route on current message only (not full history) to avoid keyword bleed
        chosen = await demo_router.route(user_message, orchestrator.agents)
        if chosen is None:
            chosen = orchestrator.fallback or "support"
        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):
                    run_id = event.metadata.get("run_id", "")
                    graph_context_info = event.metadata.get("graph_context", {})

                    # Evaluate the decision
                    decision_type = "billing_inquiry" if routed_to == "billing" else "support_request"
                    eval_result = await evaluator.evaluate(
                        agent_id=routed_to,
                        decision_type=decision_type,
                        entity_id=entity_ids[0] if entity_ids else None,
                        raw_confidence=1.0,
                    )

                    # Get calibration stats
                    cal = await calibrator.get_calibration(
                        agent_id=routed_to,
                        decision_type=decision_type,
                    )

                    meta: dict[str, Any] = {
                        "routed_to": routed_to,
                        "iterations": event.iterations,
                        "tool_calls_made": event.tool_calls_made,
                        "tool_calls": tool_calls,
                        "graph_context": graph_context_info,
                        "run_id": run_id,
                        "evaluation": {
                            "review_status": eval_result.review_status.value,
                            "raw_confidence": eval_result.raw_confidence,
                            "calibrated_confidence": eval_result.calibrated_confidence,
                            "reasons": eval_result.reasons,
                        },
                        "calibration": {
                            "factor": cal.factor,
                            "sample_count": cal.sample_count,
                            "positive_rate": cal.positive_rate,
                            "negative_rate": cal.negative_rate,
                            "fork_rate": cal.fork_rate,
                        },
                    }
                    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"}


@app.get("/api/calibration")
async def get_calibration() -> dict[str, Any]:
    assert calibrator is not None  # noqa: S101

    billing_cal = await calibrator.get_calibration(
        agent_id="billing",
        decision_type="billing_inquiry",
    )
    support_cal = await calibrator.get_calibration(
        agent_id="support",
        decision_type="support_request",
    )
    return {
        "billing": {
            "factor": billing_cal.factor,
            "sample_count": billing_cal.sample_count,
            "positive_rate": billing_cal.positive_rate,
            "negative_rate": billing_cal.negative_rate,
            "fork_rate": billing_cal.fork_rate,
        },
        "support": {
            "factor": support_cal.factor,
            "sample_count": support_cal.sample_count,
            "positive_rate": support_cal.positive_rate,
            "negative_rate": support_cal.negative_rate,
            "fork_rate": support_cal.fork_rate,
        },
    }


@app.get("/api/policies")
async def get_policies() -> list[dict[str, Any]]:
    """List all active policies in the graph."""
    assert policy_manager is not None  # noqa: S101
    return await policy_manager.list_policies()


@app.get("/api/policies/discover")
async def discover_patterns() -> list[dict[str, Any]]:
    """Discover recurring decision patterns that could become policies."""
    assert policy_manager is not None  # noqa: S101
    patterns = await policy_manager.discover_patterns(min_frequency=2, min_success_rate=0.5)
    return [
        {
            "decision_type": p.decision_type,
            "tool_sequence": p.tool_sequence,
            "frequency": p.frequency,
            "success_rate": p.success_rate,
            "suggested_name": p.suggested_policy_name,
            "description": p.description,
        }
        for p in patterns
    ]


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

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