"""Tests for scripts/prepare_test_data.py — split ratios, reproducibility, CSV format, pairs."""
import csv
from pathlib import Path

from scripts.make_sample_images import make_image
from scripts.prepare_test_data import prepare


def _make_dataset_dir(tmp: Path, layout: dict[str, int]) -> Path:
    """layout: {folder_name: image_count}. Writes real JPEGs."""
    root = tmp / "input"
    for folder, count in layout.items():
        d = root / folder
        d.mkdir(parents=True)
        for i in range(count):
            (d / f"img{i}.jpg").write_bytes(make_image(hash((folder, i)) % 1000))
    return root


def _read(path: Path) -> list[dict]:
    with path.open(newline="", encoding="utf-8") as fh:
        return list(csv.DictReader(fh))


def test_default_holdout_split(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogA": 3, "dogB": 6})
    out = tmp_path / "out"
    prepare(root, out, holdout=None, seed=1)

    known = _read(out / "known_dogs.csv")
    found = _read(out / "found_dogs.csv")
    by_folder = lambda rows, f: [r for r in rows if r["folder"] == f]  # noqa: E731

    # dogA has <5 images -> 1 holdout; dogB has 5+ -> 2 holdout.
    assert len(by_folder(found, "dogA")) == 1
    assert len(by_folder(known, "dogA")) == 2
    assert len(by_folder(found, "dogB")) == 2
    assert len(by_folder(known, "dogB")) == 4


def test_explicit_holdout(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogB": 6})
    out = tmp_path / "out"
    prepare(root, out, holdout=1, seed=1)
    assert len(_read(out / "found_dogs.csv")) == 1
    assert len(_read(out / "known_dogs.csv")) == 5


def test_reproducible_with_same_seed(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogA": 4, "dogB": 5})
    out1, out2 = tmp_path / "o1", tmp_path / "o2"
    prepare(root, out1, holdout=None, seed=7)
    prepare(root, out2, holdout=None, seed=7)
    for fname in ("known_dogs.csv", "found_dogs.csv", "pairs.csv"):
        assert (out1 / fname).read_text() == (out2 / fname).read_text()


def test_csv_columns_match_loader_contract(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogA": 2})
    out = tmp_path / "out"
    prepare(root, out, holdout=1, seed=1)
    known_cols = _read(out / "known_dogs.csv")[0].keys()
    found_cols = _read(out / "found_dogs.csv")[0].keys()
    assert set(known_cols) == {
        "folder", "image_file", "dog_name", "color", "size", "description",
        "zip", "owner_name", "owner_email", "owner_phone",
    }
    assert set(found_cols) == {
        "folder", "image_file", "description", "color", "size", "found_zip",
        "current_location", "finder_name", "finder_email", "finder_phone",
    }


def test_found_zip_equals_known_zip(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogA": 3})
    out = tmp_path / "out"
    prepare(root, out, holdout=1, seed=3)
    known_zip = _read(out / "known_dogs.csv")[0]["zip"]
    found_zip = _read(out / "found_dogs.csv")[0]["found_zip"]
    assert known_zip == found_zip  # so default-radius matching works


def test_pairs_link_known_and_found(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogA": 4})
    out = tmp_path / "out"
    prepare(root, out, holdout=1, seed=5)
    pairs = _read(out / "pairs.csv")
    assert len(pairs) == 1
    row = pairs[0]
    known_imgs = set(row["known_images"].split(";"))
    found_imgs = set(row["found_images"].split(";"))
    assert len(found_imgs) == 1
    assert len(known_imgs) == 3
    assert known_imgs.isdisjoint(found_imgs)  # an image is either registration or holdout
    assert known_imgs | found_imgs == {"img0.jpg", "img1.jpg", "img2.jpg", "img3.jpg"}


def test_skips_empty_folder(tmp_path):
    root = _make_dataset_dir(tmp_path, {"dogA": 2})
    (root / "empty_dog").mkdir()
    out = tmp_path / "out"
    result = prepare(root, out, holdout=1, seed=1)
    assert "empty_dog" in result["skipped_folders"]
    assert result["dogs"] == 1