"""CEUS 定量：对交付包里检查全部通过的病例逐例跑第 1–9 步，写进 `runs/ceus/<运行号>/`。

用法：uv run python scripts/ceus_quant.py --config config/ceus.yaml
第 1 步算出两半配准后，按视频 SHA-256 查配准目视核对的结果表（`paths.align_review`，由
`scripts/ceus_align_review.py score` 写出）：核对没通过的病例第 1 步失败；表里没有这段视频，或配准
结果与表里记的不同，说明核对的不是这次跑的东西，整次运行停下。
第 3 步（MedSAM2 分割）的结果按内容键缓存在 `<work_root>/segment_cache/<键>/masks.npz`，键相同就
直接读缓存、不调显卡。某一例在第 1–6 步失败（`CeusCaseError`）记进 `failed.csv`，接着跑下一例；
其余异常整次运行停下，留下 `.partial` 目录。终端只打印计数，不打印视频号。
"""

import argparse
import csv
import hashlib
import json
import sys
import time
from collections import Counter
from collections.abc import Callable, Hashable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Annotated, Any

import cv2
import numpy as np
import numpy.typing as npt
import yaml
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, field_validator

from cadx.ceus import (
    ANCHOR_SLOTS,
    CeusCaseError,
    CeusConfig,
    SegmentationInput,
    keyframe_images,
    pack_tic_curves,
    pack_trajectory,
    prepare_segmentation,
    quant_summary,
    quantify_case,
    render_case_overlay,
    render_tic_png,
    tic_plot_markers,
)
from cadx.delivery import CaseRecord, read_delivery
from cadx.runs import make_build_id, read_git_state, write_snapshot
from cadx.segment import SegmentConfig, cuda_peak_memory_bytes, medsam2_commit, segment_video

REPO_ROOT = Path(__file__).resolve().parents[1]

# 第 1–2 步（`prepare_segmentation`）读的 ceus 段子段：views（分屏）、registration（两半配准）、
# prompt_source（绿框颜色，配准和选提示帧都用）、pan_source（紫圈颜色，配准时修补掉）、keyframes
# （提示帧选择）。它们决定交给分割的帧和提示掩码，进分割缓存键。
_STEP_1_2_CEUS_FIELDS = frozenset(
    {"views", "registration", "prompt_source", "pan_source", "keyframes"}
)
# segment 段里不改变分割结果的开关，不进缓存键；其余字段都进。
_RESULT_NEUTRAL_SEGMENT_FIELDS = frozenset(
    {"device", "offload_video_to_cpu", "offload_state_to_cpu"}
)
_FAILED_HEADER = ("video_id", "step", "reason")
# 配准目视核对结果表 review.csv 的列：(dx, dy) 是左半 → 右半的平移
# （`HalfRegistration.correlation`），passed 取 true / false，后两列是两位判读者的判定。
REVIEW_HEADER = ("video_sha256", "dx", "dy", "passed", "verdict_a", "verdict_b")


def _reject_empty_path(value: object) -> object:
    # Path("") 会变成 "."（当前目录），空串在转换前就拦下。
    if value == "":
        raise ValueError("路径不能是空串")
    return value


# YAML 里的路径是字符串；路径字段是唯一允许从字符串转换类型的地方。
type _PathField = Annotated[Path, Field(strict=False), BeforeValidator(_reject_empty_path)]


class PathsConfig(BaseModel):
    """`paths` 段：交付包、运行目录的上级、中间文件目录、配准目视核对的运行目录上级和结果表。

    相对路径相对仓库根目录。`align_review` 在 ceus_quant 运行时必须存在，出图时用不到。
    """

    model_config = ConfigDict(strict=True, extra="forbid", frozen=True)

    delivery_dir: _PathField
    runs_root: _PathField
    work_root: _PathField
    align_runs_root: _PathField
    align_review: _PathField


class AlignCheckConfig(BaseModel):
    """`align_check` 段：配准目视核对出图用的随机数种子、阴性对照例数、对照错开距离区间（像素）。"""

    model_config = ConfigDict(strict=True, extra="forbid", frozen=True)

    seed: int
    control_count: int = Field(ge=1)
    control_shift_px: list[int]

    @field_validator("control_shift_px")
    @classmethod
    def _shift_range(cls, value: list[int]) -> list[int]:
        if len(value) != 2 or not 0 < value[0] <= value[1]:
            raise ValueError(
                f"align_check.control_shift_px 必须是 [lo, hi] 且 0 < lo ≤ hi，收到 {value}"
            )
        return value


class QuantConfig(BaseModel):
    """`ceus.yaml` 整份：`segment` 段交 segment 校验，`ceus` 段交 ceus 校验。"""

    model_config = ConfigDict(strict=True, extra="forbid", frozen=True)

    paths: PathsConfig
    align_check: AlignCheckConfig
    segment: SegmentConfig
    ceus: CeusConfig


class _UniqueKeyLoader(yaml.SafeLoader):
    """同一层出现重复键就报错：PyYAML 默认让后一个悄悄覆盖前一个。"""

    # 签名（含 deep 的默认值）照抄 PyYAML 的基类方法，覆盖时必须一致。
    def construct_mapping(self, node: yaml.MappingNode, deep: bool = False) -> dict[Hashable, Any]:
        seen: set[Hashable] = set()
        for key_node, _ in node.value:
            key = self.construct_object(key_node, deep=deep)  # pyright: ignore[reportUnknownMemberType]
            if key in seen:
                raise ValueError(f"YAML 第 {key_node.start_mark.line + 1} 行重复的键：{key!r}")
            seen.add(key)
        return super().construct_mapping(node, deep=deep)


def load_config(path: Path) -> QuantConfig:
    """读并校验 `ceus.yaml`，路径保持文件里的写法（缓存键用这一份，不随仓库位置变）。"""
    text = path.read_text(encoding="utf-8")
    return QuantConfig.model_validate(yaml.load(text, Loader=_UniqueKeyLoader))  # noqa: S506  继承 SafeLoader


def resolve_paths(config: QuantConfig) -> QuantConfig:
    """相对路径换成仓库根目录下的绝对路径（绝对路径原样保留）。"""
    paths, segment = config.paths, config.segment
    return config.model_copy(
        update={
            "paths": paths.model_copy(
                update={
                    "delivery_dir": REPO_ROOT / paths.delivery_dir,
                    "runs_root": REPO_ROOT / paths.runs_root,
                    "work_root": REPO_ROOT / paths.work_root,
                    "align_runs_root": REPO_ROOT / paths.align_runs_root,
                    "align_review": REPO_ROOT / paths.align_review,
                }
            ),
            "segment": segment.model_copy(
                update={
                    "medsam2_dir": REPO_ROOT / segment.medsam2_dir,
                    "checkpoint": REPO_ROOT / segment.checkpoint,
                }
            ),
        }
    )


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(1 << 20):
            digest.update(chunk)
    return digest.hexdigest()


@dataclass(frozen=True)
class ReviewEntry:
    """review.csv 的一行：核对时的配准结果（左半 → 右半）和是否通过。"""

    offset: tuple[int, int]
    passed: bool


def read_review(path: Path) -> dict[str, ReviewEntry]:
    """读配准目视核对结果表，视频 SHA-256 → 这一行。表头、取值不合规或同一视频出现两次就报错。"""
    with path.open(encoding="utf-8", newline="") as handle:
        rows = list(csv.reader(handle))
    if not rows or tuple(rows[0]) != REVIEW_HEADER:
        raise ValueError(f"{path} 的表头应为 {REVIEW_HEADER}")
    table: dict[str, ReviewEntry] = {}
    for number, row in enumerate(rows[1:], start=2):
        if len(row) != len(REVIEW_HEADER) or row[3] not in ("true", "false"):
            raise ValueError(f"{path} 第 {number} 行不合规：{row}")
        if row[0] in table:
            raise ValueError(f"{path} 第 {number} 行：同一段视频出现两次")
        table[row[0]] = ReviewEntry(offset=(int(row[1]), int(row[2])), passed=row[3] == "true")
    return table


def check_review(
    prepared: SegmentationInput, video_sha256: str, review: Mapping[str, ReviewEntry]
) -> None:
    """第 1 步配准结果对照目视核对表：没通过抛 `CeusCaseError`（第 1 步）；对不上表就报错停下。"""
    if video_sha256 not in review:
        raise RuntimeError(
            f"配准目视核对表里没有这段视频（SHA-256 {video_sha256}），核对的不是这批视频"
        )
    entry = review[video_sha256]
    offset = prepared.registration.correlation
    if offset != entry.offset:
        raise RuntimeError(
            f"这段视频（SHA-256 {video_sha256}）的配准结果 {offset} 与目视核对表里的 "
            f"{entry.offset} 不同，核对的不是这次跑的东西"
        )
    if not entry.passed:
        raise CeusCaseError(1, "配准目视核对未通过")


def segment_cache_key(
    video_sha256: str, config: QuantConfig, checkpoint_sha256: str, commit: str
) -> str:
    """分割缓存键：视频、第 1–3 步参数、权重、MedSAM2 提交号拼成规范 JSON 后的 SHA-256。"""
    payload = {
        "video_sha256": video_sha256,
        "ceus": config.ceus.model_dump(mode="json", include=set(_STEP_1_2_CEUS_FIELDS)),
        "segment": config.segment.model_dump(
            mode="json", exclude=set(_RESULT_NEUTRAL_SEGMENT_FIELDS)
        ),
        "checkpoint_sha256": checkpoint_sha256,
        "medsam2_commit": commit,
    }
    canonical = json.dumps(payload, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _masks(
    prepared: SegmentationInput, key: str, cache_root: Path, segment_cfg: SegmentConfig
) -> tuple[npt.NDArray[np.bool_], bool]:
    """读分割缓存；没有就调 MedSAM2，先写 `<键>.partial` 再改名。返回（掩码，是否命中缓存）。"""
    final = cache_root / key
    if final.exists():
        with np.load(final / "masks.npz", allow_pickle=False) as data:
            return data["masks"], True
    masks = segment_video(prepared.target_frames, prepared.prompt_masks, segment_cfg)
    partial = cache_root / f"{key}.partial"
    # 上次中途停下留下的 .partial 会让这里报错：work/ 可以删掉重算，删掉它再跑。
    partial.mkdir(parents=True, exist_ok=False)
    np.savez_compressed(partial / "masks.npz", masks=masks)
    partial.rename(final)
    return masks, False


def _write_png(path: Path, image: npt.NDArray[np.uint8]) -> None:
    ok, encoded = cv2.imencode(".png", image)
    if not ok:
        raise RuntimeError(f"关键帧编码 PNG 失败：{path}")
    path.write_bytes(encoded.tobytes())


def _run_case(
    case: CaseRecord,
    cfg: QuantConfig,
    key_of: Callable[[str], str],
    review: Mapping[str, ReviewEntry],
    run_dir: Path,
) -> tuple[tuple[str, ...], bool]:
    """跑一例并写出产物，返回（定位到的关键帧槽位，分割是否命中缓存）。

    `key_of` 把视频的 SHA-256 换成分割缓存键（其余几样跨病例不变）；`review` 是配准目视核对表。
    第 1–2 步都在 `prepare_segmentation` 里，所以第 2 步失败的病例不查核对表。
    """
    prepared = prepare_segmentation(case.video, cfg.ceus)
    video_sha256 = file_sha256(case.video)
    check_review(prepared, video_sha256, review)
    key = key_of(video_sha256)
    masks, hit = _masks(prepared, key, cfg.paths.work_root / "segment_cache", cfg.segment)
    result = quantify_case(prepared, masks, float(case.bolus_start_s), cfg.ceus)

    case_dir = run_dir / case.video_id
    (case_dir / "keyframes").mkdir(parents=True, exist_ok=False)
    (case_dir / "trajectory.npz").write_bytes(pack_trajectory(result.trajectory))
    (case_dir / "tic_curves.npz").write_bytes(pack_tic_curves(result.curves))
    registration = prepared.registration
    summary = {
        "video_id": case.video_id,
        "segment_cache_key": key,
        "registration": {
            "contrast_side": prepared.contrast.side,
            "contrast_to_target_offset": list(prepared.offset),
            "left_to_right_correlation": list(registration.correlation),
            "left_to_right_template": list(registration.template),
            "disagreement_px": registration.disagreement_px,
        },
        **quant_summary(result),
    }
    (case_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8"
    )
    markers = tic_plot_markers(result.anchors)
    (case_dir / "tic.png").write_bytes(render_tic_png(result.tic, markers=markers))
    (case_dir / "overlay.mp4").write_bytes(render_case_overlay(prepared, result))
    images = keyframe_images(prepared, result.anchors)
    for slot, image in images.items():
        _write_png(case_dir / "keyframes" / f"{slot}.png", image)
    return tuple(images), hit


def main(argv: Sequence[str]) -> None:
    parser = argparse.ArgumentParser(description="CEUS 定量：逐例跑第 1–9 步")
    parser.add_argument("--config", type=Path, required=True, help="ceus.yaml 的路径")
    config_path: Path = parser.parse_args(argv).config
    started_at = datetime.now().astimezone()
    clock = time.monotonic()

    written = load_config(config_path)
    cfg = resolve_paths(written)
    git = read_git_state(REPO_ROOT)
    run_id = make_build_id(started_at, git)
    run_dir = cfg.paths.runs_root / run_id
    partial = cfg.paths.runs_root / f"{run_id}.partial"
    for path in (run_dir, partial):
        if path.exists():
            raise FileExistsError(f"运行目录已存在，不覆盖：{path}")
    partial.mkdir(parents=True, exist_ok=False)

    review = read_review(cfg.paths.align_review)
    delivery = read_delivery(cfg.paths.delivery_dir, None)
    checkpoint_sha256 = file_sha256(cfg.segment.checkpoint)
    commit = medsam2_commit(cfg.segment.medsam2_dir)

    # 缓存键用配置文件里的写法（路径不随仓库位置变）。
    def key_of(video_sha256: str) -> str:
        return segment_cache_key(video_sha256, written, checkpoint_sha256, commit)

    failures: list[tuple[str, int, str]] = []
    located: Counter[str] = Counter()
    cache_hits = 0
    total = len(delivery.cases)
    for number, case in enumerate(delivery.cases, start=1):
        try:
            slots, hit = _run_case(case, cfg, key_of, review, partial)
        except CeusCaseError as error:
            failures.append((case.video_id, error.step, error.reason))
            print(f"[{number}/{total}] 失败在第 {error.step} 步", flush=True)
            continue
        located.update(slots)
        cache_hits += hit
        print(f"[{number}/{total}] 完成", flush=True)

    with (partial / "failed.csv").open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(_FAILED_HEADER)
        writer.writerows(failures)
    peak = cuda_peak_memory_bytes(cfg.segment.device)
    write_snapshot(
        snapshot_dir=partial / "snapshot",
        build_id=run_id,
        git=git,
        started_at=started_at,
        finished_at=datetime.now().astimezone(),
        uv_lock=REPO_ROOT / "uv.lock",
        configs={"ceus.yaml": config_path, "align_review.csv": cfg.paths.align_review},
    )
    segment_record = {
        "checkpoint_sha256": checkpoint_sha256,
        "medsam2_commit": commit,
        "device": cfg.segment.device,
        "cuda_peak_memory_bytes": peak,
        "segment_cache_hits": cache_hits,
    }
    (partial / "snapshot" / "segment.json").write_text(
        json.dumps(segment_record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    # POSIX 的 rename 会悄悄覆盖空目录，所以改名前再查一次。
    if run_dir.exists():
        raise FileExistsError(f"运行目录已存在，不覆盖：{run_dir}")
    partial.rename(run_dir)

    by_step = Counter(step for _, step, _ in failures)
    print(f"运行目录：{run_dir}")
    print(f"病例总数（检查全部通过）：{total}")
    print(f"成功：{total - len(failures)}")
    failed_text = "，".join(f"第 {step} 步 {n}" for step, n in sorted(by_step.items()))
    print(f"按步分的失败数：{failed_text or '无'}")
    print("关键帧按槽位分的定位数：" + "，".join(f"{s} {located[s]}" for s in ANCHOR_SLOTS))
    print(f"分割缓存命中：{cache_hits}")
    print(f"耗时：{time.monotonic() - clock:.1f} 秒")
    peak_text = "不适用（非 CUDA 设备）" if peak is None else f"{peak / 2**20:.0f} MiB"
    print(f"显存峰值（torch 分配）：{peak_text}")


if __name__ == "__main__":
    main(sys.argv[1:])