"""宿主级 Sandbox 重型执行容量。

状态保存在 workspace/.sandbox 下并由 advisory file lock 串行化，因此蓝绿和多
Web 进程共享同一组槽位。业务 DB 不承载实时状态；后台排队事实源仍是
``.zcbot_procs``，这里只记录已经获得槽位的租约和短暂的前台候选。
"""
from __future__ import annotations

import json
import os
import time
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import Callable, Dict, Iterator, Optional

from core.file_store import atomic_write_text, interprocess_file_lock

DEFAULT_MAX_ACTIVE_EXECS = 6
DEFAULT_MAX_BACKGROUND_EXECS = 4
DEFAULT_MAX_ACTIVE_EXECS_PER_USER = 2
DEFAULT_MIN_MEM_AVAILABLE_BYTES = 1024 ** 3


def _parse_bytes(value: object, default: int) -> int:
    text = str(value or "").strip().lower()
    if not text:
        return default
    units = {"k": 1024, "kb": 1024, "m": 1024**2, "mb": 1024**2,
             "g": 1024**3, "gb": 1024**3}
    for suffix, factor in sorted(units.items(), key=lambda x: -len(x[0])):
        if text.endswith(suffix):
            return int(float(text[:-len(suffix)]) * factor)
    return int(text)


def mem_available_bytes() -> Optional[int]:
    try:
        for line in Path("/proc/meminfo").read_text(encoding="ascii").splitlines():
            if line.startswith("MemAvailable:"):
                return int(line.split()[1]) * 1024
    except (OSError, ValueError, IndexError):
        return None
    return None


class ExecCapacity:
    def __init__(self, state_dir: Path, cfg: Optional[dict] = None) -> None:
        cfg = cfg or {}
        self.state_dir = Path(state_dir)
        self.state_path = self.state_dir / "exec-capacity.json"
        self.lock_path = self.state_dir / "exec-capacity.lock"
        self.max_active = max(1, min(DEFAULT_MAX_ACTIVE_EXECS, int(os.getenv("ZCBOT_MAX_ACTIVE_EXECS") or cfg.get("max_active_execs") or DEFAULT_MAX_ACTIVE_EXECS)))
        self.max_background = max(1, min(DEFAULT_MAX_BACKGROUND_EXECS, int(os.getenv("ZCBOT_MAX_BACKGROUND_EXECS") or cfg.get("max_background_execs") or DEFAULT_MAX_BACKGROUND_EXECS)))
        self.max_per_user = max(1, min(DEFAULT_MAX_ACTIVE_EXECS_PER_USER, int(os.getenv("ZCBOT_MAX_ACTIVE_EXECS_PER_USER") or cfg.get("max_active_execs_per_user") or DEFAULT_MAX_ACTIVE_EXECS_PER_USER)))
        self.min_mem_available = _parse_bytes(
            os.getenv("ZCBOT_MIN_MEM_AVAILABLE") or cfg.get("min_mem_available"),
            DEFAULT_MIN_MEM_AVAILABLE_BYTES,
        )

    def _read(self) -> dict:
        try:
            data = json.loads(self.state_path.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                data.setdefault("leases", {})
                data.setdefault("foreground_queue", [])
                data.setdefault("containers", {})
                return data
        except (OSError, ValueError):
            pass
        return {"leases": {}, "foreground_queue": [], "containers": {}}

    def _write(self, state: dict) -> None:
        atomic_write_text(self.state_path, json.dumps(state, ensure_ascii=False, sort_keys=True))

    def _prune(self, state: dict) -> None:
        leases = state["leases"]
        for key, lease in list(leases.items()):
            if lease.get("kind") == "foreground":
                pid = int(lease.get("owner_pid") or 0)
                if pid and not self._pid_alive(pid):
                    leases.pop(key, None)
        state["foreground_queue"] = [
            q for q in state["foreground_queue"]
            if self._pid_alive(int(q.get("owner_pid") or 0))
        ]

    @staticmethod
    def _pid_alive(pid: int) -> bool:
        try:
            os.kill(pid, 0)
            return True
        except PermissionError:
            return True
        except OSError:
            return False

    def _can_admit(self, state: dict, user_id: str, kind: str) -> bool:
        leases = list(state["leases"].values())
        if len(leases) >= self.max_active:
            return False
        if sum(1 for x in leases if x.get("user_id") == user_id) >= self.max_per_user:
            return False
        if kind == "background" and sum(1 for x in leases if x.get("kind") == "background") >= self.max_background:
            return False
        available = mem_available_bytes()
        return available is None or available >= self.min_mem_available

    def try_acquire(self, user_id: str, kind: str, *, lease_id: Optional[str] = None, proc_id: Optional[str] = None) -> Optional[str]:
        lease_id = lease_id or uuid.uuid4().hex
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read()
            self._prune(state)
            if lease_id in state["leases"]:
                return None
            if not self._can_admit(state, str(user_id), kind):
                self._write(state)
                return None
            state["leases"][lease_id] = {
                "lease_id": lease_id, "user_id": str(user_id), "kind": kind,
                "proc_id": proc_id, "owner_pid": os.getpid(), "started_ts": time.time(),
            }
            self._write(state)
            return lease_id

    def acquire_foreground(self, user_id: str, cancel_check: Optional[Callable[[], bool]] = None) -> Optional[str]:
        ticket = uuid.uuid4().hex
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read(); self._prune(state)
            state["foreground_queue"].append({"ticket": ticket, "user_id": str(user_id), "created_ts": time.time(), "owner_pid": os.getpid()})
            self._write(state)
        while True:
            with interprocess_file_lock(self.lock_path, timeout_seconds=None):
                state = self._read(); self._prune(state)
                queue = state["foreground_queue"]
                eligible = next((q for q in queue if q.get("user_id") == str(user_id)), None)
                if eligible and eligible.get("ticket") == ticket and self._can_admit(state, str(user_id), "foreground"):
                    queue[:] = [q for q in queue if q.get("ticket") != ticket]
                    state["leases"][ticket] = {"lease_id": ticket, "user_id": str(user_id), "kind": "foreground", "owner_pid": os.getpid(), "started_ts": time.time()}
                    self._write(state)
                    return ticket
                self._write(state)
            if cancel_check is not None and cancel_check():
                self.cancel_waiter(ticket)
                return None
            time.sleep(0.1)

    def cancel_waiter(self, ticket: str) -> None:
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read()
            state["foreground_queue"] = [q for q in state["foreground_queue"] if q.get("ticket") != ticket]
            self._write(state)

    def release(self, lease_id: Optional[str]) -> None:
        if not lease_id:
            return
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read(); state["leases"].pop(lease_id, None); self._write(state)

    def reconcile_background(self, running_proc_ids: set[str]) -> None:
        """服务重启后保留 Docker 仍在跑的租约，清启动窗口遗留的孤儿租约。"""
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read()
            for key, lease in list(state["leases"].items()):
                if (lease.get("kind") == "background"
                        and str(lease.get("proc_id") or "") not in running_proc_ids
                        and not self._pid_alive(int(lease.get("owner_pid") or 0))):
                    state["leases"].pop(key, None)
            self._write(state)

    def touch_container(self, name: str, *, active_delta: int = 0) -> None:
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read()
            row = state["containers"].setdefault(name, {"active_execs": 0})
            row["active_execs"] = max(0, int(row.get("active_execs") or 0) + active_delta)
            row["last_active_ts"] = time.time()
            self._write(state)

    def remove_container(self, name: str) -> None:
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read(); state["containers"].pop(name, None); self._write(state)

    def reconcile_containers(self, running_names: set[str]) -> None:
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read()
            for name in list(state["containers"]):
                if name not in running_names:
                    state["containers"].pop(name, None)
            self._write(state)

    @contextmanager
    def foreground(self, user_id: str, cancel_check: Optional[Callable[[], bool]] = None) -> Iterator[bool]:
        lease = self.acquire_foreground(user_id, cancel_check)
        try:
            yield lease is not None
        finally:
            self.release(lease)

    def snapshot(self) -> dict:
        with interprocess_file_lock(self.lock_path, timeout_seconds=None):
            state = self._read(); self._prune(state); self._write(state)
        leases = list(state["leases"].values())
        fg = sum(x.get("kind") == "foreground" for x in leases)
        bg = sum(x.get("kind") == "background" for x in leases)
        by_user: Dict[str, int] = {}
        for x in leases:
            by_user[x["user_id"]] = by_user.get(x["user_id"], 0) + 1
        available_mem = mem_available_bytes()
        memory_paused = available_mem is not None and available_mem < self.min_mem_available
        containers = state.get("containers", {})
        now = time.time()
        return {
            "limits": {"active": self.max_active, "background": self.max_background, "per_user": self.max_per_user},
            "foreground_running": fg, "foreground_queued": len(state["foreground_queue"]),
            "background_running": bg, "per_user": by_user,
            "admit_available": 0 if memory_paused else max(0, self.max_active - len(leases)),
            "memory_paused": memory_paused,
            "mem_available_bytes": available_mem,
            "cpu_load": list(os.getloadavg()) if hasattr(os, "getloadavg") else None,
            "active_sandbox_containers": len(containers),
            "idle_reap_candidates": sum(int(x.get("active_execs") or 0) == 0 and float(x.get("last_active_ts") or now) < now - getattr(self, "idle_ttl_seconds", 600) for x in containers.values()),
        }