tests/dataset/test_export_image_grid.py

from __future__ import annotations

from pathlib import Path

import pytest
from PIL import Image

from kneiff.datasets.export.image_grid import (
    build_existing_training_export_result,
    TRAINING_IMAGE_GRID_FILENAME,
    write_training_image_grid,
)
from kneiff.datasets.export.types import (
    TrainingExportResult,
    TrainingSamplingResult,
    TrainingSubsetExportResult,
    TrainingSubsetSamplingResult,
)


def _subset_result(subset: str) -> TrainingSubsetExportResult:
    """Return a tiny export subset summary for image-grid tests."""
    return TrainingSubsetExportResult(
        subset=subset,
        source_image_count=1,
        output_image_count=1,
        caption_file_count=1,
        min_outputs_per_source=1,
        max_outputs_per_source=1,
        average_outputs_per_source=1.0,
    )


def _export_result(export_root: Path, *, dry_run: bool = False) -> TrainingExportResult:
    """Return a tiny export result with one training subset and one SFW subset."""
    return TrainingExportResult(
        export_root=export_root,
        source_image_count=2,
        output_image_count=2,
        caption_file_count=2,
        min_outputs_per_source=1,
        max_outputs_per_source=1,
        average_outputs_per_source=1.0,
        subset_results=(_subset_result("fullbody"), _subset_result("sfw")),
        dry_run=dry_run,
    )


def _write_subset_image(
    export_root: Path,
    subset: str,
    color: str,
    *,
    filename: str = "scene__orig.png",
) -> None:
    """Write one exported image below a subset directory."""
    subset_dir = export_root / subset
    subset_dir.mkdir(parents=True, exist_ok=True)
    Image.new("RGB", (16, 16), color=color).save(subset_dir / filename)


def test_training_image_grid_excludes_mirrored_variants(tmp_path: Path) -> None:
    export_root = tmp_path / "export"
    output_path = tmp_path / "grid.png"
    _write_subset_image(export_root, "fullbody", "red")
    _write_subset_image(
        export_root,
        "fullbody",
        "green",
        filename="scene__orig__tag.png",
    )
    _write_subset_image(
        export_root,
        "fullbody",
        "blue",
        filename="scene__aug-mirror.png",
    )
    _write_subset_image(
        export_root,
        "fullbody",
        "yellow",
        filename="scene__aug-mirror__tag.png",
    )

    grid_path = write_training_image_grid(
        _export_result(export_root),
        output_path=output_path,
        thumbnail_size=16,
    )

    assert grid_path == output_path
    with Image.open(output_path) as rendered:
        colors = rendered.convert("RGB").getcolors(maxcolors=100000)
    assert colors is not None
    color_counts = {color: count for count, color in colors}
    assert color_counts.get((255, 0, 0), 0) > 0
    assert color_counts.get((0, 128, 0), 0) > 0
    assert color_counts.get((0, 0, 255), 0) == 0
    assert color_counts.get((255, 255, 0), 0) == 0


def test_training_image_grid_excludes_sfw_when_training_sampling_omits_it(
    tmp_path: Path,
) -> None:
    export_root = tmp_path / "export"
    _write_subset_image(export_root, "fullbody", "red")
    _write_subset_image(export_root, "sfw", "green")
    sampling = TrainingSamplingResult(
        sampling_method="auto-weighting",
        subset_results=(
            TrainingSubsetSamplingResult(
                subset="fullbody",
                probability=1.0,
                relative_percent=100.0,
            ),
        ),
    )

    sampled_grid = write_training_image_grid(
        _export_result(export_root),
        training_sampling=sampling,
        thumbnail_size=16,
    )
    if sampled_grid is None:
        sampled_height = 0
    else:
        with Image.open(sampled_grid) as image:
            sampled_height = image.height
    (export_root / TRAINING_IMAGE_GRID_FILENAME).unlink()
    full_grid = write_training_image_grid(
        _export_result(export_root), thumbnail_size=16
    )
    if full_grid is None:
        full_height = 0
    else:
        with Image.open(full_grid) as image:
            full_height = image.height

    assert sampled_grid == export_root / TRAINING_IMAGE_GRID_FILENAME
    assert sampled_height > 0
    assert sampled_height < full_height


def test_training_image_grid_skips_dry_run(tmp_path: Path) -> None:
    export_root = tmp_path / "export"
    _write_subset_image(export_root, "fullbody", "red")

    grid_path = write_training_image_grid(_export_result(export_root, dry_run=True))

    assert grid_path is None
    assert not (export_root / TRAINING_IMAGE_GRID_FILENAME).exists()


@pytest.mark.parametrize(
    "subset_name",
    ["../../outside", "/absolute", "nested/subset", " padded ", ""],
)
def test_existing_export_result_rejects_non_direct_subset_names(
    tmp_path: Path,
    subset_name: str,
) -> None:
    export_root = tmp_path / "export"
    export_root.mkdir()

    with pytest.raises(ValueError, match="direct directory names"):
        build_existing_training_export_result(
            export_root,
            subset_names=[subset_name],
        )


def test_existing_export_result_rejects_duplicate_subset_names(tmp_path: Path) -> None:
    export_root = tmp_path / "export"
    _write_subset_image(export_root, "fullbody", "red")

    with pytest.raises(ValueError, match="Duplicate training subset"):
        build_existing_training_export_result(
            export_root,
            subset_names=["fullbody", "fullbody"],
        )


def test_existing_export_result_rejects_symlinked_subset(tmp_path: Path) -> None:
    export_root = tmp_path / "export"
    export_root.mkdir()
    outside = tmp_path / "outside"
    _write_subset_image(outside, "images", "red")
    (export_root / "linked").symlink_to(outside / "images", target_is_directory=True)

    with pytest.raises(ValueError, match="must not be symlinks"):
        build_existing_training_export_result(export_root)


def test_existing_export_result_rejects_symlinked_image(tmp_path: Path) -> None:
    export_root = tmp_path / "export"
    subset_dir = export_root / "fullbody"
    subset_dir.mkdir(parents=True)
    outside_image = tmp_path / "outside.png"
    Image.new("RGB", (16, 16), color="red").save(outside_image)
    (subset_dir / "linked.png").symlink_to(outside_image)

    with pytest.raises(ValueError, match="Exported images must not be symlinks"):
        build_existing_training_export_result(export_root)


def test_training_image_grid_rejects_symlinked_output(tmp_path: Path) -> None:
    export_root = tmp_path / "export"
    _write_subset_image(export_root, "fullbody", "red")
    outside = tmp_path / "outside-grid.jpg"
    outside.write_text("keep", encoding="utf-8")
    output_path = export_root / TRAINING_IMAGE_GRID_FILENAME
    output_path.symlink_to(outside)

    with pytest.raises(ValueError, match="Contact sheet output must not cross"):
        write_training_image_grid(_export_result(export_root), thumbnail_size=16)

    assert outside.read_text(encoding="utf-8") == "keep"