"""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


@pytest.mark.gpu
@pytest.mark.parametrize("offload_video_to_cpu", [False, True])
def test_jpg_loader_matches_medsam2_loader_bitwise(
    tmp_path: Path, offload_video_to_cpu: bool
) -> None:
    """换上的读帧函数与 MedSAM2 原来的 `load_video_frames` 交出逐位相同的张量和原始尺寸。"""
    import torch
    from sam2.utils.misc import load_video_frames  # pyright: ignore[reportUnknownVariableType]

    from cadx.segment._medsam2 import (
        _load_jpg_frames,  # pyright: ignore[reportPrivateUsage]
        _write_jpgs,  # pyright: ignore[reportPrivateUsage]
    )

    frames, _ = _moving_disc_video(12, 120, 200, radius=18)
    noise = np.random.default_rng(0).integers(0, 40, size=(12, 120, 200, 3), dtype=np.uint8)
    _write_jpgs([frame + n for frame, n in zip(frames, noise, strict=True)], tmp_path)
    kwargs: dict[str, Any] = {
        "video_path": str(tmp_path),
        "image_size": 512,
        "offload_video_to_cpu": offload_video_to_cpu,
        "async_loading_frames": False,
        "compute_device": torch.device("cuda"),
    }
    expected, expected_h, expected_w = load_video_frames(**kwargs)  # pyright: ignore[reportUnknownVariableType]
    assert isinstance(expected, torch.Tensor)  # 不异步读时原函数交出整段张量
    images, height, width = _load_jpg_frames(**kwargs)
    assert (height, width) == (expected_h, expected_w) == (120, 200)
    assert images.device == expected.device
    assert images.dtype == torch.float32
    assert torch.equal(images, expected)