"""segment：分割设置的校验，以及在合成小视频上真跑一次 MedSAM2（`gpu`）。"""

from pathlib import Path
from typing import Any, cast

import cv2
import numpy as np
import numpy.typing as npt
import pytest
import yaml
from pydantic import ValidationError

from cadx.segment import SegmentConfig, medsam2_commit, segment_video

REPO_ROOT = Path(__file__).resolve().parents[1]
CEUS_YAML = REPO_ROOT / "config" / "ceus.yaml"


def _segment_section() -> dict[str, Any]:
    document = cast(dict[str, Any], yaml.safe_load(CEUS_YAML.read_text(encoding="utf-8")))
    return cast(dict[str, Any], document["segment"])


def _absolute(section: dict[str, Any]) -> dict[str, Any]:
    """路径字段换成仓库根目录下的绝对路径（脚本也这样做）。"""
    return {
        **section,
        "medsam2_dir": str(REPO_ROOT / section["medsam2_dir"]),
        "checkpoint": str(REPO_ROOT / section["checkpoint"]),
    }


def test_config_yaml_section_validates() -> None:
    cfg = SegmentConfig.model_validate(_segment_section())
    assert cfg.sam2_config == "configs/sam2.1_hiera_t512.yaml"


@pytest.mark.parametrize(
    "field",
    [
        "medsam2_dir",
        "checkpoint",
        "sam2_config",
        "device",
        "offload_video_to_cpu",
        "offload_state_to_cpu",
    ],
)
def test_config_missing_field_raises(field: str) -> None:
    section = _segment_section()
    del section[field]
    with pytest.raises(ValidationError, match=field):
        SegmentConfig.model_validate(section)


def test_config_extra_field_raises() -> None:
    with pytest.raises(ValidationError, match="object_id"):
        SegmentConfig.model_validate({**_segment_section(), "object_id": 1})


def test_config_bool_not_coerced() -> None:
    with pytest.raises(ValidationError, match="offload_video_to_cpu"):
        SegmentConfig.model_validate({**_segment_section(), "offload_video_to_cpu": "false"})


def _moving_disc_video(
    frame_count: int, height: int, width: int, radius: int
) -> tuple[list[npt.NDArray[np.uint8]], list[npt.NDArray[np.bool_]]]:
    """暗背景上一个亮圆从左往右匀速移动；交出逐帧 BGR 图像和每帧圆的真值掩码。"""
    frames: list[npt.NDArray[np.uint8]] = []
    truths: list[npt.NDArray[np.bool_]] = []
    yy, xx = np.mgrid[0:height, 0:width]
    for index in range(frame_count):
        cx = radius + 10 + index * 2
        cy = height // 2
        frame = np.full((height, width, 3), 30, dtype=np.uint8)
        cv2.circle(frame, (cx, cy), radius, (220, 220, 220), thickness=-1)
        frames.append(frame)
        truths.append((xx - cx) ** 2 + (yy - cy) ** 2 <= radius**2)
    return frames, truths


@pytest.mark.gpu
def test_segment_video_on_synthetic_clip() -> None:
    cfg = SegmentConfig.model_validate(_absolute(_segment_section()))
    frame_count, height, width = 40, 120, 200
    frames, truths = _moving_disc_video(frame_count, height, width, radius=18)
    # 两个提示帧，最早的不在第 0 帧，前向和反向传播都要走到。
    prompt_masks = {5: truths[5], 25: truths[25]}

    masks = segment_video(frames, prompt_masks, cfg)

    assert masks.shape == (frame_count, height, width)
    assert masks.dtype == np.bool_
    assert len(medsam2_commit(cfg.medsam2_dir)) == 40