"""Validation of generated output images."""

from __future__ import annotations

from pathlib import Path

from PIL import Image

IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}


def index_images(directory: Path) -> dict[str, Path]:
    if not directory.is_dir():
        raise FileNotFoundError(f"directory does not exist: {directory}")
    files = sorted(path for path in directory.iterdir() if path.suffix.lower() in IMAGE_SUFFIXES)
    output: dict[str, Path] = {}
    for path in files:
        if path.stem in output:
            raise ValueError(f"duplicate stem {path.stem!r} in {directory}")
        output[path.stem] = path
    if not output:
        raise ValueError(f"no supported images in {directory}")
    return output


def validate_outputs(input_dir: Path, output_dir: Path, expected_count: int, size: int = 512) -> dict[str, int]:
    inputs = index_images(input_dir)
    outputs = index_images(output_dir)
    if len(inputs) != expected_count:
        raise ValueError(f"input count {len(inputs)} != expected {expected_count}")
    if len(outputs) != expected_count:
        raise ValueError(f"output count {len(outputs)} != expected {expected_count}")
    if set(inputs) != set(outputs):
        missing = sorted(set(inputs) - set(outputs))[:10]
        extra = sorted(set(outputs) - set(inputs))[:10]
        raise ValueError(f"output stem mismatch: missing={missing}, extra={extra}")
    for stem, path in outputs.items():
        try:
            with Image.open(path) as image:
                if image.format != "PNG":
                    raise ValueError(f"not a PNG: {path} ({image.format})")
                if image.size != (size, size):
                    raise ValueError(f"wrong dimensions: {path} is {image.size}, expected {(size, size)}")
                extrema = image.convert("RGB").getextrema()
        except OSError as error:
            raise ValueError(f"invalid image: {path}") from error
        if all(low == high for low, high in extrema):
            raise ValueError(f"constant-color output: {path}")
    return {"inputs": len(inputs), "outputs": len(outputs), "size": size}