"""`scripts/ceus_quant.py`：逐例循环、failed.csv、输出布局、分割缓存。

假交付包两例：第 1 例是能过质量门的合成分屏造影视频（造影在右半，绿框是病灶的空心描边），第 2 例
同样的画面但没有绿框。分割换成计数的假函数（monkeypatch 脚本里引用的 `segment_video`），交出病灶
方块挪到左半局部坐标的掩码，不调显卡。

手算的预期（视频 400×200、10 fps、40 秒，推注时间 X = 5 秒，同 test_ceus_pipeline）：
- 第 2 例两半都没有绿框 → 第 2 步失败，进 failed.csv，没有输出目录，也不调分割；
- 第 1 例 F1 在视频第 20 − 5 = 15 秒，有关键帧图；F2 在第 55 秒，超出视频，未定位，没有图。
"""

import csv
import importlib.util
import json
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from types import ModuleType
from typing import Any

import cv2
import numpy as np
import numpy.typing as npt
import pytest
import yaml

from cadx.ceus import ANCHOR_SLOTS, LognormalParameters, evaluate_lognormal
from conftest import CEUS_YAML
from fake_delivery import make_delivery

_SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "ceus_quant.py"
_WIDTH, _HEIGHT, _SPLIT = 400, 200, 200
_FPS, _FRAMES = 10.0, 400
_LESION = (250, 60, 369, 159)  # 整帧坐标 x1, y1, x2, y2（含端点）
_CYST = (290, 95, 320, 125)
_BOLUS = LognormalParameters(area=1300.0, mu=2.4, sigma=0.45, arrival_offset_s=4.0, baseline=40.0)


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("ceus_quant", _SCRIPT)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _write_video(path: Path, green_box: bool) -> None:
    rng = np.random.default_rng(3)
    intensity = evaluate_lognormal(np.arange(_FRAMES) / _FPS, _BOLUS)
    writer = cv2.VideoWriter(str(path), cv2.VideoWriter.fourcc(*"mp4v"), _FPS, (_WIDTH, _HEIGHT))
    x1, y1, x2, y2 = _LESION
    for index in range(_FRAMES):
        frame = np.zeros((_HEIGHT, _WIDTH, 3), dtype=np.uint8)
        frame[:, :_SPLIT] = 120
        frame[:, _SPLIT + 1 :] = 30
        lesion = np.clip(
            intensity[index] + rng.normal(0.0, 6.0, (y2 - y1 + 1, x2 - x1 + 1)), 0, 255
        )
        frame[y1 : y2 + 1, x1 : x2 + 1] = lesion.astype(np.uint8)[:, :, None]
        cx1, cy1, cx2, cy2 = _CYST
        frame[cy1 : cy2 + 1, cx1 : cx2 + 1] = 10
        if green_box:
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), thickness=3)
        writer.write(frame)
    writer.release()


def _config_dict(tmp_path: Path, delivery_dir: Path) -> dict[str, Any]:
    document: dict[str, Any] = yaml.safe_load(CEUS_YAML.read_text(encoding="utf-8"))
    document["paths"] = {
        "delivery_dir": str(delivery_dir),
        "runs_root": str(tmp_path / "runs"),
        "work_root": str(tmp_path / "work"),
    }
    checkpoint = tmp_path / "fake_checkpoint.pt"
    checkpoint.write_bytes(b"fake weights")
    document["segment"].update(
        {"medsam2_dir": str(tmp_path / "MedSAM2"), "checkpoint": str(checkpoint), "device": "cpu"}
    )
    return document


def _write_config(tmp_path: Path, document: dict[str, Any], name: str) -> Path:
    path = tmp_path / name
    path.write_text(yaml.safe_dump(document, allow_unicode=True), encoding="utf-8")
    return path


class _CountingSegment:
    """假分割：每帧都是病灶方块挪到左半局部坐标后的位置（略往里收 3 像素）。"""

    def __init__(self) -> None:
        self.calls = 0

    def __call__(
        self,
        frames: Sequence[npt.NDArray[np.uint8]],
        prompt_masks: Mapping[int, npt.NDArray[np.bool_]],
        cfg: object,
    ) -> npt.NDArray[np.bool_]:
        self.calls += 1
        height, width = frames[0].shape[:2]
        masks = np.zeros((len(frames), height, width), dtype=bool)
        x1, y1, x2, y2 = _LESION
        masks[:, y1 + 3 : y2 - 2, x1 - _SPLIT + 3 : x2 - _SPLIT - 2] = True
        return masks


@pytest.fixture
def setup(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]]:
    """写假交付包（两例真视频），载入脚本并换掉分割、MedSAM2 提交号和运行号。"""
    delivery_dir = make_delivery(2).write(tmp_path / "delivery")
    _write_video(delivery_dir / "CEUS_agent" / "001" / "U001_roi2.mp4", green_box=True)
    _write_video(delivery_dir / "CEUS_agent" / "002" / "U002_roi2.mp4", green_box=False)
    module = _load_script()
    segment = _CountingSegment()
    monkeypatch.setattr(module, "segment_video", segment)
    run_ids: list[str] = []

    def fake_commit(medsam2_dir: Path) -> str:
        return "c" * 40

    def fake_run_id(now: object, git: object) -> str:
        return run_ids[-1]

    monkeypatch.setattr(module, "medsam2_commit", fake_commit)
    monkeypatch.setattr(module, "make_build_id", fake_run_id)
    return module, segment, _config_dict(tmp_path, delivery_dir), run_ids.append


def test_run_writes_outputs_and_failed_case(
    tmp_path: Path,
    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
    capsys: pytest.CaptureFixture[str],
) -> None:
    module, segment, document, next_run_id = setup
    next_run_id("run-a")
    module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])

    run_dir = tmp_path / "runs" / "run-a"
    assert sorted(p.name for p in (tmp_path / "runs").iterdir()) == ["run-a"]
    with (run_dir / "failed.csv").open(encoding="utf-8", newline="") as handle:
        rows = list(csv.reader(handle))
    assert rows[0] == ["video_id", "step", "reason"]
    assert [row[:2] for row in rows[1:]] == [["U002", "2"]]
    assert "都没有" in rows[1][2]
    assert not (run_dir / "U002").exists()
    assert segment.calls == 1

    case_dir = run_dir / "U001"
    for name in ("trajectory.npz", "tic_curves.npz", "summary.json", "tic.png", "overlay.mp4"):
        assert (case_dir / name).stat().st_size > 0, name
    summary = json.loads((case_dir / "summary.json").read_text(encoding="utf-8"))
    assert summary["video_id"] == "U001"
    assert summary["injection_offset_s"] == 5.0
    assert list(summary["anchor_frames"]) == list(ANCHOR_SLOTS)
    assert summary["perfusion_metrics"]["roi_perf"] is not None
    assert "pan_detected_fraction" in summary
    assert summary["anchor_frames"]["F2"]["status"] == "not_covered"
    located = {s for s, a in summary["anchor_frames"].items() if a["frame_index"] is not None}
    assert "F1" in located and "F2" not in located
    keyframes = {p.stem for p in (case_dir / "keyframes").iterdir()}
    assert keyframes == located
    image = cv2.imread(str(case_dir / "keyframes" / "F1.png"))
    assert image is not None and image.shape == (_HEIGHT, _WIDTH, 3)  # 整幅原始画面

    snapshot = run_dir / "snapshot"
    assert (snapshot / "ceus.yaml").read_bytes() == (tmp_path / "ceus.yaml").read_bytes()
    segment_record = json.loads((snapshot / "segment.json").read_text(encoding="utf-8"))
    assert segment_record["medsam2_commit"] == "c" * 40
    assert segment_record["cuda_peak_memory_bytes"] is None  # device 是 cpu
    assert (snapshot / "snapshot.json").is_file()

    out = capsys.readouterr().out
    assert "成功：1" in out and "第 2 步 1" in out
    assert "U001" not in out and "U002" not in out


def test_unreadable_video_fails_at_step_1_and_next_case_runs(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
) -> None:
    module, segment, document, next_run_id = setup
    # 第 1 例换成损坏的视频（cv2 打不开），第 2 例换成正常视频。
    videos = Path(document["paths"]["delivery_dir"]) / "CEUS_agent"
    (videos / "001" / "U001_roi2.mp4").write_bytes(b"not a real video")
    _write_video(videos / "002" / "U002_roi2.mp4", green_box=True)
    next_run_id("run-a")
    module.main(["--config", str(_write_config(tmp_path, document, "ceus.yaml"))])

    run_dir = tmp_path / "runs" / "run-a"
    with (run_dir / "failed.csv").open(encoding="utf-8", newline="") as handle:
        rows = list(csv.reader(handle))
    assert [row[:2] for row in rows[1:]] == [["U001", "1"]]
    assert "打不开" in rows[1][2]
    assert not (run_dir / "U001").exists()
    assert (run_dir / "U002" / "summary.json").is_file()
    assert segment.calls == 1


def test_segment_cache_reused_and_keyed_on_step_1_3_params(
    tmp_path: Path,
    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
) -> None:
    module, segment, document, next_run_id = setup
    config = _write_config(tmp_path, document, "ceus.yaml")
    next_run_id("run-a")
    module.main(["--config", str(config)])
    assert segment.calls == 1
    cache = tmp_path / "work" / "segment_cache"
    assert [p.name.endswith(".partial") for p in cache.iterdir()] == [False]

    next_run_id("run-b")
    module.main(["--config", str(config)])
    assert segment.calls == 1  # 同样的输入，读缓存
    # 同一进程里第二次跑（cv2 的随机数已被第一次推进过）结果也逐位相同：quantify_case 每例重设种子。
    first = json.loads((tmp_path / "runs/run-a/U001/summary.json").read_text(encoding="utf-8"))
    second = json.loads((tmp_path / "runs/run-b/U001/summary.json").read_text(encoding="utf-8"))
    assert first == second

    # 改设备不改结果，不重算；改第 1 步的参数，键变了，重新分割。
    document["segment"]["device"] = "cuda:0"
    document["segment"]["offload_state_to_cpu"] = True
    next_run_id("run-c")
    module.main(["--config", str(_write_config(tmp_path, document, "device.yaml"))])
    assert segment.calls == 1
    document["ceus"]["views"]["split_sample_frames"] = 20
    next_run_id("run-d")
    module.main(["--config", str(_write_config(tmp_path, document, "views.yaml"))])
    assert segment.calls == 2
    assert len(list(cache.iterdir())) == 2


def test_existing_run_directory_is_not_overwritten(
    tmp_path: Path,
    setup: tuple[ModuleType, _CountingSegment, dict[str, Any], Callable[[str], None]],
) -> None:
    module, _, document, next_run_id = setup
    config = _write_config(tmp_path, document, "ceus.yaml")
    next_run_id("run-a")
    module.main(["--config", str(config)])
    before = sorted(str(p) for p in (tmp_path / "runs").rglob("*"))
    with pytest.raises(FileExistsError, match="运行目录已存在"):
        module.main(["--config", str(config)])
    assert sorted(str(p) for p in (tmp_path / "runs").rglob("*")) == before


def _key(module: ModuleType, document: dict[str, Any]) -> str:
    config = module.QuantConfig.model_validate(document)
    return module.segment_cache_key("v" * 64, config, "w" * 64, "c" * 40)


@pytest.mark.parametrize(
    ("section", "path", "value", "changes"),
    [
        ("ceus", ("views", "split_search_band"), [0.3, 0.7], True),
        ("ceus", ("prompt_source", "min_area"), 9000, True),
        ("ceus", ("keyframes", "iou_threshold"), 0.9, True),
        ("segment", ("sam2_config",), "configs/other.yaml", True),
        ("segment", ("checkpoint",), "models/other.pt", True),
        ("segment", ("device",), "cuda:1", False),
        ("segment", ("offload_video_to_cpu",), True, False),
        ("ceus", ("smoothing", "temporal_window"), 9, False),  # 第 4 步
        ("ceus", ("tic_analysis", "fitting", "minimum_r2"), 0.5, False),  # 第 7 步
    ],
)
def test_cache_key_covers_exactly_step_1_3_params(
    tmp_path: Path,
    section: str,
    path: tuple[str, ...],
    value: object,
    changes: bool,
) -> None:
    module = _load_script()
    document = _config_dict(tmp_path, tmp_path)
    base = _key(module, document)
    node = document[section]
    for name in path[:-1]:
        node = node[name]
    node[path[-1]] = value
    assert (_key(module, document) != base) is changes
    assert _key(module, document) == _key(module, document)  # 同输入同键
    other = module.QuantConfig.model_validate(document)
    assert module.segment_cache_key("x" * 64, other, "w" * 64, "c" * 40) != _key(module, document)