"""セグメンテーション用のスクリプト（ダミーの ONNX・合成データ）のテスト.

仕様（実装サマリーを仕様として扱う）:

- `make_dummy_seg_models.build_resunet()`: 入力 `(1, 3, 256, 256)` float32 → 出力
  `(1, 2, 256, 256)` float32（本物の ResUNet と入出力の形を合わせる）
- `make_dummy_seg_models.build_detector()`: 入力 `image_tensor:0`（`(1, H, W, 3)` uint8）→
  `detection_boxes:0`（`(1, N, 4)`）・`detection_scores:0`（`(1, N)`）
- `make_synthetic_segmentation.generate()`: 実データと同じ構成（`train`・`test/cropped`・
  `test_extra/cropped`・`test/fullsize` の `sample/`・`label/`）．同じ seed で同じ出力．
  出力先が空でなければ `--overwrite` が必要で，目印の無い出力先は削除しない
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import pytest
from PIL import Image

SET_DIRS = ["train", "test/cropped", "test_extra/cropped", "test/fullsize"]


def _session(model: onnx.ModelProto) -> ort.InferenceSession:
    return ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])


def test_dummy_resunet_io_shapes(script_loader) -> None:
    script = script_loader("make_dummy_seg_models")
    sess = _session(script.build_resunet())
    (inp,) = sess.get_inputs()
    assert inp.shape == [1, 3, 256, 256]
    x = np.random.default_rng(0).random((1, 3, 256, 256), dtype=np.float32)
    (out,) = sess.run(None, {inp.name: x})
    assert out.shape == (1, 2, 256, 256)
    assert out.dtype == np.float32
    np.testing.assert_allclose(out[0, 1], -out[0, 0], atol=1e-5)


def test_dummy_detector_io_shapes(script_loader) -> None:
    script = script_loader("make_dummy_seg_models")
    sess = _session(script.build_detector())
    assert [i.name for i in sess.get_inputs()] == ["image_tensor:0"]
    names = [o.name for o in sess.get_outputs()]
    assert {"detection_boxes:0", "detection_scores:0"} <= set(names)
    img = np.zeros((1, 50, 70, 3), np.uint8)
    img[0, 10:20, 30:40] = (200, 90, 100)  # RGB の赤
    boxes, scores = sess.run(["detection_boxes:0", "detection_scores:0"], {"image_tensor:0": img})
    assert boxes.shape == (1, 2, 4)
    assert scores.shape == (1, 2)
    np.testing.assert_allclose(boxes[0, 0], [10 / 50, 30 / 70, 20 / 50, 40 / 70], rtol=1e-6)


def test_main_writes_both_models(script_loader, tmp_path: Path) -> None:
    script = script_loader("make_dummy_seg_models")
    assert script.main(["--out", str(tmp_path / "m")]) == 0
    assert (tmp_path / "m" / "ResUNet.onnx").is_file()
    assert (tmp_path / "m" / "Detection_Normal.onnx").is_file()


@pytest.fixture(scope="module")
def synthetic(script_loader, tmp_path_factory: pytest.TempPathFactory) -> tuple:
    script = script_loader("make_synthetic_segmentation")
    out = tmp_path_factory.mktemp("synseg") / "a"
    n = script.generate(out, 3, seed=5)
    return script, out, n


def test_generate_layout(synthetic: tuple) -> None:
    _, out, n = synthetic
    assert n == 3 * len(SET_DIRS)
    for set_dir in SET_DIRS:
        samples = sorted((out / "input" / set_dir / "sample").iterdir())
        labels = sorted((out / "input" / set_dir / "label").iterdir())
        assert len(samples) == len(labels) == 3
        assert [p.stem for p in samples] == [p.stem for p in labels]
        for s, lb in zip(samples, labels, strict=True):
            with Image.open(s) as si, Image.open(lb) as li:
                assert si.size == li.size
                assert si.mode == "RGB"
                assert li.mode == "L"
                assert (np.asarray(li) >= 128).any()


def test_generate_is_reproducible(synthetic: tuple, tmp_path: Path) -> None:
    script, out, _ = synthetic
    other = tmp_path / "b"
    script.generate(other, 3, seed=5)
    for p in (out / "input").rglob("*.jpg"):
        q = other / "input" / p.relative_to(out / "input")
        assert p.read_bytes() == q.read_bytes()


def test_main_refuses_non_empty_without_overwrite(script_loader, tmp_path: Path) -> None:
    script = script_loader("make_synthetic_segmentation")
    out = tmp_path / "x"
    out.mkdir()
    (out / "keep.txt").write_text("x", encoding="utf-8")
    assert script.main(["--out", str(out), "--per-set", "1"]) == 1
    # 目印が無ければ --overwrite でも消さない
    assert script.main(["--out", str(out), "--per-set", "1", "--overwrite"]) == 1
    assert (out / "keep.txt").is_file()


def test_main_overwrite_regenerates_marked_output(script_loader, tmp_path: Path) -> None:
    script = script_loader("make_synthetic_segmentation")
    out = tmp_path / "y"
    assert script.main(["--out", str(out), "--per-set", "1"]) == 0
    assert script.main(["--out", str(out), "--per-set", "2", "--overwrite"]) == 0
    assert len(list((out / "input" / "train" / "sample").iterdir())) == 2