"""合成データ生成スクリプトのテスト.

仕様: SPEC_01「データ契約」（9 ディレクトリ + manifest.csv）・TEST_01「画像の性質」
（PNG・RGB でアルファなし・マスク外は黒 (0, 0, 0)）・同じ seed で同じ出力．
"""

import importlib.util
import sys
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest
from PIL import Image

from tias_judge_model.data.classes import parse_dir_name
from tias_judge_model.data.manifest import read_manifest, resolve_image_path
from tias_judge_model.data.split import make_kfold_splits, split_train_val_test

SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "make_synthetic_dataset.py"

TEST01_DIRS = {
    "light_red_normal",
    "dark_red_strong",
    "dark_red_mild",
    "purple_mild",
    "red_strong",
    "red_mild",
    "purple_strong",
    "pale_strong",
    "pale_mild",
}


@pytest.fixture(scope="module")
def script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("make_synthetic_dataset", SCRIPT_PATH)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _run(script: ModuleType, out: Path, per_dir: int = 3, seed: int = 0) -> None:
    code = script.main(["--out", str(out), "--per-dir", str(per_dir), "--seed", str(seed)])
    assert code == 0


def _snapshot(root: Path) -> dict[str, bytes]:
    return {
        p.relative_to(root).as_posix(): p.read_bytes()
        for p in sorted(root.rglob("*"))
        if p.is_file() and p.suffix in {".png", ".csv"}
    }


def test_outputs_nine_spec_directories_and_manifest(script: ModuleType, tmp_path: Path) -> None:
    out = tmp_path / "syn"
    _run(script, out, per_dir=3)

    images = out / "images"
    assert {p.name for p in images.iterdir() if p.is_dir()} == TEST01_DIRS
    for d in images.iterdir():
        assert len(list(d.glob("*.png"))) == 3
    assert (out / "manifest.csv").is_file()


def test_manifest_passes_reading_and_validation(script: ModuleType, tmp_path: Path) -> None:
    out = tmp_path / "syn"
    _run(script, out, per_dir=3)

    records = read_manifest(out / "manifest.csv")
    assert len(records) == 27
    assert len({r.patient_id for r in records}) == 27  # 1 患者 1 枚（実データと同じ前提）
    for r in records:
        dir_name = r.image_path.split("/")[0]
        assert (r.label, r.severity) == parse_dir_name(dir_name)
        assert resolve_image_path(out / "images", r).is_file()


def test_same_seed_produces_identical_output(script: ModuleType, tmp_path: Path) -> None:
    _run(script, tmp_path / "a", seed=5)
    _run(script, tmp_path / "b", seed=5)
    a, b = _snapshot(tmp_path / "a"), _snapshot(tmp_path / "b")
    assert a and a == b


def test_different_seed_produces_different_images(script: ModuleType, tmp_path: Path) -> None:
    _run(script, tmp_path / "a", seed=1)
    _run(script, tmp_path / "b", seed=2)
    a, b = _snapshot(tmp_path / "a"), _snapshot(tmp_path / "b")
    assert {k: v for k, v in a.items() if k.endswith(".png")} != {
        k: v for k, v in b.items() if k.endswith(".png")
    }


def test_images_are_rgb_png_with_black_corners(script: ModuleType, tmp_path: Path) -> None:
    out = tmp_path / "syn"
    _run(script, out, per_dir=3)

    pngs = sorted((out / "images").rglob("*.png"))
    assert len(pngs) == 27
    for png in pngs:
        with Image.open(png) as img:
            assert img.format == "PNG"
            assert img.mode == "RGB"  # アルファチャンネルなし（TEST_01）
            pixels = np.asarray(img)  # (H, W, 3)
        h, w, _ = pixels.shape
        for y, x in ((0, 0), (0, w - 1), (h - 1, 0), (h - 1, w - 1)):
            assert tuple(int(v) for v in pixels[y, x]) == (0, 0, 0)
        # 舌の代わりの前景（黒以外）も描かれている
        assert (pixels.sum(axis=2) > 0).any()


def test_synthetic_manifest_can_be_split_by_patient(script: ModuleType, tmp_path: Path) -> None:
    out = tmp_path / "syn"
    _run(script, out, per_dir=5)
    records = read_manifest(out / "manifest.csv")

    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=0)
    assert not set(result.train) & set(result.val)
    assert not set(result.train) & set(result.test)
    assert not set(result.val) & set(result.test)

    folds = make_kfold_splits(records, n_folds=5, seed=0)
    held_out = [p for f in folds for p in f.held_out]
    assert sorted(held_out) == sorted({r.patient_id for r in records})