"""One gateway owns reconciliation and durable conversation affinity.

Never move conversations from a live worker. A draining worker is stopped only
after its run registry AND outstanding HTTP requests are empty. Failure does
not replay a request, because a model/tool may already have produced effects.
"""

from __future__ import annotations

import asyncio
import json
import os
import sqlite3
import time
from urllib.parse import quote

import httpx

from .docker import DockerEngine, DockerError, worker_spec
from .settings import SettingsStore


class PoolUnavailable(RuntimeError):
    pass


class Controller:
    def __init__(self, store: SettingsStore):
        self.store = store
        self.lock = asyncio.Lock()
        self.route_lock = asyncio.Lock()
        self.nodes: dict[str, dict] = {}
        self.error: str | None = None
        self.task: asyncio.Task | None = None
        self._lock_handle = None
        self.http = httpx.AsyncClient(timeout=httpx.Timeout(15, read=None), trust_env=False, limits=httpx.Limits(max_connections=256, max_keepalive_connections=64))
        self.db: sqlite3.Connection | None = None
        self.pending = 0
        self.inflight: dict[str, int] = {}
        self.local_threads: dict[str, int] = {}
        self._python_probe_key: tuple | None = None
        self._python_probe_result: str | None = None
        self._worker_image_key: tuple | None = None
        self._worker_image_result: str | None = None
        self.discovery_complete = not bool(store.load().source_container)

    def _ensure_db(self):
        if self.db is None:
            self.db = sqlite3.connect(self.store.directory / "affinity.sqlite3")
            self.db.execute("PRAGMA journal_mode=WAL")
            self.db.execute("CREATE TABLE IF NOT EXISTS affinity (thread TEXT PRIMARY KEY, node TEXT NOT NULL)")
            self.db.commit()
        return self.db

    def _acquire_controller_lock(self):
        if self._lock_handle is not None:
            return
        self.store.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
        handle = (self.store.directory / "controller.lock").open("a+b")
        try:
            if os.name == "nt":
                import msvcrt

                handle.seek(0)
                if not handle.read(1):
                    handle.write(b"0")
                    handle.flush()
                handle.seek(0)
                msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
            else:
                import fcntl

                fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except OSError:
            handle.close()
            raise RuntimeError("Worker 池只能由一个主服务进程管理；请关闭重复的 Worker 池控制器") from None
        self._lock_handle = handle

    def start(self):
        settings = self.store.load()
        if settings.enabled:
            self._acquire_controller_lock()
            self._ensure_db()
            if self.task is None or self.task.done():
                self.task = asyncio.create_task(self.loop())
        if hasattr(self, "diagnostics") and self.diagnostics.enabled:
            self.diagnostics.start()

    async def close(self):
        if self.task:
            self.task.cancel()
            await asyncio.gather(self.task, return_exceptions=True)
        if hasattr(self, "diagnostics"):
            await self.diagnostics.close()
        await self.http.aclose()
        if self.db is not None:
            self.db.close()
            self.db = None
        self._release_controller_lock()

    async def loop(self):
        try:
            while True:
                try:
                    await self.reconcile()
                except Exception as exc:
                    self.error = str(exc) if isinstance(exc, (DockerError, PoolUnavailable)) else f"节点检查失败（{type(exc).__name__}）"
                settings = self.store.load()
                if not settings.enabled and not self.nodes:
                    return
                await asyncio.sleep(5)
        finally:
            if not self.store.load().enabled:
                self._release_controller_lock()

    def _release_controller_lock(self):
        if self._lock_handle is None:
            return
        if os.name == "nt":
            import msvcrt

            self._lock_handle.seek(0)
            msvcrt.locking(self._lock_handle.fileno(), msvcrt.LK_UNLCK, 1)
        else:
            import fcntl

            fcntl.flock(self._lock_handle, fcntl.LOCK_UN)
        self._lock_handle.close()
        self._lock_handle = None

    async def inspect_template(self, settings):
        config = settings
        if config.desired * config.cpu > config.total_cpu or config.desired * config.memory_mb > config.total_memory_mb:
            raise PoolUnavailable("Worker 总资源超过设置的预算，请减少数量或调整预算")
        if config.desired * (config.database_pool_size + 2 * config.postgres_pool_size) > config.database_connection_budget:
            raise PoolUnavailable("Worker 数据库连接总额度超过预算，请减少数量或单节点连接数")
        if not config.source_container or not config.network:
            raise PoolUnavailable("请填写源容器名称和 Docker 网络")
        engine = DockerEngine(config)
        source = await asyncio.to_thread(engine.request, "GET", f"/containers/{quote(config.source_container, safe='')}/json?size=true")
        if source.get("Config", {}).get("Labels", {}).get("deerflow.worker-pool"):
            raise PoolUnavailable("源容器必须是主服务，不能选择已有 Worker")
        backend = await asyncio.to_thread(
            engine.configured_storage_backend, source, config.python_executable, config.working_dir,
        )
        probe_key = (
            config.transport, config.host, config.port, config.username, config.source_container,
            source.get("Id"), source.get("Image"), config.python_executable,
            config.working_dir, source.get("Config", {}).get("User", ""), backend,
        )
        if probe_key == self._python_probe_key and self._python_probe_result:
            python_executable = self._python_probe_result
        else:
            python_executable = await asyncio.to_thread(
                engine.discover_worker_python, source, config.python_executable, config.working_dir,
            )
            self._python_probe_key = probe_key
            self._python_probe_result = python_executable
        if python_executable != config.python_executable:
            config = config.model_copy(update={"python_executable": python_executable})
            self.store.save(config)
            self._python_probe_key = (
                config.transport, config.host, config.port, config.username, config.source_container,
                source.get("Id"), source.get("Image"), config.python_executable,
                config.working_dir, source.get("Config", {}).get("User", ""), backend,
            )
        runtime_image = None
        if config.enabled and config.desired > 0:
            image_key = (source.get("Id"), source.get("Image"), python_executable, backend)
            if image_key != self._worker_image_key or not self._worker_image_result:
                self._worker_image_result = await asyncio.to_thread(
                    engine.snapshot_worker_runtime, source, self.store.raw()["instance"],
                )
                self._worker_image_key = image_key
            runtime_image = self._worker_image_result
        spec = worker_spec(
            source, config, self.store.raw()["instance"], self.store.raw()["token"],
            str(self.store.directory), runtime_image=runtime_image,
            require_shared_archive=backend == "postgres" and config.enabled and config.desired > 0,
        )
        return engine, spec

    async def reconcile(self):
        async with self.lock:
            settings = self.store.load()
            if not settings.enabled and not settings.source_container:
                return
            db = self._ensure_db()
            engine, spec = await self.inspect_template(settings)
            instance = self.store.raw()["instance"]
            filters = quote(json.dumps({"label": [f"deerflow.worker-pool={instance}"]}))
            containers = await asyncio.to_thread(engine.request, "GET", f"/containers/json?all=1&filters={filters}")
            self.discovery_complete = True
            desired = settings.desired if settings.enabled else 0
            wanted = {f"dfw-{instance}-{i}" for i in range(desired)}
            existing = {row["Names"][0].lstrip("/"): row for row in containers}
            stale_templates = set()
            for name, row in list(existing.items()):
                self.nodes.setdefault(name, {"name": name, "id": row["Id"], "ready": False, "active": 0, "draining": name not in wanted, "url": f"http://{name}:8001"})
                if name in wanted and row.get("State") == "running" and row.get("ImageID") and row["ImageID"] != spec["Image"]:
                    stale_templates.add(name)
                    self.nodes[name]["draining"] = True

            # A stopped container cannot be executing a request. If its saved
            # startup template has an old Python path, replace it with the
            # resolved template instead of restarting the same broken command.
            for name, row in list(existing.items()):
                if name not in wanted or row["State"] == "running":
                    continue
                detail = await asyncio.to_thread(engine.request, "GET", f"/containers/{quote(row['Id'], safe='')}/json")
                worker_config = detail.get("Config", {})
                stale_template = (
                    detail.get("Image") != spec["Image"]
                    or worker_config.get("Entrypoint") != spec["Entrypoint"]
                    or worker_config.get("WorkingDir") != spec["WorkingDir"]
                )
                if not stale_template:
                    continue
                if hasattr(self, "diagnostics"):
                    await self.diagnostics.archive_logs(engine, name)
                await asyncio.to_thread(engine.request, "DELETE", f"/containers/{quote(row['Id'], safe='')}")
                db.execute("DELETE FROM affinity WHERE node=?", (name,))
                db.commit()
                existing.pop(name, None)
                self.nodes.pop(name, None)
            for name, row in list(existing.items()):
                previous = self.nodes.get(name, {})
                node = {
                    "name": name,
                    "id": row["Id"],
                    "ready": previous.get("ready", False),
                    "ever_ready": previous.get("ever_ready", False),
                    "active": previous.get("active", 0),
                    "capacity": previous.get("capacity"),
                    "draining": name not in wanted or name in stale_templates,
                    "url": f"http://{name}:8001",
                    "checked_at": time.time(),
                }
                for key in ("running", "waiting", "process_started_at", "runtime_revision", "event_loop_lag_ms", "http"):
                    if key in previous:
                        node[key] = previous[key]
                self.nodes[name] = node
                if row["State"] != "running":
                    node["ready"] = False
                    node["ever_ready"] = False
                    # A confirmed stopped process cannot keep executing. Old
                    # conversations can move only after this confirmation.
                    db.execute("DELETE FROM affinity WHERE node=?", (name,))
                    db.commit()
                    if name in wanted:
                        await asyncio.to_thread(engine.request, "POST", f"/containers/{row['Id']}/start")
                    else:
                        if hasattr(self, "diagnostics"):
                            await self.diagnostics.archive_logs(engine, name)
                        await asyncio.to_thread(engine.request, "DELETE", f"/containers/{row['Id']}")
                        self.nodes.pop(name, None)
                    continue
                try:
                    diagnostics_enabled = bool(getattr(getattr(self, "diagnostics", None), "enabled", False))
                    response = await self.http.get(
                        node["url"] + "/internal/worker-status",
                        headers={
                            "x-deerflow-worker-token": self.store.raw()["token"],
                            "x-deerflow-diagnostics-enabled": "true" if diagnostics_enabled else "false",
                        },
                        timeout=3,
                    )
                    response.raise_for_status()
                    status = response.json()
                    node.update(ready=status["ready"], active=status["active"], capacity=status["capacity"], ever_ready=True)
                    for key in ("running", "waiting", "process_started_at", "runtime_revision", "event_loop_lag_ms", "http"):
                        if key in status:
                            node[key] = status[key]
                except (httpx.HTTPError, ValueError, KeyError):
                    node["ready"] = False
                    node["error"] = "等待就绪或连接失败；请检查主容器是否位于同一网络、镜像代码和聊天存储配置"
                    continue
                if node["draining"] and not node["active"] and not self.inflight.get(name, 0):
                    # Refuse all new reservations while stop is in flight.
                    node["ready"] = False
                    node["ever_ready"] = False
                    await asyncio.to_thread(engine.request, "POST", f"/containers/{row['Id']}/stop?t=30")
                    if hasattr(self, "diagnostics"):
                        await self.diagnostics.archive_logs(engine, name)
                    await asyncio.to_thread(engine.request, "DELETE", f"/containers/{row['Id']}")
                    db.execute("DELETE FROM affinity WHERE node=?", (name,))
                    db.commit()
                    self.nodes.pop(name, None)
                    existing.pop(name, None)
            # Replacements are created only after the previous container has
            # drained and been removed. This also refreshes Workers after the
            # main service's runtime image or configured storage backend changes.
            for name in sorted(wanted - existing.keys()):
                row = await asyncio.to_thread(engine.request, "POST", f"/containers/create?name={name}", spec)
                self.nodes[name] = {
                    "name": name,
                    "id": row["Id"],
                    "ready": False,
                    "active": 0,
                    "draining": False,
                    "url": f"http://{name}:8001",
                }
                await asyncio.to_thread(engine.request, "POST", f"/containers/{row['Id']}/start")
                existing[name] = {"Id": row["Id"], "State": "running", "ImageID": spec["Image"]}
            # Missing nodes are not silently reassigned: an operator must
            # reconcile the actual container state before changing affinity.
            self.error = None

    async def reserve(self, thread: str, *, create: bool) -> dict | None:
        settings = self.store.load()
        if not settings.enabled and not self.nodes:
            return None
        db = self._ensure_db()
        # Container creation/SSH can be slow; routing must never wait on that lock.
        async with self.route_lock:
            row = db.execute("SELECT node FROM affinity WHERE thread=?", (thread,)).fetchone()
            if row is None and create and self.local_threads.get(thread, 0):
                raise PoolUnavailable("该会话正在主服务执行，请等待当前轮完成")
            if row:
                node = self.nodes.get(row[0])
                if node and node.get("draining") and create:
                    return {"_fallback_to_main": True, "_fallback_reason": "worker_draining", "_worker": dict(node)}
                # A delayed health probe stops new work, but must not hide an
                # existing run's result. Keep reads/cancellation on its owner;
                # stopped containers clear ever_ready before affinity changes.
                if not node or (not node["ready"] and (create or not node.get("ever_ready"))):
                    raise PoolUnavailable("该会话的执行节点暂不可用，请稍后重试；系统不会重复提交任务")
                if node["draining"] and create:
                    raise PoolUnavailable("该会话节点正在停止接单，请等待当前任务完成后重试")
                if create:
                    capacity = max(1, int(node.get("capacity") or settings.concurrency))
                    load = max(int(node.get("active") or 0), self.inflight.get(node["name"], 0))
                    if load >= capacity:
                        return {"_fallback_to_main": True, "_fallback_reason": "worker_saturated", "_worker": dict(node)}
            elif not settings.enabled or not create:
                return None
            else:
                choices = [n for n in self.nodes.values() if n["ready"] and not n["draining"]]
                if not choices:
                    return {"_fallback_to_main": True, "_fallback_reason": "workers_unavailable"}
                available = []
                for candidate in choices:
                    capacity = max(1, int(candidate.get("capacity") or settings.concurrency))
                    load = max(int(candidate.get("active") or 0), self.inflight.get(candidate["name"], 0))
                    if load < capacity:
                        available.append((load, candidate))
                if not available:
                    return {"_fallback_to_main": True, "_fallback_reason": "worker_pool_saturated"}
                _, node = min(available, key=lambda item: (item[0], item[1]["name"]))
                db.execute("INSERT INTO affinity VALUES (?, ?)", (thread, node["name"]))
                db.commit()
            self.inflight[node["name"]] = self.inflight.get(node["name"], 0) + 1
            return dict(node)

    def release(self, name):
        self.inflight[name] = max(0, self.inflight.get(name, 0) - 1)

    def status(self):
        return {"settings": self.store.public(), "nodes": [{k: v for k, v in n.items() if k != "url"} for n in self.nodes.values()], "error": self.error}