"""快照预热池 WarmPool（实验性，非生产路径，feat-chv-arch95 定案）。

定案（架构决策，见 docs/chv-architecture.md ADR-5）：WarmPool
不再宣称为生产快速启动特性。原因：预热 VM 的磁盘是干净快照
盘，不是任务盘；dispatch 前查池若要真正接线，必须在派任务
前给已暂停 VM 重新供给任务盘（热换盘/二盘挂载/host.py 磁盘
语义三者都未实现），把预热盘直接当任务盘用会击穿“任务盘
隔离”不变量。acquire 内 resume 只是实验性接口，worker
dispatch 永远走冷启动（wake_vm），_run_with_vm 绝不查池。
pool 保留的唯一生产价值是 reconcile_orphans()（上一世残留
清理），worker 启动/关闭仍调它。warm_pool_wired_to_dispatch
恒为 False，任何把它改成 True 的改动都必须先实现任务盘
重供给并补不变量测试，否则不得合入。


隔离红线（与 #216 一致，绝不妥协）：预热条目是一次性的
干净快照 VM，acquire 消耗掉它，release 一律销毁，绝不把
一个任务用过的 VM（带任务文件、带 agent 痕迹）翻出来给
下一个任务/租户复用；池子空了就回落冷启动（wake_vm），
绝不因为池空就降级 Docker/Local。

Registry（本轨新增）：
- 文件：state_dir/warm-pool.json（registry_path 可覆盖），
  内容为 {<合成名>: {"state", "vm", "task_id", "snapshot_url"}}。
  每次 prewarm/acquire/release 后立即落盘（临时文件 + 原子
  rename），worker 崩溃后重启仍能发现上一世的条目。
- 载入时绝不复用：新池实例从 registry 读到的条目一律标为
  orphan（不计入 warm_count，acquire 绝不发给任务），因为
  旧 VMM 进程已随旧 worker 消失，无法验证它还是当初那个
  干净快照 VM。它们只能走 reconcile_orphans() 清理。
- reconcile_orphans()：worker 启动/清理时的显式入口，对
  每个 orphan/assigned 条目尽力调 executor.shutdown_vm
  （清 socket 与残留登记），条目清掉并落盘，返回已清理的
  VM 名列表。registry 里没有的残留 socket 仍由执行器的
  cleanup_stale_vms() 扫 state_dir 处理，两者互补，不互相
  替代；reconcile 绝不在构造时自动执行（同一 state_dir 只
  能有一个活池，构造时自动清理会误杀另一个活池的 VM）。
- 损坏 fail-safe：registry 不是合法 JSON / 顶层不是 dict 时，
  整个文件 quarantine 成 <path>.corrupt（尽力而为），池子
  从空开始，绝不用损坏数据复用 VM；单条目结构非法时只
  丢该条（记 registry_corrupt_entries），不连坐其它条目。
  registry 写失败时只记日志警告、不炸任务流（池子内存态
  仍可用，重启后孤儿改由 socket 扫描兜底）。
"""

import json
import logging
import os

_LOGGER = logging.getLogger("cloud.cloud_hypervisor")

_REGISTRY_FILENAME = "warm-pool.json"


class WarmPool:
    """快照预热池（实验性，见模块头定案）。

    wired_to_dispatch = False：生产 dispatch 不查池。
    

    目标：冷启动（preflight + 起 VMM + create + boot + SSH
    就绪）是慢路径；预热池提前把干净 VM 起好并暂停，任务
    来时 acquire 直接 resume，省掉冷启动的大头。
    状态机：cold（未预热）→ prewarming（wake 中）→ warm
    （已暂停待分配）→ assigned（已分配给某任务）→
    running（已 resume）→ recycled/destroyed（release 后）。
    prewarm 的实现：对每个条目调 executor.wake_vm（冷启动）
    再 pause_vm，得到 paused 的干净 VM；mock 单测里 wake/
    pause 均可注入假实现，真机上这一步需要 KVM 宿主（见
    就绪清单“需 KVM 宿主实测”）。snapshot_url 未设时
    prewarm 返回 []（无基准快照可预热，不是静默造快照）。
    """

    wired_to_dispatch = False

    def __init__(self, executor, size=None, snapshot_url=None,
                 registry_path=None):
        self.executor = executor
        self.size = int(size if size is not None
                        else getattr(executor, "warm_pool_size", 0))
        self.snapshot_url = snapshot_url \
            if snapshot_url is not None \
            else getattr(executor, "snapshot_url", None)
        self._entries = {}  # 合成名 -> {"state", "vm", "task", ...}
        self.registry_path = registry_path
        if self.registry_path is None:
            state_dir = getattr(executor, "state_dir", None)
            if state_dir:
                self.registry_path = os.path.join(
                    state_dir, _REGISTRY_FILENAME)
        self.registry_corrupt = False
        self.registry_corrupt_entries = 0
        self._load_registry()

    # -- registry ------------------------------------------------------
    def _load_registry(self):
        """载入 registry：读到的条目一律标 orphan（见模块
        docstring，绝不直接当 warm 发给任务）。无文件时为空
        池；损坏时 quarantine 并从空开始。"""
        if not self.registry_path:
            return
        try:
            with open(self.registry_path, "r") as fh:
                data = json.load(fh)
        except FileNotFoundError:
            return
        except (OSError, ValueError):
            self.registry_corrupt = True
            self._quarantine_registry()
            return
        if not isinstance(data, dict):
            self.registry_corrupt = True
            self._quarantine_registry()
            return
        for name, entry in data.items():
            if not isinstance(entry, dict) \
                    or not entry.get("vm") \
                    or entry.get("state") not in ("warm", "assigned"):
                self.registry_corrupt_entries += 1
                continue
            self._entries[name] = {
                "state": "orphan",
                "vm": entry["vm"],
                "task": {"id": entry.get("task_id") or name},
                "snapshot_url": entry.get("snapshot_url"),
            }

    def _quarantine_registry(self):
        """把损坏的 registry 文件挪到 <path>.corrupt 旁（尽力
        而为），原路径不再留损坏数据给下一次载入。"""
        if not self.registry_path:
            return
        try:
            os.replace(self.registry_path,
                       self.registry_path + ".corrupt")
        except OSError:
            _LOGGER.warning("chv_pool_registry_quarantine_failed "
                            "path=%s", self.registry_path)

    def _save_registry(self):
        """落盘当前条目（只写 warm/assigned 态；orphan 是
        待清理的上一世条目，reconcile 后即删，不写回）。
        写失败只记日志，不炸任务流。"""
        if not self.registry_path:
            return
        data = {}
        for name, entry in self._entries.items():
            if entry["state"] not in ("warm", "assigned"):
                continue
            task = entry.get("task") or {}
            task_id = task.get("id") if isinstance(task, dict) \
                else None
            data[name] = {
                "state": entry["state"],
                "vm": entry.get("vm"),
                "task_id": task_id,
                "snapshot_url": entry.get("snapshot_url"),
            }
        try:
            directory = os.path.dirname(self.registry_path)
            if directory:
                os.makedirs(directory, exist_ok=True)
            tmp_path = self.registry_path + ".tmp"
            with open(tmp_path, "w") as fh:
                json.dump(data, fh)
            os.replace(tmp_path, self.registry_path)
        except OSError:
            _LOGGER.warning("chv_pool_registry_save_failed path=%s",
                            self.registry_path)

    def reconcile_orphans(self) -> list:
        """清理上一世残留条目：对每个 orphan/assigned 条目
        尽力 shutdown_vm（清 socket 与残留登记），条目清掉
        并落盘，返回已清理的 VM 名列表。worker 启动/清理时
        显式调它（见模块 docstring 的单活池假设）。"""
        cleaned = []
        for name, entry in list(self._entries.items()):
            if entry["state"] not in ("orphan", "assigned"):
                continue
            vm = entry.get("vm") or name
            try:
                self.executor.shutdown_vm(vm)
            except Exception:
                pass
            self._entries.pop(name, None)
            cleaned.append(vm)
        if cleaned:
            self._save_registry()
            try:
                self.executor._emit(
                    "warm_pool_reconcile", {"id": None},
                    "cleaned=%s" % ",".join(cleaned))
            except Exception:
                pass
        return cleaned

    # -- 池操作 ----------------------------------------------------------
    def warm_count(self) -> int:
        return sum(1 for e in self._entries.values()
                   if e["state"] == "warm")

    def orphan_count(self) -> int:
        """上一世残留条目数（待 reconcile，不可分配）。"""
        return sum(1 for e in self._entries.values()
                   if e["state"] in ("orphan", "assigned"))

    def prewarm(self, count=None):
        """预热到目标数：已有 warm 条目不重复建；每个新条目
        用一个合成任务名 wake + pause（合成任务 id 为
        warm-<序号>，绝不与真实任务 id 撞名）。返回新增的
        warm VM 名列表。"""
        if self.snapshot_url is None:
            return []
        target = int(count) if count is not None else self.size
        added = []
        i = 0
        while self.warm_count() < target:
            i += 1
            name = "warm-%d" % i
            if name in self._entries:
                continue
            task = {"id": name}
            self._entries[name] = {"state": "prewarming",
                                   "task": task}
            try:
                self.executor.wake_vm(task)
                self.executor.pause_vm(task)
            except Exception:
                self._entries.pop(name, None)
                raise
            vm = self.executor.vm_name(task)
            self._entries[name] = {
                "state": "warm", "task": task, "vm": vm,
                "snapshot_url": self.snapshot_url}
            added.append(name)
        if added:
            self._save_registry()
        return added

    def acquire(self, task):
        """分配一个预热 VM 给任务（feat-chv-fix，P1-3 接线）：

        有 warm 条目时先调 executor.resume_vm() 把暂停的 VM
        真正唤醒（resume 失败即条目作废、继续找下一个，
        绝不把 paused 的 VM 当 running 发给任务），再标
        assigned 并返回 mode=warm；resume 全部失败或无 warm
        条目时返回 mode=cold（调用方回落 wake_vm 冷启动）。
        注意：本方法只解决“acquire 不再悬空 resume 责任”
        这一半；worker dispatch 前查池的另一半仍未接线
        （见执行器/就绪文档的“未接线缺口”标注），预热收益
        在生产路径上仍不是已交付特性。
        feat-chv-secfix（L5，接线时的盘对应关系定义——接线
        本身仍未做，不假装已接）：预热 VM 的磁盘是干净快照
        盘，不是任务盘。将来 dispatch 前查池接线时必须做到：
        acquire 返回的 VM 在派任务前仍须按任务重新供给磁盘
        （provision_task_disk 的 copy 盘，或校验快照来源与该
        任务的对应关系后才能用快照盘），绝不能把预热盘直接
        当任务盘用；接线时须补“任务盘来源”不变量的测试，
        否则 acquire 的 warm 路径不得合入生产路径。"""
        for name, entry in list(self._entries.items()):
            if entry["state"] != "warm":
                continue
            resume = getattr(self.executor, "resume_vm", None)
            if resume is not None:
                try:
                    resume(entry.get("task") or task)
                except Exception:
                    self._entries.pop(name, None)
                    self._save_registry()
                    continue
            # 无 resume_vm 的鸭子型执行器（旧测试替身）：
            # 保持旧 assigned 语义，不因此丢弃条目；真实
            # CloudHypervisorExecutor 一律有 resume_vm，走
            # 上面分支，责任不悬空。
            entry["state"] = "assigned"
            entry["task"] = task
            self._save_registry()
            try:
                self.executor._emit(
                    "warm_pool_acquire", task,
                    "vm=%s" % name)
            except Exception:
                pass
            return {"mode": "warm", "vm": name,
                    "pool_task": entry["task"]}
        return {"mode": "cold", "vm": None, "pool_task": None}

    def release(self, task, pool_vm=None):
        """释放：无论 warm 还是 cold 路径，VM 一律销毁
        （executor.shutdown_vm），池条目清掉；绝不把用过
        的 VM 放回 warm 池。返回 True 表示确实销毁了一个
        条目或 VM。"""
        destroyed = False
        if pool_vm and pool_vm in self._entries:
            self._entries.pop(pool_vm, None)
            destroyed = True
        try:
            if self.executor.shutdown_vm(task):
                destroyed = True
        except Exception:
            pass
        self._save_registry()
        self.executor._emit(
            "warm_pool_release", task,
            "vm=%s destroyed=%s" % (pool_vm, destroyed))
        return destroyed


def reconcile_orphans_for_executor(executor) -> list:
    """生产路径唯一的 WarmPool 入口：只做孤儿清理。

    worker 启动/关闭调它，不构造 WarmPool 类、不调
    prewarm/acquire/release——实验接口与生产路径彻底分离：
    生产代码 import 本函数即可，acquire/prewarm 只存在于
    实验用的 WarmPool 类与测试里。"""
    return WarmPool(executor).reconcile_orphans()