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

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

落盘保证（长训必需）：
- ``--checkpoint-every`` 每 N 个环境步在 ``<模型目录>/checkpoints/`` 落一个可续训
  的 checkpoint（zip + 同名 json 元数据），只保留最近 ``--checkpoint-keep`` 个；
- 训练前端的「停止」是对**进程组**发 SIGTERM，因此 subproc worker 在 :func:`_make_env`
  里显式忽略 SIGTERM/SIGINT，主进程收到信号后在 step 边界收尾并保存模型，
  不会被 worker 先行退出打断；
- 任何异常（含被强杀之外的失败）都会先落一份 ``*_emergency.zip`` 再抛出，
  不再出现"跑了很久什么都没有"。

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 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, training_config
from mcm26b.hrl.constants import (ACTION_LAYOUT_VERSION, FEATURE_VERSION,
                                  HIGH_OBS_DIM, N_CANDIDATES, REWARD_VERSION)


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 _ignore_stop_signals_in_worker() -> None:
    """subproc worker 忽略 SIGTERM/SIGINT，保证主进程能优雅收尾落盘。

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


def _make_env(problem: int, stage_index: int, seed: 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, args.seed + rank,
                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 _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,
        "reward_version": REWARD_VERSION,
        "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),
        "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_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)
    model_path.with_suffix(".json").write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    return metadata


def _prune_checkpoints(directory: Path, stem: str, keep: int) -> list[str]:
    """只保留最近 ``keep`` 个 checkpoint（zip+json 同删），返回被删的 zip 路径。"""
    files = sorted(directory.glob(f"{stem}_*.zip"),
                   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, reward, 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, problem: 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)

    # ------------------------------------------------------------ 优雅停止
    # 训练前端「停止」= killpg(SIGTERM)：worker 已在 _make_env 里忽略该信号，
    # 主进程在此装处理器，等当前 step 结束后收尾落盘。
    stop_state: dict[str, Any] = {"requested": False, "signal": None}
    previous_handlers: dict[int, Any] = {}

    def _on_stop_signal(signum, _frame):
        stop_state["requested"] = True
        stop_state["signal"] = signal.Signals(signum).name
        print(json.dumps({"event": "stop_requested",
                          "signal": signal.Signals(signum).name}), flush=True)

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

    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):
            super().__init__()
            self.max_wall_s = max_wall_s
            self.start = time.monotonic()

        def _on_step(self) -> bool:
            if time.monotonic() - self.start >= self.max_wall_s:
                print(json.dumps({"event": "wall_clock_stop",
                                  "elapsed_s": time.monotonic() - self.start}))
                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
            values = getattr(self.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 value != value or value in (float("inf"), float("-inf")):
                    return default
                return value

            ep_rewards = []
            ep_lengths = []
            for item in getattr(self.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
            payload = {
                "event": "progress",
                "timesteps": current,
                "start_timesteps": self.start,
                "n_updates": int(getattr(self.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"),
            }
            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):
        """每 N 个环境步落一个可续训 checkpoint（zip + 同名元数据）。

        只保留最近 ``keep`` 个；checkpoint 与最终模型共用同名 json 契约，
        因此 ``--resume <checkpoint.zip>`` 直接可用。
        """

        def __init__(self, every_steps: int, keep: int):
            super().__init__()
            self.every_steps = max(1, int(every_steps))
            self.keep = max(1, int(keep))
            self.last_saved = int(model.num_timesteps)
            self.saved: list[str] = []

        def _on_step(self) -> bool:
            current = int(self.model.num_timesteps)
            if current - self.last_saved < self.every_steps:
                return True
            self.last_saved = current
            self._write(current)
            return True

        def _write(self, current: int) -> None:
            checkpoints_dir.mkdir(parents=True, exist_ok=True)
            target = checkpoints_dir / f"{checkpoint_stem}_{current:09d}.zip"
            self.model.save(str(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=self.model,
                extra={"checkpoint": True, "checkpoint_timesteps": current,
                       "checkpoint_of": str(model_path)})
            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 = None
    if args.checkpoint_every > 0:
        checkpoint_cb = _CheckpointJSON(args.checkpoint_every,
                                        args.checkpoint_keep)
    callbacks = [_GracefulStop(), _WallClockStop(args.max_wall_s)]
    if checkpoint_cb is not None:
        callbacks.append(checkpoint_cb)
    if args.progress_json:
        callbacks.insert(0, _ProgressJSON(progress_every))

    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:
            model.save(str(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"]})
            failure["emergency_model"] = str(emergency)
        except Exception as save_exc:  # 落盘失败不能再掩盖原始异常
            failure["emergency_save_error"] = (
                f"{type(save_exc).__name__}: {save_exc}")
        (run_dir / "training_failure.json").write_text(
            json.dumps(failure, ensure_ascii=False, indent=2), encoding="utf-8")
        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"])
    model.save(str(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
                                    if checkpoint_cb is not None else [])})
    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,
                problem=args.problem)
        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("--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("--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.checkpoint_keep < 1:
        parser.error("checkpoint-keep 至少为 1")
    return run_training(args)


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