"""The assistant: its provider, and the conversation with it.

The provider is instance-wide and superuser-only to set, because the key it
holds is spent by everyone who chats. The conversation is per person and runs
as them: every tool call carries their own token, so the assistant can do
nothing they could not do in the browser.
"""

from __future__ import annotations

import json
import logging
import uuid
from collections.abc import AsyncIterator
from typing import Any, Literal

import httpx
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from fluksio.ai import agent, client
from fluksio.ai import config as ai_config
from fluksio.api.deps import (
    CurrentUser,
    FlowControllerDep,
    TokenDep,
    get_current_active_superuser,
)
from fluksio.core.config import settings
from fluksio.flow.messages import qualify
from fluksio.flow.schemas import NAME_PATTERN, NodeDef
from fluksio.mcp import server as tools

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/ai", tags=["ai"])

#: Whose turn is running: a second message while one is answered is refused
#: rather than interleaved into the same transcript.
_busy: set[uuid.UUID] = set()


class AIConfigPublic(BaseModel):
    api_base: str
    model: str
    max_rounds: int
    api_key_set: bool
    api_key_masked: str | None


class AIConfigUpdate(BaseModel):
    api_base: str = Field(min_length=1, max_length=512)
    #: Blank or absent keeps the stored key.
    api_key: str | None = Field(default=None, max_length=512)
    model: str = Field(default="", max_length=128)
    max_rounds: int = Field(default=12, ge=1, le=50)


class AIModelsBody(BaseModel):
    api_base: str = Field(min_length=1, max_length=512)
    #: Absent tries the stored key, so a saved provider can be re-listed.
    api_key: str | None = Field(default=None, max_length=512)


class AIStatus(BaseModel):
    available: bool
    model: str | None


class ChatItem(BaseModel):
    role: str
    content: str | None = None
    tool_calls: list[dict[str, Any]] | None = None
    #: On a tool's answer: whether it worked, and the error if not.
    ok: bool | None = None
    error: str | None = None


type WatchKind = Literal["partial", "text", "reply", "actions"]


class ControlWatch(BaseModel):
    """A message the conversation shows, as one kind of line."""

    msg: str
    kind: WatchKind
    #: The node whose values on it count. Empty counts every value, as the
    #: prompt does: a command typed here or spoken to a device.
    node: str = ""


class ControlNode(BaseModel):
    """A conversation the chat can hold: where a command goes, and what to show."""

    #: Qualified, ``flow.node``.
    id: str
    flow: str
    title: str
    prompt_msg: str
    watch: list[ControlWatch]


class TurnBody(BaseModel):
    message: str = Field(min_length=1, max_length=20_000)
    #: The flow open in the editor, when the chat is opened from one.
    flow: str | None = Field(default=None, pattern=NAME_PATTERN.pattern)


def _public(config: ai_config.AIConfig | None) -> AIConfigPublic:
    if config is None:
        return AIConfigPublic(
            api_base="",
            model="",
            max_rounds=12,
            api_key_set=False,
            api_key_masked=None,
        )
    return AIConfigPublic(**config.public())


@router.get(
    "/config",
    dependencies=[Depends(get_current_active_superuser)],
    response_model=AIConfigPublic,
)
def read_ai_config() -> Any:
    return _public(ai_config.load())


@router.put(
    "/config",
    dependencies=[Depends(get_current_active_superuser)],
    response_model=AIConfigPublic,
)
def update_ai_config(body: AIConfigUpdate) -> Any:
    try:
        client.check_base(body.api_base)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    current = ai_config.load()
    key = body.api_key or (current.api_key if current else "")
    config = ai_config.AIConfig(
        api_base=body.api_base.rstrip("/"),
        api_key=key,
        model=body.model.strip(),
        max_rounds=body.max_rounds,
    )
    ai_config.save(config)
    return _public(config)


@router.delete("/config", dependencies=[Depends(get_current_active_superuser)])
def delete_ai_config() -> Any:
    ai_config.delete()
    return _public(None)


@router.post(
    "/config/models",
    dependencies=[Depends(get_current_active_superuser)],
    response_model=list[str],
)
async def list_ai_models(body: AIModelsBody) -> Any:
    """What the provider offers, tried with the key given or the one stored."""
    try:
        client.check_base(body.api_base)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    current = ai_config.load()
    key = body.api_key or (current.api_key if current else "")
    try:
        return await client.list_models(body.api_base.rstrip("/"), key)
    except (client.ProviderError, httpx.HTTPError) as exc:
        raise HTTPException(
            status_code=400, detail=f"The provider answered {exc}"
        ) from exc


@router.get("/status", response_model=AIStatus)
def read_ai_status(_: CurrentUser) -> Any:
    """Whether there is an assistant to show. Readable by anyone signed in."""
    config = ai_config.load()
    if config is None or not config.ready:
        return AIStatus(available=False, model=None)
    return AIStatus(available=True, model=config.model)


#: The node types a command can be typed to.
_TALKERS = ("wren", "assistant")


def _wired(flow: str, node: NodeDef, port: str) -> str:
    """The message a node's port is wired to, qualified, or empty."""
    for spec in [*node.requires, *node.provides]:
        if spec.port == port and spec.name:
            return qualify(flow, spec.name)
    return ""


@router.get("/control", response_model=list[ControlNode])
def read_control_nodes(_: CurrentUser, controller: FlowControllerDep) -> Any:
    """The conversations of the running live flows, for the chat's Control tab.

    One per wren or assistant node a command reaches first. A node whose
    prompt is another's ``unhandled`` answers in that one's conversation,
    after it. The conversation shows the prompt, the words heard so far (the
    ``partial`` of a ``stt`` node writing the prompt), and each node's
    ``reply`` and ``actions``, so typed and spoken commands are one.
    """
    found = []
    for flow in controller.store.read_all():
        if flow.mode != "live" or flow.name in controller.disabled:
            continue
        talkers = [
            node
            for node in flow.nodes
            if node.type in _TALKERS and _wired(flow.name, node, "prompt")
        ]
        passed = {_wired(flow.name, node, "unhandled") for node in talkers}
        for node in talkers:
            prompt = _wired(flow.name, node, "prompt")
            if prompt in passed:
                continue
            watch = [ControlWatch(msg=prompt, kind="text")]
            watch += [
                ControlWatch(msg=partial, kind="partial", node=f"{flow.name}.{stt.id}")
                for stt in flow.nodes
                if stt.type == "stt"
                and _wired(flow.name, stt, "text") == prompt
                and (partial := _wired(flow.name, stt, "partial"))
            ]
            chain = [node]
            for talker in chain:
                unhandled = _wired(flow.name, talker, "unhandled")
                chain += [
                    other
                    for other in talkers
                    if unhandled
                    and _wired(flow.name, other, "prompt") == unhandled
                    and other not in chain
                ]
            answers: tuple[WatchKind, ...] = ("reply", "actions")
            watch += [
                ControlWatch(msg=msg, kind=kind, node=f"{flow.name}.{talker.id}")
                for talker in chain
                for kind in answers
                if (msg := _wired(flow.name, talker, kind))
            ]
            if any(w.kind in answers for w in watch):
                found.append(
                    ControlNode(
                        id=f"{flow.name}.{node.id}",
                        flow=flow.name,
                        title=f"{flow.title or flow.name} · {node.title or node.id}",
                        prompt_msg=prompt,
                        watch=watch,
                    )
                )
    return found


def _chat_file(user_id: uuid.UUID) -> Any:
    return settings.CHATS_DIR / f"{user_id}.json"


def _load_chat(user_id: uuid.UUID) -> list[dict[str, Any]]:
    try:
        messages: list[dict[str, Any]] = json.loads(_chat_file(user_id).read_text())
    except (OSError, ValueError):
        return []
    return messages


def _save_chat(user_id: uuid.UUID, messages: list[dict[str, Any]]) -> None:
    path = _chat_file(user_id)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(messages))


@router.get("/chat", response_model=list[ChatItem])
def read_chat(user: CurrentUser) -> Any:
    """The transcript, with each tool's answer reduced to whether it worked."""
    items = []
    for message in _load_chat(user.id):
        if message.get("role") == "tool":
            items.append(
                ChatItem(role="tool", **agent.summarize(message.get("content")))
            )
        else:
            items.append(
                ChatItem(
                    role=message["role"],
                    content=message.get("content"),
                    tool_calls=message.get("tool_calls"),
                )
            )
    return items


@router.delete("/chat", response_model=list[ChatItem])
def clear_chat(user: CurrentUser) -> Any:
    _chat_file(user.id).unlink(missing_ok=True)
    return []


def _sse(event: dict[str, Any]) -> str:
    return f"data: {json.dumps(event)}\n\n"


@router.post("/chat")
def chat(
    body: TurnBody, request: Request, user: CurrentUser, token: TokenDep
) -> StreamingResponse:
    """One turn, streamed as server-sent events.

    ``delta`` as the model writes, ``message`` with a round's whole text,
    ``tool_call`` and ``tool_result`` as it works, ``error`` when something
    stops it, and ``done`` last either way.
    """
    config = ai_config.load()
    if config is None or not config.ready:
        # 404 rather than 403: an unconfigured instance says nothing about it.
        raise HTTPException(status_code=404, detail="The assistant is not set up")
    if user.id in _busy:
        raise HTTPException(status_code=409, detail="Still answering the last one")
    tools.connect(request.app)
    messages = _load_chat(user.id)

    async def generate() -> AsyncIterator[str]:
        _busy.add(user.id)
        try:
            yield _sse({"type": "start"})
            async for event in agent.run_turn(
                config,
                messages,
                body.message,
                tools=await agent.tool_specs(),
                system=agent.system_prompt(body.flow),
                call_tool=agent.mcp_caller(f"Bearer {token}"),
            ):
                yield _sse(event)
        except Exception as exc:  # noqa: BLE001 - the stream is the error channel
            logger.exception("The assistant's turn failed")
            yield _sse({"type": "error", "detail": f"{type(exc).__name__}: {exc}"})
        finally:
            # What happened so far, even if the browser left part-way.
            _save_chat(user.id, messages)
            _busy.discard(user.id)
        yield _sse({"type": "done"})

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )