"""MCP tools over one or more T3 Code servers.

`t3_wait` blocks on purpose: Claude Code has no MCP Tasks support and moves
any tool call still running after two minutes into a background task by
itself, so a plain long poll is the cheapest correct shape here.
"""

from __future__ import annotations

import asyncio
import contextlib
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Annotated, Any

import httpx
import websockets
from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from pydantic import Field
from starlette.responses import JSONResponse

from t3code_mcp import t3, views
from t3code_mcp.config import Machine, ModelSpec
from t3code_mcp.t3 import T3Client, T3Error
from t3code_mcp.views import BUSY
from t3code_mcp.watch import Tracked, Tracker

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Callable, Coroutine

    from starlette.requests import Request

    from t3code_mcp.t3 import RpcCall

TITLE_MAX = 80
WAIT_MAX = 3600.0


@dataclass
class Registry:
    machines: dict[str, Machine]
    transport: httpx.AsyncBaseTransport | None = None
    rpc: Callable[[str], RpcCall] | None = None
    clients: dict[str, T3Client] = field(default_factory=dict)
    threads: dict[str, str] = field(default_factory=dict)
    tracker: Tracker = field(default_factory=Tracker)

    def machine(self, name: str) -> Machine:
        try:
            return self.machines[name]
        except KeyError:
            msg = f"unknown machine {name!r}; known: {', '.join(self.machines)}"
            raise ToolError(msg) from None

    def client(self, name: str) -> T3Client:
        machine = self.machine(name)
        if name not in self.clients:
            self.clients[name] = T3Client(
                machine.url,
                machine.token,
                transport=self.transport,
                rpc=self.rpc(name) if self.rpc else None,
            )
        return self.clients[name]

    async def locate(self, thread_id: str) -> tuple[str, dict[str, Any]]:
        known = self.threads.get(thread_id) or (
            self.tracker.threads[thread_id].machine
            if thread_id in self.tracker.threads
            else None
        )
        names = [known] if known in self.machines else list(self.machines)
        for name in names:
            detail = await self.client(name).thread(thread_id)
            if detail is not None:
                self.threads[thread_id] = name
                return name, detail
        msg = f"thread {thread_id} not found on {', '.join(names)}"
        raise ToolError(msg)

    async def aclose(self) -> None:
        for client in self.clients.values():
            await client.aclose()
        self.clients.clear()


def _title(prompt: str, title: str | None) -> str:
    text = (title or prompt).strip().splitlines()[0].strip()
    if len(text) > TITLE_MAX:
        text = text[:TITLE_MAX].rsplit(" ", 1)[0] + "…"
    return text or "t3code-mcp"


def _find_project(
    machine: Machine, projects: list[dict[str, Any]], key: str
) -> dict[str, Any]:
    allowed = [p for p in projects if machine.allows(p)]
    for project in allowed:
        if key in (project["id"], project["title"], project["workspaceRoot"]):
            return project
    names = ", ".join(p["title"] for p in allowed) or "(none allowlisted)"
    msg = f"project {key!r} is not available on {machine.name}; allowed: {names}"
    raise ToolError(msg)


def _project_view(
    project: dict[str, Any], threads: list[dict[str, Any]]
) -> dict[str, Any]:
    own = [
        t
        for t in threads
        if t["projectId"] == project["id"] and not t.get("archivedAt")
    ]
    running = sum(1 for t in own if views.shell_state(t) in BUSY)
    default = project.get("defaultModelSelection")
    return {
        "id": project["id"],
        "title": project["title"],
        "path": project["workspaceRoot"],
        "default_model": str(ModelSpec.from_selection(default)) if default else None,
        "threads": len(own),
        "running": running,
    }


def _watch_view(thread_id: str, tracked: Tracked) -> dict[str, Any]:
    return {
        "thread_id": thread_id,
        "machine": tracked.machine,
        "project": tracked.project,
        "title": tracked.title,
        "state": tracked.state,
    }


def build_server(
    registry: Registry,
    *,
    poll: float = 5.0,
    background: Callable[[], Coroutine[Any, Any, None]] | None = None,
) -> FastMCP:
    @contextlib.asynccontextmanager
    async def lifespan(_: FastMCP) -> AsyncIterator[dict[str, Any]]:
        task = asyncio.create_task(background()) if background is not None else None
        try:
            yield {}
        finally:
            if task is not None:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
            await registry.aclose()

    mcp = FastMCP(
        "t3code",
        instructions=(
            "Coding threads on Бобёр's machines through T3 Code. Flow: t3_machines → "
            "t3_projects(machine) → t3_dispatch, then either t3_wait (block) or just "
            "carry on: the gateway injects an event when a dispatched thread finishes, "
            "fails or asks a question (answer with t3_answer; t3_watch turns those "
            "events off for a thread). Prompts run "
            "autonomously (full-access, no approvals): say what to change, what to "
            "check, what not to touch, and that the answer should end with a report."
        ),
        lifespan=lifespan,
    )

    @mcp.custom_route("/healthz", methods=["GET"])
    async def healthz(_: Request) -> JSONResponse:
        return JSONResponse({"status": "ok", "machines": list(registry.machines)})

    @mcp.tool
    async def t3_machines() -> list[dict[str, Any]]:
        """Machines with a T3 Code server: reachable now or not, project allowlist."""
        result = []
        for name, machine in registry.machines.items():
            entry: dict[str, Any] = {
                "name": name,
                "url": machine.url,
                "projects": machine.projects,
            }
            try:
                descriptor = await registry.client(name).descriptor()
            except (httpx.HTTPError, T3Error) as exc:
                entry.update(online=False, error=str(exc)[:200])
            else:
                entry.update(
                    online=True,
                    label=descriptor.get("label"),
                    version=descriptor.get("serverVersion"),
                    os=(descriptor.get("platform") or {}).get("os"),
                )
            result.append(entry)
        return result

    @mcp.tool
    async def t3_projects(
        machine: Annotated[str, Field(description="Machine name from t3_machines")],
    ) -> list[dict[str, Any]]:
        """Allowlisted projects on a machine: default model, thread counts."""
        spec = registry.machine(machine)
        shell = await _call(registry.client(machine).shell())
        return [
            _project_view(p, shell["threads"])
            for p in shell["projects"]
            if spec.allows(p)
        ]

    @mcp.tool
    async def t3_dispatch(
        machine: Annotated[str, Field(description="Machine name from t3_machines")],
        project: Annotated[
            str, Field(description="Project title, path or id from t3_projects")
        ],
        prompt: Annotated[
            str,
            Field(
                description="The task for the coding agent, complete and self-contained"
            ),
        ],
        title: Annotated[
            str | None,
            Field(description="Thread title; default - first line of the prompt"),
        ] = None,
        model: Annotated[
            str | None,
            Field(
                description=(
                    "`instance/model`, e.g. claudeAgent/claude-opus-5; "
                    "default from config or the project"
                )
            ),
        ] = None,
        thread_id: Annotated[
            str | None,
            Field(description="Continue this existing thread instead of creating one"),
        ] = None,
        watch: Annotated[  # noqa: FBT002 - MCP arguments are always named
            bool,
            Field(
                description=(
                    "Report this thread back as an event; false for a thread you "
                    "read yourself. Every dispatch resets the flag - see t3_watch"
                )
            ),
        ] = True,
    ) -> dict[str, Any]:
        """Start a coding thread on a machine (or follow up in one); returns at once.

        The thread runs in the background: t3_wait blocks until it finishes,
        t3_thread peeks without blocking.
        """
        text = prompt.strip()
        if not text:
            msg = "prompt is empty"
            raise ToolError(msg)
        spec = registry.machine(machine)
        client = registry.client(machine)
        shell = await _call(client.shell())
        target = _find_project(spec, shell["projects"], project)
        if model:
            selection = ModelSpec.parse(
                model, spec.model.options if spec.model else None
            )
        elif spec.model:
            selection = spec.model
        elif target.get("defaultModelSelection"):
            selection = ModelSpec.from_selection(target["defaultModelSelection"])
        else:
            msg = f"no model: set `model` in config for {machine} or pass model="
            raise ToolError(msg)
        heading = _title(text, title)
        if thread_id:
            existing = next((t for t in shell["threads"] if t["id"] == thread_id), None)
            if existing is None or existing["projectId"] != target["id"]:
                msg = f"thread {thread_id} is not in {target['title']} on {machine}"
                raise ToolError(msg)
            if views.shell_state(existing) in BUSY:
                msg = f"thread {thread_id} is running; t3_wait or t3_interrupt first"
                raise ToolError(msg)
            tid = thread_id
            await _call(
                client.dispatch(t3.message_dispatch(tid, text, selection.selection()))
            )
        else:
            tid = t3.new_id()
            await _call(
                client.dispatch(
                    t3.thread_create(tid, target["id"], heading, selection.selection())
                )
            )
            await _call(
                client.dispatch(t3.message_dispatch(tid, text, title_seed=heading))
            )
        registry.threads[tid] = machine
        registry.tracker.track(tid, machine, target["title"], heading, watch=watch)
        return {
            "thread_id": tid,
            "machine": machine,
            "project": target["title"],
            "path": target["workspaceRoot"],
            "title": heading,
            "model": str(selection),
            "state": "running",
            "watch": watch,
        }

    @mcp.tool
    async def t3_watch(
        thread_id: Annotated[
            str | None,
            Field(description="Thread to subscribe or unsubscribe; omit to only list"),
        ] = None,
        enabled: Annotated[
            bool | None,
            Field(description="True - report the thread back again, false - stop"),
        ] = None,
    ) -> dict[str, Any]:
        """Which dispatched threads report back, and turn that on or off per thread.

        t3_dispatch subscribes its thread, so its completion, failure and
        questions arrive here as events. Unsubscribe a thread someone reads
        directly: it keeps running, it just stops reporting. Called with no
        arguments this only lists.
        """
        if thread_id is not None:
            if enabled is None:
                msg = "pass enabled=true or enabled=false along with thread_id"
                raise ToolError(msg)
            try:
                registry.tracker.set_watch(thread_id, enabled=enabled)
            except KeyError:
                msg = f"thread {thread_id} was not dispatched from here; nothing to "
                raise ToolError(msg + ("watch" if enabled else "unwatch")) from None
        threads = registry.tracker.threads
        return {
            "watched": [_watch_view(k, v) for k, v in threads.items() if v.watch],
            "unwatched": [_watch_view(k, v) for k, v in threads.items() if not v.watch],
        }

    @mcp.tool
    async def t3_thread(
        thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
        turns: Annotated[
            int,
            Field(
                description="How many recent runs (user turns) to include", ge=1, le=20
            ),
        ] = 2,
    ) -> dict[str, Any]:
        """State of a thread without waiting.

        Turn state, last messages, tools run, files changed, pending questions.
        """
        machine, detail = await registry.locate(thread_id)
        shell = await _call(registry.client(machine).shell())
        summary = views.summary(
            machine, detail, project=views.project_title(shell, detail), turns=turns
        )
        summary["pending"] = views.pending_detail(detail["projection"])
        return summary

    @mcp.tool
    async def t3_wait(
        thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
        timeout: Annotated[  # noqa: ASYNC109 - the tool's contract, not asyncio's
            float,
            Field(
                description="Seconds to wait before giving the current state back",
                ge=5,
                le=WAIT_MAX,
            ),
        ] = 900,
    ) -> dict[str, Any]:
        """Block until the turn finishes, asks a question, or timeout runs out.

        Ends on completed / interrupted / error, on a pending approval or
        user-input request, or after `timeout` seconds; returns the t3_thread view.
        """
        machine, detail = await registry.locate(thread_id)
        client = registry.client(machine)
        started = time.monotonic()
        project = views.project_title(await _call(client.shell()), detail)
        while True:
            projection = detail["projection"]
            state = views.state(projection)
            pending = views.pending_detail(projection)
            stuck = views.latest_run(projection) is None and views.error(projection)
            done = state not in BUSY or pending is not None or stuck
            timed_out = time.monotonic() - started >= timeout
            if done or timed_out:
                summary = views.summary(machine, detail, project=project)
                summary["pending"] = pending
                summary["timed_out"] = timed_out and not done
                summary["waited"] = round(time.monotonic() - started, 1)
                return summary
            await asyncio.sleep(poll)
            refreshed = await _call(client.thread(thread_id))
            if refreshed is None:
                msg = f"thread {thread_id} vanished from {machine}"
                raise ToolError(msg)
            detail = refreshed

    @mcp.tool
    async def t3_answer(
        thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
        answers: Annotated[
            dict[str, str | list[str]],
            Field(
                description=(
                    "question id → chosen option label (list for multi_select), "
                    "ids and options from the `question` of t3_thread / the event"
                )
            ),
        ],
    ) -> dict[str, Any]:
        """Answer the pending question of a thread; its turn resumes."""
        machine, detail = await registry.locate(thread_id)
        pending = views.question(detail["projection"])
        if pending is None or not pending.get("request_id"):
            msg = f"thread {thread_id} has no pending question"
            raise ToolError(msg)
        known = {q["id"] for q in pending["questions"]}
        unknown = set(answers) - known
        if unknown:
            msg = f"unknown question ids {sorted(unknown)}; expected {sorted(known)}"
            raise ToolError(msg)
        await _call(
            registry.client(machine).dispatch(
                t3.runtime_request_respond(
                    thread_id, pending["request_id"], dict(answers)
                )
            )
        )
        return {
            "thread_id": thread_id,
            "machine": machine,
            "answered": sorted(answers),
            "state": "running",
        }

    @mcp.tool
    async def t3_interrupt(
        thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
    ) -> dict[str, Any]:
        """Interrupt the running turn; continue later via t3_dispatch(thread_id=...)."""
        machine, detail = await registry.locate(thread_id)
        latest = views.latest_run(detail["projection"])
        if latest is None or latest["status"] not in BUSY:
            return {
                "thread_id": thread_id,
                "machine": machine,
                "state": latest["status"] if latest else "idle",
            }
        await _call(
            registry.client(machine).dispatch(t3.run_interrupt(thread_id, latest["id"]))
        )
        return {"thread_id": thread_id, "machine": machine, "state": "interrupting"}

    return mcp


async def _call(awaitable: Any) -> Any:
    try:
        return await awaitable
    except T3Error as exc:
        raise ToolError(str(exc)) from exc
    except (httpx.HTTPError, websockets.WebSocketException, OSError) as exc:
        msg = f"T3 unreachable: {exc}"
        raise ToolError(msg) from exc