"""HTTP 上一对多 Runtime 连接（与 ``aicd host`` 同款 dispatch）。"""

from __future__ import annotations

import asyncio
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from aicd.host_dispatch import dispatch_host_request
from aicd.runtime import RuntimeContext, build_runtime


@dataclass
class _ConnectionEntry:
    ctx: RuntimeContext
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)


class RuntimeConnectionManager:
    """每个 ``connection_id`` 独占一个 ``RuntimeContext``（隔离内存态与插件加载）。"""

    def __init__(
        self,
        *,
        home: Path,
        max_connections: int = 16,
        skip_setup: bool = True,
        setup_tty_relaxed: bool = True,
        no_plugins: bool = False,
        plugin_policy: str | None = None,
    ) -> None:
        self._home = home
        self._max = max(1, int(max_connections))
        self._skip_setup = skip_setup
        self._setup_tty_relaxed = setup_tty_relaxed
        self._no_plugins = no_plugins
        self._plugin_policy = plugin_policy
        self._entries: dict[str, _ConnectionEntry] = {}
        self._global_lock = asyncio.Lock()

    @property
    def connection_count(self) -> int:
        return len(self._entries)

    def find_runtime_for_chat_session(
        self, chat_session_id: str
    ) -> RuntimeContext | None:
        """若有 HTTP connection 绑定了该 ``chat_session_id``，返回对应运行时装载（含 LLM 客户端）。"""
        sid = (chat_session_id or "").strip()
        if not sid:
            return None
        for ent in self._entries.values():
            rsid = getattr(ent.ctx.router, "_chat_session_id", None)
            if rsid is not None and str(rsid).strip() == sid:
                return ent.ctx
        return None

    async def shutdown_all(self) -> None:
        async with self._global_lock:
            items = list(self._entries.items())
            self._entries.clear()
        for _cid, ent in items:
            await ent.ctx.shutdown_graceful()

    async def create(
        self,
        *,
        chat_session_id: str | None = None,
        no_plugins: bool | None = None,
        plugin_policy: str | None = None,
    ) -> tuple[str, str]:
        np = self._no_plugins if no_plugins is None else bool(no_plugins)
        pp = self._plugin_policy if plugin_policy is None else plugin_policy
        async with self._global_lock:
            if len(self._entries) >= self._max:
                raise RuntimeError("max_connections_reached")
            cid = str(uuid.uuid4())
            ctx = await build_runtime(
                home=self._home,
                session_id=chat_session_id,
                skip_first_setup=self._skip_setup,
                setup_tty_relaxed=self._setup_tty_relaxed,
                cli_no_plugins=np,
                plugin_policy_override=pp,
            )
            self._entries[cid] = _ConnectionEntry(ctx=ctx)
        sid = ctx.router._chat_session_id or ""
        return cid, str(sid)

    async def destroy(self, connection_id: str) -> bool:
        async with self._global_lock:
            ent = self._entries.pop(connection_id, None)
        if ent is None:
            return False
        await ent.ctx.shutdown_graceful()
        return True

    async def invoke(
        self,
        connection_id: str,
        req: dict[str, Any],
        *,
        require_nonce: bool,
    ) -> dict[str, Any]:
        async with self._global_lock:
            ent = self._entries.get(connection_id)
        if ent is None:
            return {
                "ok": False,
                "id": req.get("id"),
                "error": "unknown connection_id",
                "code": "connection",
            }
        state_dir = ent.ctx.layout.get("plugins_state")
        async with ent.lock:
            return await dispatch_host_request(
                ent.ctx,
                req,
                state_dir=state_dir,
                require_nonce=require_nonce,
            )
