"""高层 MaskablePPO 训练（P1）。

架构与边界：
- 环境 = :class:`mcm26b.hrl.hierarchical_env.HierarchicalEnv`，每一步是一个高层
  宏动作（选择 0–119 槽位：clear/locate/station/fallback/保留）；
- 训练/部署共用 :mod:`mcm26b.hrl.scheduler` 的候选、观测和动作掩码；
- 低层默认使用几何选点；``--local-selector rl`` 仅作为接口预留并记录，
  当前 HTTP 执行路径仍由高层宏动作终点决定，不把八方向推演当正式收益。

落盘保证（长训必需）：
- **急停可恢复**：``--recovery-every``（默认 2000 环境步）把一份滚动恢复点
  ``<模型目录>/checkpoints/<stem>_latest.zip`` 原子覆盖写；训练开始前先写一份
  起步基线。因此即使进程被 SIGKILL / ``taskkill /F`` 强杀（没有任何 Python
  收尾机会），也**最多只丢 ``--recovery-every`` 个环境步**，且磁盘上始终有一份
  完整、可加载、可 ``--resume`` 的产物；
- **原子落盘**：模型 zip、旁挂 json 与 ``<run_dir>/resume.json`` 都是"临时文件 +
  ``os.replace``"，保存途中被杀不会留下半截文件，也不会破坏上一份完整快照；
- ``--checkpoint-every`` 每 N 个环境步在 ``<模型目录>/checkpoints/`` 另留一份历史
  checkpoint（``<stem>_<步数>.zip``），只保留最近 ``--checkpoint-keep`` 个，供
  返祖对比与多档恢复；
- ``<run_dir>/resume.json`` 是恢复的**权威指针**：无论上一刻是滚动恢复点、历史
  checkpoint、优雅收尾还是异常兜底，它总指向一份可续训产物，并记录增量语义
  （``timesteps`` 是追加步数）与生效过的 PPO 超参；
- ``--control-file`` 是训练前端的**跨平台**控制通道：前端原子写 JSON
  （``paused`` / ``stop``），本脚本每个 step 轮询。暂停时训练在 step 边界挂起，
  停止时先保存模型与元数据再退出；
- 信号（SIGTERM/SIGINT，Windows 上另有 SIGBREAK）仍可用：直接 kill 本脚本或
  前端发进程组信号时同样走优雅收尾；subproc worker 在 :func:`_make_env` 里忽略
  这些信号，保证它们活到主进程 ``vec_env.close()`` 之后再退出；
- 任何异常都会先落一份 ``*_emergency.zip`` + ``resume.json`` 再抛出，不再出现
  "跑了很久什么都没有"。

为什么暂停/停止不靠信号：Windows 没有 SIGSTOP，且 ``os.kill(pid, SIGTERM)`` 会
退化成 ``TerminateProcess`` 硬杀——收尾分支根本不会执行。控制文件两端都只是
普通文件 IO，POSIX 与 Windows 行为一致（见 :mod:`mcm26b.process_control`）。

P1 smoke 示例：
    python scripts/train_hrl_high.py \
        --problem 3 --stage 1 --timesteps 2048 --seed 320001 \
        --envs 1 --vec-env dummy --device mps \
        --config-profile production --model output/hrl/p1-smoke/q3.zip \
        --eval-episodes 2 --verbose 0
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import platform
import random
import signal
import sys
import time
from pathlib import Path
from typing import Any

from mcm26b.config import ProblemConfig
from mcm26b.hrl.curriculum import STAGES
from mcm26b.hrl.constants import (ACTION_LAYOUT_VERSION, FEATURE_VERSION,
                                  HIGH_OBS_DIM, N_CANDIDATES, REWARD_VERSION)
from mcm26b.process_control import ControlState, read_control


def _stage_from_index(index: int):
    if not 1 <= int(index) <= len(STAGES):
        raise ValueError(f"stage 必须在 1..{len(STAGES)}")
    return STAGES[int(index) - 1]


def _profile_fidelity(profile: str) -> str:
    if profile == "production":
        return "production"
    if profile == "coarse-pretrain":
        return "fast"
    raise ValueError(f"未知 config-profile: {profile}")


def _reward_from_profile(profile: str):
    from mcm26b.hrl.hierarchical_env import HRLReward

    if profile == "time":
        return HRLReward.time_only()
    if profile == "prompt":
        return HRLReward()
    raise ValueError(f"未知 reward-profile: {profile}")


def _stop_signals() -> list[int]:
    """本平台可用于"请求收尾"的信号集合。

    Windows 上没有 SIGTERM 的软语义（``os.kill`` 直接 TerminateProcess），但
    ``SIGBREAK`` 可以由 ``CTRL_BREAK_EVENT`` 触发并进入 Python 处理器；POSIX 上
    SIGTERM/SIGINT 均可捕获。控制文件才是主通道，这里只是让"直接 kill 脚本"
    也尽量走优雅收尾。
    """
    signals = [signal.SIGTERM, signal.SIGINT]
    sigbreak = getattr(signal, "SIGBREAK", None)
    if sigbreak is not None:
        signals.append(sigbreak)
    return signals


def _ignore_stop_signals_in_worker() -> None:
    """subproc worker 忽略收尾信号，保证主进程能优雅落盘。

    无论信号来自进程组 SIGTERM 还是 Windows ``CTRL_BREAK_EVENT``，若 worker
    同时退出，主进程会在保存模型之前就拿到断管/EOF 错误，整局训练白跑。worker
    活在管道后面，主进程收尾时通过 ``vec_env.close()`` 的 ``close`` 指令正常
    结束它们。
    """
    try:
        import multiprocessing
        if multiprocessing.current_process().name == "MainProcess":
            return
        for signum in _stop_signals():
            signal.signal(signum, signal.SIG_IGN)
    except (ImportError, ValueError, OSError):  # 非主线程/平台不支持：不影响训练
        return


class _ControlPoll:
    """控制文件轮询状态机（纯标准库，可脱离 SB3 单测）。

    ``poll()`` 每 ``poll_every`` 次调用才真正读盘，避免每个 env step 都做一次
    文件 IO；``wait_while_paused()`` 阻塞到暂停解除或收到 stop，并累计暂停秒数
    （``--max-wall-s`` 要把这段扣掉，否则"挂起一小时回来就判超时"）。
    """

    def __init__(self, path: Path, *, poll_every: int = 32,
                 poll_s: float = 0.2):
        self.path = Path(path)
        self.poll_every = max(1, int(poll_every))
        self.poll_s = max(0.02, float(poll_s))
        self.paused_total_s = 0.0
        self._calls = 0
        self._state = ControlState()

    def poll(self) -> ControlState:
        """按抽样间隔返回最近一次控制状态。"""
        self._calls += 1
        if self._calls % self.poll_every:
            return self._state
        self._state = read_control(self.path)
        return self._state

    def wait_while_paused(self, on_pause, on_resume) -> bool:
        """阻塞到继续或停止；返回 True 表示期间收到了 stop。

        ``on_pause`` / ``on_resume`` 用来向前端回显事件（训练脚本的 JSON 行）。
        """
        on_pause()
        started = time.monotonic()
        stopped = False
        while True:
            time.sleep(self.poll_s)
            self._state = read_control(self.path)
            if self._state.stop:
                stopped = True
                break
            if not self._state.paused:
                break
        self.paused_total_s += time.monotonic() - started
        on_resume()
        return stopped


def _make_env(problem: int, stage_index: int,
              monitor_path: str | None, profile: str,
              reward_profile: str):
    """顶层工厂：供 DummyVecEnv / SubprocVecEnv 调用（必须是可 pickle 的）。"""
    from stable_baselines3.common.monitor import Monitor

    from mcm26b.hrl.hierarchical_env import HierarchicalEnv

    _ignore_stop_signals_in_worker()
    stage = _stage_from_index(stage_index)
    env = HierarchicalEnv(
        problem=problem,
        stage=stage,
        fidelity=_profile_fidelity(profile),
        max_episode_steps=300,
        real_budget_s=1800.0,
        reward=_reward_from_profile(reward_profile),
    )
    return Monitor(env, filename=monitor_path)


def _build_vec_env(args, monitor_paths: list[str]):
    from stable_baselines3.common.vec_env import (DummyVecEnv, SubprocVecEnv)

    from functools import partial

    factories = [
        partial(_make_env, args.problem, args.stage,
                monitor_paths[rank], args.config_profile, args.reward_profile)
        for rank in range(args.envs)
    ]
    if args.vec_env == "subproc":
        # 每个 worker 独立 Monitor 文件；spawn 入口在 __main__ 保护内。
        return SubprocVecEnv(factories, start_method="spawn")
    return DummyVecEnv(factories)


def _hash_file(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _policy_hyperparams(model) -> dict[str, Any]:
    """记录**实际生效**的 PPO 超参。

    ``--resume``/``--init`` 时 ``MaskablePPO.load`` 会一并恢复 checkpoint 内的
    n_steps/batch_size/n_epochs 等，命令行同名参数被忽略；把它们写进元数据，
    事后才能核对一局到底用什么配置训的。
    """

    def _scalar(value):
        # SB3 2.9 的 learning_rate/clip_range 可能是 Schedule 对象
        current = getattr(value, "current_value", None)
        if callable(current):
            try:
                value = current()
            except Exception:
                return str(value)
        try:
            return float(value)
        except (TypeError, ValueError):
            return str(value)

    return {
        "n_steps": int(getattr(model, "n_steps", 0) or 0),
        "batch_size": int(getattr(model, "batch_size", 0) or 0),
        "n_epochs": int(getattr(model, "n_epochs", 0) or 0),
        "gamma": _scalar(getattr(model, "gamma", 0.0)),
        "gae_lambda": _scalar(getattr(model, "gae_lambda", 0.0)),
        "learning_rate": _scalar(getattr(model, "learning_rate", 0.0)),
        "clip_range": _scalar(getattr(model, "clip_range", 0.0)),
        "ent_coef": _scalar(getattr(model, "ent_coef", 0.0)),
        "vf_coef": _scalar(getattr(model, "vf_coef", 0.0)),
        "max_grad_norm": _scalar(getattr(model, "max_grad_norm", 0.0)),
    }


def _progress_payload(model, start: int) -> dict[str, Any]:
    """组装 ``--progress-json`` 的单行负载。

    纯函数（只读 model 的 logger / ep_info_buffer），便于脱离训练循环测试；
    NaN / ±inf 一律折算成 None，保证 ``json.dumps(allow_nan=False)`` 不会炸。
    """
    values = getattr(model.logger, "name_to_value", {}) or {}

    def _num(key, default=None):
        value = values.get(key)
        if value is None:
            return default
        try:
            value = float(value)
        except (TypeError, ValueError):
            return default
        if not math.isfinite(value):
            return default
        return value

    ep_rewards: list[float] = []
    ep_lengths: list[float] = []
    for item in getattr(model, "ep_info_buffer", []) or []:
        try:
            ep_rewards.append(float(item.get("r")))
            ep_lengths.append(float(item.get("l")))
        except (TypeError, ValueError, AttributeError):
            continue
    return {
        "event": "progress",
        "timesteps": int(model.num_timesteps),
        "start_timesteps": int(start),
        "n_updates": int(getattr(model, "_n_updates", 0) or 0),
        "ep_rew_mean": (sum(ep_rewards) / len(ep_rewards)
                        if ep_rewards else None),
        "ep_len_mean": (sum(ep_lengths) / len(ep_lengths)
                        if ep_lengths else None),
        "value_loss": _num("train/value_loss"),
        "policy_gradient_loss": _num("train/policy_gradient_loss"),
        "approx_kl": _num("train/approx_kl"),
        "entropy_loss": _num("train/entropy_loss"),
        "explained_variance": _num("train/explained_variance"),
        "fps": _num("time/fps"),
    }


def _model_metadata(*, args, wall_s: float, n_envs: int,
                    actual_timesteps: int) -> dict[str, Any]:
    import numpy
    import stable_baselines3
    import sb3_contrib
    import torch

    stage = _stage_from_index(args.stage)
    return {
        "level": "high",
        "problem": int(args.problem),
        "stage": int(args.stage),
        "stage_name": stage.name,
        "algorithm": "maskable_ppo",
        "feature_version": FEATURE_VERSION,
        "action_layout_version": ACTION_LAYOUT_VERSION,
        "reward_version": REWARD_VERSION,
        "obs_dim": HIGH_OBS_DIM,
        "act_dim": N_CANDIDATES,
        "seed": int(args.seed),
        "requested_timesteps": int(args.timesteps),
        "timesteps_semantics": ("additional" if (args.resume or args.init)
                                else "total"),
        "actual_timesteps": int(actual_timesteps),
        "n_envs": int(n_envs),
        "vec_env": args.vec_env,
        "config_profile": args.config_profile,
        "fidelity": _profile_fidelity(args.config_profile),
        "reward_profile": args.reward_profile,
        "local_selector": args.local_selector,
        "local_model": args.local_model,
        "checkpoint_every": int(getattr(args, "checkpoint_every", 0) or 0),
        "checkpoint_keep": int(getattr(args, "checkpoint_keep", 0) or 0),
        "checkpoint_dir": getattr(args, "checkpoint_dir", None),
        "recovery_every": int(getattr(args, "recovery_every", 0) or 0),
        "device": args.device,
        "wall_seconds": round(wall_s, 3),
        "max_wall_s": float(args.max_wall_s),
        "created_unix_s": time.time(),
        "python": sys.version.split()[0],
        "platform": platform.platform(),
        "torch": torch.__version__,
        "numpy": numpy.__version__,
        "stable_baselines3": stable_baselines3.__version__,
        "sb3_contrib": sb3_contrib.__version__,
        "mps_available": bool(getattr(getattr(torch.backends, "mps", None),
                                          "is_available", lambda: False)()),
    }


def _write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
    """原子写 JSON：同目录 ``.tmp-`` 前缀临时文件 + ``os.replace``。

    急停（SIGKILL / ``taskkill /F``）会绕过所有 Python 清理逻辑，保存途中被杀
    是完全可能的。原子替换保证**真实文件名下要么是旧的完整文件、要么是新的
    完整文件**，不会出现半截 JSON 让续训读到一半就崩。
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(f".tmp-{path.name}")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
                   encoding="utf-8")
    os.replace(tmp, path)


def _save_model_atomic(model, target: Path) -> Path:
    """把 SB3 模型原子地写到 ``target``，返回实际落盘路径。

    先写 ``.tmp-<name>``（同样以 ``.zip`` 结尾，SB3 不会再加后缀），完整写完后
    ``os.replace`` 覆盖目标。于是：

    - 保存途中被强杀 → 目标路径仍是上一份**完整**的快照，可继续复训；
    - 保存成功 → 目标路径立即变成新快照，不存在"看起来存在但打不开"的窗口。

    临时名带 ``.tmp-`` 前缀，因此不会被 ``*_ckpt_*.zip`` 的清理逻辑误当成正式
    checkpoint。
    """
    target = Path(target)
    target.parent.mkdir(parents=True, exist_ok=True)
    tmp = target.with_name(f".tmp-{target.name}")
    model.save(str(tmp))
    if not tmp.exists():        # SB3 对未知后缀可能自行补 .zip
        candidate = tmp.with_name(tmp.name + ".zip")
        if candidate.exists():
            tmp = candidate
    os.replace(tmp, target)
    return target


def _write_resume_hint(*, run_dir: Path, model_path: Path, timesteps: int,
                       reason: str, model=None) -> dict[str, Any]:
    """把"从哪继续"写进 ``<run_dir>/resume.json``。

    这是急停恢复的**权威指针**：无论上一刻是周期 checkpoint、滚动恢复点、
    优雅收尾还是异常兜底，这个文件总指向一份完整、可 ``--resume`` 的产物，
    并带上续训语义（``timesteps`` 是增量步数）与生效过的 PPO 超参。
    """
    hint: dict[str, Any] = {
        "event": "resume_hint",
        "model": str(model_path),
        "timesteps": int(timesteps),
        "timesteps_semantics": "additional",
        "reason": str(reason),
        "created_unix_s": time.time(),
        "resume_args": ["--resume", "--model", str(model_path)],
    }
    if model is not None:
        hint["ppo_hyperparams"] = _policy_hyperparams(model)
    _write_json_atomic(run_dir / "resume.json", hint)
    return hint


def _write_metadata(*, model_path: Path, args, wall_s: float, n_envs: int,
                    actual_timesteps: int, monitor_paths: list[str],
                    model=None,
                    extra: dict[str, Any] | None = None) -> dict[str, Any]:
    """写模型旁挂元数据（含 sha256），训练前端/审计都读它。"""
    metadata = _model_metadata(args=args, wall_s=wall_s, n_envs=n_envs,
                               actual_timesteps=actual_timesteps)
    if model is not None:
        metadata["ppo_hyperparams"] = _policy_hyperparams(model)
    metadata["model_sha256"] = _hash_file(model_path)
    metadata["monitor_files"] = list(monitor_paths)
    metadata["reward_note"] = (
        "reward_profile=time 使用 HRLReward.time_only（P0.4 纯时间口径）；"
        "reward_profile=prompt 使用原始 §5.2 塑形。模型元数据已固化口径。")
    if extra:
        metadata.update(extra)
    _write_json_atomic(model_path.with_suffix(".json"), metadata)
    return metadata


def _prune_checkpoints(directory: Path, stem: str, keep: int) -> list[str]:
    """只保留最近 ``keep`` 个 checkpoint（zip+json 同删），返回被删的 zip 路径。

    只匹配 ``<stem>_<步数>.zip``；滚动恢复点 ``<stem>_latest.zip`` 与临时文件
    ``.tmp-*.zip`` 都不匹配，因此永远不会被历史清理误删。
    """
    files = sorted((path for path in directory.glob(f"{stem}_[0-9]*.zip")
                    if path.stem.rsplit("_", 1)[-1].isdigit()),
                   key=lambda path: path.stat().st_mtime, reverse=True)
    removed: list[str] = []
    for stale in files[max(1, int(keep)):]:
        for candidate in (stale, stale.with_suffix(".json")):
            if candidate.exists():
                candidate.unlink()
        removed.append(str(stale))
    return removed


def _evaluate_env(model, *, problem: int, stage_index: int,
                  episodes: int, seed: int, profile: str,
                  reward_profile: str) -> dict[str, Any]:
    from mcm26b.hrl.hierarchical_env import HierarchicalEnv

    stage = _stage_from_index(stage_index)
    env = HierarchicalEnv(problem=problem, stage=stage,
                          fidelity=_profile_fidelity(profile),
                          max_episode_steps=300, real_budget_s=1800.0,
                          reward=_reward_from_profile(reward_profile))
    wins = 0
    rows = []
    if episodes <= 0:
        env.close()
        return {"episodes": 0, "successes": 0, "success_rate": None, "rows": []}
    for episode in range(episodes):
        obs, info = env.reset(seed=seed + episode)
        done = False
        steps = 0
        while not done:
            action, _state = model.predict(
                obs, action_masks=env.action_masks(), deterministic=True)
            obs, _, terminated, truncated, info = env.step(int(action))
            done = bool(terminated or truncated)
            steps += 1
        # 课程阶段的小场景只需证明全部频道 accounted；完整阶段至少应清掉
        # 场景最少的题面源数。这里不读取真值，只使用环境自身统计。
        success = bool(
            info.get("accounted", 0) >= 20
            and info.get("cleared", 0) >= stage.sources_min
            and not info.get("anomalies")
        )
        wins += int(success)
        rows.append({
            "episode": episode,
            "seed": seed + episode,
            "steps": steps,
            "cleared": info.get("cleared", 0),
            "accounted": info.get("accounted", 0),
            "virtual_time_s": info.get("episode_virtual_cost_s"),
            "stop_reason": info.get("stop_reason"),
            "anomalies": info.get("anomalies"),
            "success": success,
        })
    env.close()
    return {
        "episodes": len(rows),
        "successes": wins,
        "success_rate": wins / len(rows) if rows else 0.0,
        "rows": rows,
    }


def _evaluate_manifest(model_path: Path, manifest_path: str, *,
                       limit: int) -> dict[str, Any]:
    """可选：用 benchmark_hrl 的同一套真值隔离评测在固定 manifest 上跑已训模型。"""
    scripts_dir = Path(__file__).resolve().parent
    if str(scripts_dir) not in sys.path:
        sys.path.insert(0, str(scripts_dir))
    import benchmark_hrl as bench  # scripts/benchmark_hrl.py

    manifest = json.loads(Path(manifest_path).read_text(encoding="utf-8"))
    scenes = bench.scenes_from_manifest(manifest)
    if limit > 0:
        scenes = scenes[:limit]
    rows = []
    for scene in scenes:
        run_dir = Path(model_path).parent / "eval-manifest" / scene.scene_id
        result = bench.run_scene_variant(
            scene, "hrl", config=ProblemConfig(),
            hrl_model=str(model_path), run_dir=run_dir)
        rows.append(result["row"])
    aggregate = bench._aggregate(rows)
    return aggregate


def run_training(args) -> int:
    import numpy as np
    import torch
    from sb3_contrib import MaskablePPO
    from stable_baselines3.common.callbacks import BaseCallback
    from stable_baselines3.common.utils import set_random_seed

    torch.set_num_threads(1)
    set_random_seed(args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    model_path = Path(args.model)
    if model_path.exists() and not (args.resume or args.force):
        print(f"模型已存在：{model_path}；如需覆盖请加 --force，续训请加 --resume。",
              file=sys.stderr)
        return 2
    model_path.parent.mkdir(parents=True, exist_ok=True)
    run_dir = model_path.parent
    monitor_paths = [
        str(run_dir / f"high_curve_{model_path.stem}_env{rank}.monitor.csv")
        for rank in range(args.envs)
    ]
    for path in monitor_paths:
        Path(path).parent.mkdir(parents=True, exist_ok=True)

    vec_env = _build_vec_env(args, monitor_paths)

    if args.resume:
        if not model_path.exists():
            print(f"--resume 指定的模型不存在：{model_path}", file=sys.stderr)
            return 2
        model = MaskablePPO.load(str(model_path), env=vec_env,
                                 device=args.device,
                                 print_system_info=False)
    elif args.init:
        init_path = Path(args.init)
        if not init_path.exists():
            print(f"--init 指定的模型不存在：{init_path}", file=sys.stderr)
            return 2
        model = MaskablePPO.load(str(init_path), env=vec_env,
                                 device=args.device,
                                 print_system_info=False)
    else:
        model = MaskablePPO(
            "MlpPolicy",
            vec_env,
            policy_kwargs={"net_arch": [128, 128]},
            n_steps=args.rollout_steps,
            batch_size=args.batch_size,
            n_epochs=args.n_epochs,
            gamma=1.0,
            gae_lambda=0.95,
            learning_rate=3e-4,
            clip_range=0.2,
            ent_coef=0.005,
            vf_coef=0.5,
            max_grad_norm=0.5,
            seed=args.seed,
            device=args.device,
            verbose=args.verbose,
        )
        if args.init:
            pass  # 已在上面处理；此分支仅保持语义清晰

    # resume/init 会恢复 checkpoint 内的 PPO 超参，命令行同名三项目前被忽略；
    # 显式说出差异，避免"以为换了 n_steps 其实没换"。
    if args.resume or args.init:
        effective = _policy_hyperparams(model)
        requested = {"n_steps": int(args.rollout_steps),
                     "batch_size": int(args.batch_size),
                     "n_epochs": int(args.n_epochs)}
        if {key: effective[key] for key in requested} != requested:
            print(json.dumps({
                "event": "hyperparams_from_checkpoint",
                "effective": {key: effective[key] for key in requested},
                "requested": requested,
                "note": ("resume/init 沿用 checkpoint 内的 PPO 超参，"
                         "命令行 rollout-steps/batch-size/n-epochs 被忽略")},
                ensure_ascii=False), file=sys.stderr, flush=True)

    # ------------------------------------------------------ 优雅暂停 / 停止
    # 主通道是控制文件（跨平台）；信号是兜底：直接 kill 本进程时也走同一段收尾。
    stop_state: dict[str, Any] = {"requested": False, "signal": None}
    previous_handlers: dict[int, Any] = {}

    def _request_stop(source: str) -> None:
        if stop_state["requested"]:
            return
        stop_state["requested"] = True
        stop_state["signal"] = source
        print(json.dumps({"event": "stop_requested", "signal": source}), flush=True)

    def _on_stop_signal(signum, _frame):
        _request_stop(signal.Signals(signum).name)

    for _signum in _stop_signals():
        try:
            previous_handlers[_signum] = signal.signal(_signum, _on_stop_signal)
        except (ValueError, OSError):  # 非主线程/平台不支持：不阻塞训练
            pass

    class _ControlChannel(BaseCallback):
        """把 :class:`_ControlPoll` 接到 SB3 的 step 回调上。

        为什么不用信号：Windows 没有 SIGSTOP，且 ``os.kill(SIGTERM)`` 会退化成
        ``TerminateProcess`` 硬杀，收尾分支根本不会执行。控制文件在 POSIX 与
        Windows 上都只是普通文件 IO，语义完全一致。

        暂停是**协作式**的：在 ``_on_step`` 内阻塞轮询，训练因此停在 step 边界，
        进程仍存活、显存仍占用，但不消耗算力；收到 stop 立刻转入收尾落盘。
        """

        def __init__(self, path: Path, *, poll_every: int = 32,
                     poll_s: float = 0.2):
            super().__init__()
            self.poller = _ControlPoll(path, poll_every=poll_every,
                                       poll_s=poll_s)

        def paused_seconds(self) -> float:
            return self.poller.paused_total_s

        def _event(self, name: str) -> None:
            print(json.dumps({"event": name,
                              "timesteps": int(self.model.num_timesteps)}),
                  flush=True)

        def _on_step(self) -> bool:
            state = self.poller.poll()
            if state.stop:
                _request_stop("control_file")
                return False
            if state.paused:
                stopped = self.poller.wait_while_paused(
                    lambda: self._event("paused"),
                    lambda: self._event("resumed"))
                if stopped:
                    _request_stop("control_file")
                    return False
                return not stop_state["requested"]
            return True

    control_path = getattr(args, "control_file", None)
    control_cb = _ControlChannel(Path(control_path)) if control_path else None

    class _GracefulStop(BaseCallback):
        """收到停止请求后在最近一次 step 边界结束 rollout，让模型能落盘。"""

        def _on_step(self) -> bool:
            if stop_state["requested"]:
                print(json.dumps({"event": "graceful_stop",
                                  "signal": stop_state["signal"],
                                  "timesteps": int(self.model.num_timesteps)}),
                      flush=True)
                return False
            return True

    class _WallClockStop(BaseCallback):
        """墙钟硬上限；暂停时间不计入，避免"挂起一小时回来就判超时"。"""

        def __init__(self, max_wall_s: float, paused_clock=None):
            super().__init__()
            self.max_wall_s = max_wall_s
            self.start = time.monotonic()
            self.paused_clock = paused_clock

        def paused_seconds(self) -> float:
            return float(self.paused_clock()) if self.paused_clock else 0.0

        def _on_step(self) -> bool:
            elapsed = time.monotonic() - self.start - self.paused_seconds()
            if elapsed >= self.max_wall_s:
                print(json.dumps({"event": "wall_clock_stop",
                                  "elapsed_s": elapsed}))
                return False
            return True

    class _ProgressJSON(BaseCallback):
        """向训练前端输出单行 JSON 进度；不影响 SB3 自身 logger。

        前端首个 curve 点越早出现越好，因此开头几步按更小间隔发送（前 8 个
        进度点逐步放大到 ``every_steps``），之后稳定在 ``every_steps``。
        这样短程 smoke 也能看到完整曲线，长训不会刷屏。
        """

        def __init__(self, every_steps: int):
            super().__init__()
            self.every_steps = max(1, int(every_steps))
            # init/resume 时计数器从上一阶段继续，前端要用增量算进度
            self.start = int(model.num_timesteps)
            self.last_emit = self.start
            self.emitted = 0

        def _next_gap(self) -> int:
            warmup = max(0.2, 1.0 - self.emitted / 8.0)
            return max(1, min(self.every_steps, int(round(self.every_steps * warmup))))

        def _on_step(self) -> bool:
            current = int(self.model.num_timesteps)
            if current - self.last_emit < self._next_gap():
                return True
            self.last_emit = current
            self.emitted += 1
            payload = _progress_payload(self.model, self.start)
            print(json.dumps(payload, ensure_ascii=False, allow_nan=False),
                  flush=True)
            return True

    started = time.monotonic()
    checkpoints_dir = (Path(args.checkpoint_dir) if args.checkpoint_dir
                       else run_dir / "checkpoints")
    checkpoint_stem = f"{model_path.stem}_ckpt"

    class _CheckpointJSON(BaseCallback):
        """可续训落盘：滚动恢复点 + 周期 checkpoint + resume.json 指针。

        - **滚动恢复点** ``<stem>_latest.zip``：每 ``recovery_every`` 个环境步
          原子覆盖写一次，训练开始前也先写一份。它保证"急停"（SIGKILL /
          ``taskkill /F``，没有任何 Python 清理机会）之后**最多只丢
          ``recovery_every`` 个环境步**，而且文件始终完整可加载。
        - **周期 checkpoint** ``<stem>_<步数>.zip``：每 ``every_steps`` 步留一份，
          只保留最近 ``keep`` 个，用于返祖对比与多档恢复。
        - 每次落盘都刷新 ``<run_dir>/resume.json``：无论怎么停，都有一份权威的
          "从哪继续"指针（含增量语义与生效过的 PPO 超参）。
        """

        def __init__(self, every_steps: int, keep: int, recovery_every: int):
            super().__init__()
            self.every_steps = max(0, int(every_steps))
            self.keep = max(1, int(keep))
            self.recovery_every = max(0, int(recovery_every))
            self.last_saved = int(model.num_timesteps)
            self.last_recovery = int(model.num_timesteps)
            self.saved: list[str] = []
            self.latest_path = checkpoints_dir / f"{checkpoint_stem}_latest.zip"

        def write_recovery(self, reason: str) -> None:
            """立刻写一份滚动恢复点；训练开始前调用即"起步基线"。

            用闭包里的 ``model`` 而不是 ``self.model``：``BaseCallback.model``
            要等 ``learn()`` 时才被赋值，而起步基线恰恰在 ``learn()`` 之前写。
            """
            current = int(model.num_timesteps)
            self.last_recovery = current
            self._write_recovery(current, reason)

        def _on_step(self) -> bool:
            current = int(model.num_timesteps)
            if (self.recovery_every > 0
                    and current - self.last_recovery >= self.recovery_every):
                self.last_recovery = current
                self._write_recovery(current, "interval")
            if (self.every_steps > 0
                    and current - self.last_saved >= self.every_steps):
                self.last_saved = current
                self._write_interval(current)
            return True

        def _persist(self, target: Path, current: int, reason: str,
                     extra: dict[str, Any]) -> None:
            """原子落盘 zip + 元数据，并刷新 resume.json 指针。"""
            _save_model_atomic(model, target)
            _write_metadata(
                model_path=target, args=args,
                wall_s=time.monotonic() - started, n_envs=args.envs,
                actual_timesteps=current, monitor_paths=monitor_paths,
                model=model,
                extra={"checkpoint_timesteps": current,
                       "checkpoint_reason": reason,
                       "checkpoint_of": str(model_path), **extra})
            _write_resume_hint(run_dir=run_dir, model_path=target,
                               timesteps=current, reason=reason,
                               model=model)

        def _write_recovery(self, current: int, reason: str) -> None:
            checkpoints_dir.mkdir(parents=True, exist_ok=True)
            self._persist(self.latest_path, current, reason,
                          {"recovery": True})
            print(json.dumps({"event": "recovery_checkpoint",
                              "reason": reason, "timesteps": current,
                              "model": str(self.latest_path)},
                             ensure_ascii=False), flush=True)

        def _write_interval(self, current: int) -> None:
            checkpoints_dir.mkdir(parents=True, exist_ok=True)
            target = checkpoints_dir / f"{checkpoint_stem}_{current:09d}.zip"
            self._persist(target, current, "interval", {"checkpoint": True})
            removed = _prune_checkpoints(checkpoints_dir, checkpoint_stem,
                                         self.keep)
            self.saved.append(str(target))
            print(json.dumps({"event": "checkpoint", "timesteps": current,
                              "model": str(target), "pruned": removed},
                             ensure_ascii=False), flush=True)

    progress_every = (args.progress_every if args.progress_every > 0
                      else args.rollout_steps * args.envs)
    checkpoint_cb = _CheckpointJSON(args.checkpoint_every, args.checkpoint_keep,
                                    args.recovery_every)
    if args.recovery_every > 0:
        # 起步基线：训练还没跑就急停，也至少有一份可加载、可续训的产物
        checkpoint_cb.write_recovery("baseline")
    # 顺序有语义：控制通道最先（暂停/停止要在本步立刻生效），进度次之，
    # 墙钟上限要能读到控制通道累计的暂停秒数，checkpoint 放最后。
    callbacks: list[Any] = []
    if control_cb is not None:
        callbacks.append(control_cb)
    if args.progress_json:
        callbacks.append(_ProgressJSON(progress_every))
    callbacks.append(_GracefulStop())
    callbacks.append(_WallClockStop(
        args.max_wall_s,
        paused_clock=(control_cb.paused_seconds if control_cb is not None
                      else None)))
    callbacks.append(checkpoint_cb)

    actual_timesteps = int(model.num_timesteps)
    try:
        model.learn(total_timesteps=args.timesteps,
                    callback=callbacks if len(callbacks) > 1 else callbacks[0],
                    progress_bar=False,
                    reset_num_timesteps=not (args.resume or args.init))
    except BaseException as exc:  # 训练失败也必须留下可续训产物，不静默
        failure = {
            "event": "training_failed",
            "error": f"{type(exc).__name__}: {exc}",
            "model": str(model_path),
            "timesteps": int(model.num_timesteps),
        }
        emergency = run_dir / f"{model_path.stem}_emergency.zip"
        try:
            _save_model_atomic(model, emergency)
            _write_metadata(model_path=emergency, args=args,
                            wall_s=time.monotonic() - started,
                            n_envs=args.envs,
                            actual_timesteps=int(model.num_timesteps),
                            monitor_paths=monitor_paths, model=model,
                            extra={"emergency": True,
                                   "error": failure["error"]})
            _write_resume_hint(run_dir=run_dir, model_path=emergency,
                               timesteps=int(model.num_timesteps),
                               reason="emergency", model=model)
            failure["emergency_model"] = str(emergency)
        except Exception as save_exc:  # 落盘失败不能再掩盖原始异常
            failure["emergency_save_error"] = (
                f"{type(save_exc).__name__}: {save_exc}")
        _write_json_atomic(run_dir / "training_failure.json", failure)
        print(json.dumps(failure, ensure_ascii=False), file=sys.stderr,
              flush=True)
        raise
    wall_s = time.monotonic() - started
    actual_timesteps = int(model.num_timesteps)
    stopped_by_signal = bool(stop_state["requested"])
    _save_model_atomic(model, model_path)
    metadata = _write_metadata(
        model_path=model_path, args=args, wall_s=wall_s, n_envs=args.envs,
        actual_timesteps=actual_timesteps, monitor_paths=monitor_paths,
        model=model,
        extra={"stopped_by_signal": stopped_by_signal,
               "stop_signal": stop_state["signal"],
               "checkpoint_files": checkpoint_cb.saved,
               "recovery_model": str(checkpoint_cb.latest_path)})
    # 最终模型也刷新一次 resume.json：优雅收尾后"从哪继续"指向最终模型
    _write_resume_hint(run_dir=run_dir, model_path=model_path,
                       timesteps=actual_timesteps,
                       reason=("stopped" if stopped_by_signal else "finished"),
                       model=model)
    metadata_path = model_path.with_suffix(".json")
    vec_env.close()
    # 收尾阶段恢复默认处理器：再次 Ctrl+C / 停止可立即强杀，不再等待评估。
    for _signum, _handler in previous_handlers.items():
        try:
            signal.signal(_signum, _handler)
        except (ValueError, OSError):
            pass
    if stopped_by_signal:
        print(json.dumps({"event": "saved_after_stop",
                          "model": str(model_path),
                          "timesteps": actual_timesteps}), flush=True)

    evaluation = _evaluate_env(
        model, problem=args.problem, stage_index=args.stage,
        episodes=args.eval_episodes, seed=args.seed + 1_000_000,
        profile=args.config_profile, reward_profile=args.reward_profile)
    manifest_eval = None
    if args.eval_manifest:
        try:
            manifest_eval = _evaluate_manifest(
                model_path, args.eval_manifest, limit=args.eval_manifest_limit)
        except Exception as exc:
            manifest_eval = {"error": f"{type(exc).__name__}: {exc}"}
    metadata["evaluation"] = evaluation
    metadata["eval_manifest"] = manifest_eval
    metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2),
                             encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))
    return 0


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--problem", type=int, choices=(3, 4), required=True)
    parser.add_argument("--stage", type=int, choices=(1, 2, 3, 4),
                        default=4)
    parser.add_argument("--timesteps", type=int, required=True)
    parser.add_argument("--seed", type=int, default=320001)
    parser.add_argument("--envs", type=int, default=1)
    parser.add_argument("--vec-env", choices=("dummy", "subproc"), default="dummy")
    parser.add_argument("--device", choices=("cpu", "cuda", "mps", "auto"),
                        default="auto")
    parser.add_argument("--config-profile",
                        choices=("production", "coarse-pretrain"),
                        default="production")
    parser.add_argument("--model", required=True)
    parser.add_argument("--resume", action="store_true",
                        help="从 --model 续训（沿用 checkpoint 内的 PPO 超参）")
    parser.add_argument("--init", default=None,
                        help="可选：Q3→Q4 或继续预训练权重路径")
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--eval-manifest", default=None)
    parser.add_argument("--eval-manifest-limit", type=int, default=0)
    parser.add_argument("--eval-episodes", type=int, default=5)
    parser.add_argument("--local-selector", choices=("geometry", "rl"),
                        default="geometry")
    parser.add_argument("--local-model", default=None)
    parser.add_argument("--rollout-steps", type=int, default=256)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--n-epochs", type=int, default=10)
    parser.add_argument("--max-wall-s", type=float, default=1200.0,
                        help="训练墙钟硬上限；默认 20 分钟")
    parser.add_argument("--reward-profile", choices=("time", "prompt"),
                        default="time",
                        help="time=P0.4 纯时间口径；prompt=原始 §5.2 塑形对照")
    parser.add_argument("--progress-json", action="store_true",
                        help="向训练前端输出单行 JSON 进度（UI 自动启用）")
    parser.add_argument("--control-file", default=None,
                        help="训练前端控制文件（JSON：paused/stop）；"
                             "跨平台暂停与优雅停止的主通道")
    parser.add_argument("--progress-every", type=int, default=0,
                        help="JSON 进度间隔环境步数；0=每个 rollout")
    parser.add_argument("--checkpoint-every", type=int, default=20000,
                        help="每 N 个环境步留一个历史 checkpoint；0=关闭（默认 20000）")
    parser.add_argument("--recovery-every", type=int, default=2000,
                        help="每 N 个环境步原子覆盖写一份滚动恢复点 "
                             "<模型目录>/checkpoints/<stem>_latest.zip；"
                             "急停后最多丢 N 步。0=关闭（默认 2000）")
    parser.add_argument("--checkpoint-keep", type=int, default=3,
                        help="checkpoint 保留最近 K 个（默认 3）")
    parser.add_argument("--checkpoint-dir", default=None,
                        help="checkpoint 目录；默认 <模型目录>/checkpoints")
    parser.add_argument("--verbose", type=int, default=1)
    args = parser.parse_args(argv)
    if args.timesteps <= 0 or args.envs <= 0:
        parser.error("timesteps 与 envs 必须为正")
    if args.envs * args.rollout_steps < 2:
        parser.error("PPO rollout 至少需要 2 个样本")
    if args.eval_episodes < 0:
        parser.error("eval-episodes 不能为负")
    if args.checkpoint_every < 0:
        parser.error("checkpoint-every 不能为负")
    if args.recovery_every < 0:
        parser.error("recovery-every 不能为负")
    if args.checkpoint_keep < 1:
        parser.error("checkpoint-keep 至少为 1")
    return run_training(args)


if __name__ == "__main__":
    raise SystemExit(main())