"""高层训练落盘保证：checkpoint、控制文件优雅停止、信号兜底、worker 信号屏蔽。

背景：训练前端的「停止」以前直接向进程组发 SIGTERM，而旧脚本只在 ``learn()``
正常返回后才保存，导致"停一下就全丢"；在 Windows 上 ``os.killpg`` 更是直接
``AttributeError``。现在停止/暂停走**控制文件**（跨平台），信号只作兜底。这里把
保证固定成回归测试：
1. 定期 checkpoint 可续训（zip + 同名元数据）；
2. 控制文件发出 stop 后仍能保存最终模型与元数据（Windows 同样成立）；
3. POSIX 下进程组 SIGTERM 也能优雅收尾；
4. subproc worker 忽略停止信号，不抢在主进程收尾之前退出。
"""
from __future__ import annotations

import importlib.util
import json
import multiprocessing
import os
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path

import pytest

from mcm26b.process_control import spawn_kwargs, terminate_process_tree, write_control


def _train_module():
    """按路径加载训练脚本（不触发 torch 导入，torch 只在函数内部惰性导入）。"""
    path = Path(__file__).parents[1] / "scripts" / "train_hrl_high.py"
    spec = importlib.util.spec_from_file_location("train_hrl_high_ckpt", path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _require_rl():
    pytest.importorskip("stable_baselines3")
    pytest.importorskip("sb3_contrib")
    pytest.importorskip("torch")


class _Run:
    """一次子进程训练：stdout 走文件，避免管道写满把子进程卡死。

    与训练前端一致：用控制文件请求暂停/停止，进程组隔离走
    :func:`mcm26b.process_control.spawn_kwargs`（POSIX ``setsid`` /
    Windows ``CREATE_NEW_PROCESS_GROUP``）。
    """

    def __init__(self, run_dir: Path, *extra: str):
        script = Path(__file__).parents[1] / "scripts" / "train_hrl_high.py"
        self.run_dir = run_dir
        self.log_path = run_dir / "train_stdout.log"
        self.log_file = self.log_path.open("w", encoding="utf-8")
        self.control_path = run_dir / "control.json"
        write_control(self.control_path)
        cmd = [sys.executable, str(script),
               "--problem", "3", "--stage", "1", "--timesteps", "200000",
               "--seed", "7", "--envs", "2", "--vec-env", "subproc",
               "--device", "cpu", "--config-profile", "coarse-pretrain",
               "--reward-profile", "time", "--eval-episodes", "1",
               "--rollout-steps", "32", "--batch-size", "32", "--n-epochs", "1",
               "--max-wall-s", "600", "--force", "--verbose", "0",
               "--control-file", str(self.control_path),
               "--model", str(run_dir / "q3.zip"), *extra]
        self.proc = subprocess.Popen(   # pylint: disable=consider-using-with
            cmd, stdout=self.log_file, stderr=subprocess.STDOUT, text=True,
            **spawn_kwargs())

    def checkpoints(self) -> list[Path]:
        """历史 checkpoint（``<stem>_<步数>.zip``）。

        刻意排除滚动恢复点 ``<stem>_latest.zip``：它随写随覆盖，不是"第几个
        checkpoint"，混进来会让序号断言失去意义。
        """
        found = (self.run_dir / "checkpoints").glob("q3_ckpt_[0-9]*.zip")
        return sorted(path for path in found
                      if path.stem.rsplit("_", 1)[-1].isdigit())

    def recovery(self) -> Path:
        return self.run_dir / "checkpoints" / "q3_ckpt_latest.zip"

    def resume_hint(self) -> dict:
        return json.loads((self.run_dir / "resume.json").read_text(
            encoding="utf-8"))

    def wait_for(self, predicate, timeout_s: float) -> bool:
        deadline = time.monotonic() + timeout_s
        while time.monotonic() < deadline:
            if predicate():
                return True
            if self.proc.poll() is not None:
                return predicate()
            time.sleep(0.4)
        return predicate()

    def logs(self) -> str:
        try:
            return self.log_path.read_text(encoding="utf-8", errors="replace")
        except OSError:  # pragma: no cover - 文件被并发写时的极端情况
            return ""

    def pause(self) -> None:
        write_control(self.control_path, paused=True)

    def resume(self) -> None:
        write_control(self.control_path, paused=False)

    def _finish(self) -> str:
        self.log_file.close()
        return self.logs()

    def stop(self) -> str:
        """复现训练前端的停止语义：写控制文件请求优雅收尾（跨平台）。"""
        write_control(self.control_path, stop=True)
        forced = terminate_process_tree(self.proc, grace_s=240.0)
        if forced:
            self.log_file.close()
            pytest.fail("控制文件停止请求后训练进程未在 240s 内退出（优雅停止失败）")
        return self._finish()

    def stop_by_signal(self) -> str:
        """POSIX 兜底路径：对整个进程组发 SIGTERM（Windows 不支持 killpg）。"""
        if os.name == "nt":  # pragma: no cover - 仅在 Windows 上走到
            pytest.skip("Windows 无 os.killpg，信号路径不适用")
        try:
            os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
            self.proc.terminate()
        try:
            self.proc.wait(timeout=240)
        except subprocess.TimeoutExpired:
            self.proc.kill()
            self.proc.wait(timeout=30)
            self.log_file.close()
            pytest.fail("SIGTERM 后训练进程未在 240s 内退出（优雅停止失败）")
        return self._finish()

    def hard_kill(self) -> None:
        """模拟"急停"：立即强杀整棵进程树，不给训练脚本任何收尾机会。"""
        terminate_process_tree(self.proc, grace_s=0.0)
        self.log_file.close()


def _run_dir(tmp_path: Path) -> Path:
    run_dir = tmp_path / "run"
    run_dir.mkdir(parents=True, exist_ok=True)
    return run_dir


# --------------------------------------------------------------- 纯函数回归
def test_prune_checkpoints_keeps_newest(tmp_path):
    module = _train_module()
    directory = tmp_path / "checkpoints"
    directory.mkdir()
    for index in range(5):
        for suffix in (".zip", ".json"):
            target = directory / f"q3_ckpt_{index:09d}{suffix}"
            target.write_bytes(b"x")
            stamp = 1_700_000_000 + index
            os.utime(target, (stamp, stamp))
    removed = module._prune_checkpoints(directory, "q3_ckpt", 2)
    assert len(removed) == 3
    remaining = sorted(path.name for path in directory.iterdir())
    assert remaining == ["q3_ckpt_000000003.json", "q3_ckpt_000000003.zip",
                         "q3_ckpt_000000004.json", "q3_ckpt_000000004.zip"]


def test_worker_ignores_stop_signals_but_main_does_not(monkeypatch):
    module = _train_module()
    original = signal.getsignal(signal.SIGTERM)

    class _Worker:
        name = "SpawnProcess-3"

    monkeypatch.setattr(multiprocessing, "current_process", _Worker)
    try:
        module._ignore_stop_signals_in_worker()
        assert signal.getsignal(signal.SIGTERM) == signal.SIG_IGN
        assert signal.getsignal(signal.SIGINT) == signal.SIG_IGN
    finally:
        signal.signal(signal.SIGTERM, original)
        signal.signal(signal.SIGINT, signal.default_int_handler)

    class _Main:
        name = "MainProcess"

    monkeypatch.setattr(multiprocessing, "current_process", _Main)
    module._ignore_stop_signals_in_worker()
    assert signal.getsignal(signal.SIGTERM) == original


def test_policy_hyperparams_records_effective_values():
    """元数据必须记实际生效的超参（resume 时来自 checkpoint，而非命令行）。"""
    module = _train_module()

    class _Schedule:
        def __init__(self, value):
            self._value = value

        def current_value(self):
            return self._value

    class _Model:
        n_steps = 256
        batch_size = 256
        n_epochs = 10
        gamma = 1.0
        gae_lambda = 0.95
        learning_rate = _Schedule(3e-4)
        clip_range = _Schedule(0.2)
        ent_coef = 0.005
        vf_coef = 0.5
        max_grad_norm = 0.5

    hyper = module._policy_hyperparams(_Model())
    assert hyper["n_steps"] == 256 and hyper["batch_size"] == 256
    assert hyper["n_epochs"] == 10
    assert hyper["learning_rate"] == 3e-4
    assert hyper["clip_range"] == 0.2
    assert hyper["gamma"] == 1.0 and hyper["ent_coef"] == 0.005
    assert module._policy_hyperparams(object())["n_steps"] == 0


def test_control_poll_state_machine(tmp_path):
    """控制文件轮询状态机：抽样读盘、暂停阻塞、stop 打断、暂停时长累计。"""
    module = _train_module()
    path = tmp_path / "control.json"
    write_control(path, paused=True)
    poll = module._ControlPoll(path, poll_every=1, poll_s=0.05)
    assert poll.poll().paused is True

    events = []

    def _resume_later():
        time.sleep(0.15)
        write_control(path, paused=False)

    threading.Thread(target=_resume_later, daemon=True).start()
    stopped = poll.wait_while_paused(lambda: events.append("pause"),
                                     lambda: events.append("resume"))
    assert stopped is False
    assert events == ["pause", "resume"]
    assert poll.paused_total_s >= 0.1

    # 暂停期间收到 stop：wait 必须返回 True，让调用方转入收尾
    write_control(path, paused=True)

    def _stop_later():
        time.sleep(0.15)
        write_control(path, stop=True)

    threading.Thread(target=_stop_later, daemon=True).start()
    assert poll.wait_while_paused(lambda: None, lambda: None) is True
    assert poll.poll().stop is True


def test_atomic_model_save_keeps_previous_snapshot_on_failure(tmp_path):
    """保存途中失败/被杀：目标路径仍是上一份完整快照，不会出现半截 zip。"""
    module = _train_module()
    target = tmp_path / "q3_ckpt_latest.zip"
    target.write_bytes(b"OLD-COMPLETE")

    class _FailingModel:
        def save(self, path):
            Path(path).write_bytes(b"PARTIAL")
            raise RuntimeError("磁盘写满")

    with pytest.raises(RuntimeError):
        module._save_model_atomic(_FailingModel(), target)
    assert target.read_bytes() == b"OLD-COMPLETE"
    leftovers = [p.name for p in tmp_path.iterdir() if p.name != target.name]
    assert leftovers and all(name.startswith(".tmp-") for name in leftovers)


def test_write_json_atomic_replaces_in_one_step(tmp_path):
    module = _train_module()
    path = tmp_path / "resume.json"
    module._write_json_atomic(path, {"model": "a.zip", "timesteps": 1})
    assert json.loads(path.read_text(encoding="utf-8"))["timesteps"] == 1
    module._write_json_atomic(path, {"model": "b.zip", "timesteps": 2})
    assert json.loads(path.read_text(encoding="utf-8"))["model"] == "b.zip"
    assert not list(tmp_path.glob(".tmp-*"))


def test_prune_checkpoints_never_deletes_recovery_point(tmp_path):
    """滚动恢复点与临时文件不参与历史清理。"""
    module = _train_module()
    directory = tmp_path / "checkpoints"
    directory.mkdir()
    latest = directory / "q3_ckpt_latest.zip"
    latest.write_bytes(b"RECOVERY")
    (directory / ".tmp-q3_ckpt_000000009.zip").write_bytes(b"PARTIAL")
    for index in range(5):
        (directory / f"q3_ckpt_{index:09d}.zip").write_bytes(b"x")
    module._prune_checkpoints(directory, "q3_ckpt", 1)
    assert latest.exists() and latest.read_bytes() == b"RECOVERY"
    assert (directory / ".tmp-q3_ckpt_000000009.zip").exists()


# --------------------------------------------------------------- 集成回归
def test_recovery_point_exists_before_training_and_is_loadable(tmp_path):
    """起步基线：训练刚起就有可加载、可续训的产物（--checkpoint-every 0 也成立）。"""
    _require_rl()
    from sb3_contrib import MaskablePPO

    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "0")
    try:
        assert run.wait_for(lambda: run.recovery().exists(), 120.0), \
            "起步恢复点未落盘"
        loaded = MaskablePPO.load(str(run.recovery()))
        assert loaded.num_timesteps >= 0
        hint = run.resume_hint()
        assert Path(hint["model"]) == run.recovery()
        assert hint["timesteps_semantics"] == "additional"
    finally:
        run.stop()


def test_hard_kill_still_leaves_resumable_recovery_point(tmp_path):
    """急停（强杀、无收尾）后仍能续训：滚动恢复点 + resume.json 都在且可加载。"""
    _require_rl()
    from sb3_contrib import MaskablePPO

    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "0",
               "--recovery-every", "32")
    try:
        assert run.wait_for(lambda: run.recovery().exists(), 120.0), \
            "起步恢复点未落盘"
        first = run.recovery().stat().st_mtime_ns
        assert run.wait_for(
            lambda: run.recovery().stat().st_mtime_ns != first, 300.0), \
            "滚动恢复点未按 --recovery-every 刷新"
        assert run.proc.poll() is None, "训练已自行结束，无法验证急停"
    finally:
        run.hard_kill()          # 急停：强杀整棵进程树，没有任何 Python 收尾

    hint = run.resume_hint()
    assert Path(hint["model"]).is_file(), hint
    assert Path(hint["model"]) == run.recovery(), hint
    loaded = MaskablePPO.load(str(run.recovery()))
    assert loaded.num_timesteps >= 32, loaded.num_timesteps
    metadata = json.loads(run.recovery().with_suffix(".json").read_text(
        encoding="utf-8"))
    assert metadata["recovery"] is True
    assert metadata["checkpoint_reason"] in ("baseline", "interval")
    # 急停没有 graceful 收尾，所以不会有"stopped_by_signal"的最终模型
    assert not (run.run_dir / "q3.zip").exists()


def test_checkpoint_is_written_and_loadable(tmp_path):
    _require_rl()
    from sb3_contrib import MaskablePPO

    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
    try:
        assert run.wait_for(run.checkpoints, 300.0), \
            "训练未在 300s 内产出 checkpoint"
    finally:
        run.stop()

    archives = run.checkpoints()
    assert archives, "checkpoint 目录为空"
    loaded = MaskablePPO.load(str(archives[-1]))
    assert loaded.num_timesteps >= 64
    metadata = json.loads(archives[-1].with_suffix(".json").read_text(
        encoding="utf-8"))
    assert metadata["checkpoint"] is True
    assert metadata["checkpoint_timesteps"] == loaded.num_timesteps
    assert metadata["model_sha256"]
    hyper = metadata["ppo_hyperparams"]
    assert (hyper["n_steps"], hyper["batch_size"], hyper["n_epochs"]) == (32, 32, 1)


def test_control_file_stop_saves_model(tmp_path):
    """控制文件 stop 是训练前端的默认停止路径，Windows 上同样必须成立。"""
    _require_rl()
    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
    model_path = run.run_dir / "q3.zip"
    try:
        assert run.wait_for(run.checkpoints, 300.0), \
            "训练未在 300s 内进入 rollout 阶段"
        assert run.proc.poll() is None, "训练进程提前退出，无法验证优雅停止"
    finally:
        output = run.stop()

    assert run.proc.returncode == 0, \
        f"控制文件停止应正常退出，实际 {run.proc.returncode}\n{output}"
    assert model_path.exists(), "控制文件停止后最终模型未落盘"
    metadata = json.loads(model_path.with_suffix(".json").read_text(
        encoding="utf-8"))
    assert metadata["stopped_by_signal"] is True
    assert metadata["stop_signal"] == "control_file"
    assert metadata["checkpoint_files"], "元数据应记录已产出的 checkpoint"
    assert not (run.run_dir / "training_failure.json").exists()
    assert "saved_after_stop" in output


def test_control_file_pause_then_resume_keeps_training_alive(tmp_path):
    """暂停是协作式的：进程存活、停在 step 边界，continue 后照常收尾。"""
    _require_rl()
    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
    try:
        assert run.wait_for(run.checkpoints, 300.0), \
            "训练未在 300s 内进入 rollout 阶段"
        run.pause()
        assert run.wait_for(lambda: '"event": "paused"' in run.logs(), 120.0), \
            "训练脚本未在 120s 内进入暂停"
        assert run.proc.poll() is None, "暂停不应结束进程"
        run.resume()
        assert run.wait_for(lambda: '"event": "resumed"' in run.logs(), 120.0), \
            "训练脚本未在 120s 内恢复"
    finally:
        run.stop()
    assert run.proc.returncode == 0


@pytest.mark.skipif(os.name == "nt", reason="Windows 没有 os.killpg，信号兜底路径不适用")
def test_sigterm_saves_model_instead_of_losing_the_run(tmp_path):
    _require_rl()
    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
    model_path = run.run_dir / "q3.zip"
    try:
        assert run.wait_for(run.checkpoints, 300.0), \
            "训练未在 300s 内进入 rollout 阶段"
        assert run.proc.poll() is None, "训练进程提前退出，无法验证优雅停止"
    finally:
        output = run.stop_by_signal()

    assert run.proc.returncode == 0, \
        f"优雅停止应正常退出，实际 {run.proc.returncode}\n{output}"
    assert model_path.exists(), "SIGTERM 后最终模型未落盘"
    metadata = json.loads(model_path.with_suffix(".json").read_text(
        encoding="utf-8"))
    assert metadata["stopped_by_signal"] is True
    assert metadata["stop_signal"] == "SIGTERM"
    assert metadata["checkpoint_files"], "元数据应记录已产出的 checkpoint"
    assert not (run.run_dir / "training_failure.json").exists()
    assert "saved_after_stop" in output