#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["pytest", "Pillow"]
# ///

from functools import partial
from io import BytesIO
from pathlib import Path
import shutil
import subprocess
import sys
from zipfile import ZIP_DEFLATED, ZipFile

from PIL import Image
import pytest

import imageset_thumbnails as script


@pytest.fixture
def rendered_sheets(monkeypatch: pytest.MonkeyPatch) -> list[Image.Image]:
    sheets = []
    original_save = Image.Image.save

    def capture_sheet(image: Image.Image, **keywords: object) -> None:
        if Path(keywords["fp"]).suffix == ".jpg":
            sheets.append(image.copy())
        original_save(image, **keywords)

    monkeypatch.setattr(Image.Image, "save", capture_sheet)
    return sheets


def test_directory_and_zip_natural_order(tmp_path: Path) -> None:
    directory = tmp_path / "images"
    directory.mkdir()
    for name, color in [("prefix_10.bmp", "blue"), ("prefix_2.bmp", "green"), ("prefix_1.bmp", "red"), ("nested_10/prefix_1.bmp", "yellow"), ("nested_2/prefix_1.bmp", "black")]:
        path = directory / name
        path.parent.mkdir(parents=True, exist_ok=True)
        with Image.new(mode="RGB", size=(20, 20), color=color) as image:
            image.save(fp=path)
    (directory / "broken.bmp").write_bytes(data=b"invalid")
    (directory / "notes.txt").write_text(data="not an image")
    archive_path = tmp_path / "images.zip"
    with ZipFile(file=archive_path, mode="w") as archive:
        for path in sorted(directory.rglob("*"), reverse=True):
            if path.is_file():
                archive.write(filename=path, arcname=path.relative_to(directory).as_posix())
    output = tmp_path / "directory.jpg"
    sources = script.discover_sources(input_path=directory, output_path=output, archive=None)
    assert [source.name for source in sources] == ["nested_2/prefix_1.bmp", "nested_10/prefix_1.bmp", "prefix_1.bmp", "prefix_2.bmp", "prefix_10.bmp"]
    script.create_sheet(input_path=directory, output_path=output, width=68, target_height=20, gap=2)
    zip_output = tmp_path / "zip.jpg"
    script.create_sheet(input_path=archive_path, output_path=zip_output, width=68, target_height=20, gap=2)
    with Image.open(fp=output) as first, Image.open(fp=zip_output) as second:
        assert first.size == second.size
        assert first.tobytes() == second.tobytes()
        assert first.getpixel(xy=(10, 10)) == pytest.approx((0, 0, 0), abs=5)
        assert first.getpixel(xy=(32, 10)) == pytest.approx((255, 255, 0), abs=5)
        assert first.getpixel(xy=(54, 10)) == pytest.approx((255, 0, 0), abs=5)


@pytest.mark.parametrize("ratios,width,target,gap", [
    ([1.0] * 7, 100, 30, 2),
    ([0.5, 2.0, 1.0, 0.75, 1.5], 100, 25, 3),
    ([100.0], 100, 30, 2),
    ([1.0] * 5, 90, 30, 0),
    ([0.001, 100.0, 0.001], 100, 30, 1),
    ([0.5] * 4, 3, 20, 1),
])
def test_layout_geometry(ratios: list[float], width: int, target: int, gap: int) -> None:
    sources = [script.Source(name=str(index), ratio=ratio) for index, ratio in enumerate(ratios)]
    rows = script.layout_rows(sources=sources, width=width, target_height=target, gap=gap)
    assert [source for row in rows for source in row.sources] == sources
    for index, row in enumerate(rows):
        assert row.height >= 1
        assert all(cell_width >= 1 for cell_width in row.widths)
        occupied = sum(row.widths) + gap * (len(row.sources) + 1)
        assert occupied <= width
        if index < len(rows) - 1:
            assert occupied == width


def test_render_gaps_equal_heights_and_last_row(tmp_path: Path, rendered_sheets: list[Image.Image]) -> None:
    directory = tmp_path / "images"
    directory.mkdir()
    for index in range(4):
        with Image.new(mode="RGB", size=(40, 40), color="red") as image:
            image.save(fp=directory / f"image_{index}.bmp")
    output = tmp_path / "sheet.jpg"
    script.create_sheet(input_path=directory, output_path=output, width=100, target_height=30, gap=2)
    with rendered_sheets[0] as sheet:
        assert sheet.size == (100, 67)
        for left in (2, 35, 67):
            assert sheet.getpixel(xy=(left, 2)) == (255, 0, 0)
            assert sheet.getpixel(xy=(left, 32)) == (255, 0, 0)
        for position in [(0, 0), (33, 10), (65, 10), (99, 10), (10, 33), (32, 40)]:
            assert sheet.getpixel(xy=position) == (255, 255, 255)
        assert sheet.getpixel(xy=(2, 35)) == (255, 0, 0)


def test_orientation_and_transparency(tmp_path: Path) -> None:
    path = tmp_path / "oriented.jpg"
    with Image.new(mode="RGB", size=(40, 20), color="red") as image:
        exif = Image.Exif()
        exif[274] = 6
        image.save(fp=path, exif=exif)
    with script.read_image(source=script.Source(name=path.name, path=path), archive=None) as image:
        assert image.size == (20, 40)
    path = tmp_path / "transparent.gif"
    with Image.new(mode="RGBA", size=(10, 10), color=(255, 0, 0, 0)) as image:
        image.save(fp=path)
    with script.read_image(source=script.Source(name=path.name, path=path), archive=None) as image:
        assert image.mode == "RGBA"
        assert image.getpixel(xy=(0, 0)) == (255, 0, 0, 0)


@pytest.mark.parametrize("bg_color,expected", [("white", (255, 255, 255)), ("black", (0, 0, 0)), ("#123456", (18, 52, 86))])
@pytest.mark.parametrize("use_vips", [False, True])
def test_sheet_background(
    tmp_path: Path, rendered_sheets: list[Image.Image], bg_color: str,
    expected: tuple[int, int, int], use_vips: bool,
) -> None:
    vips_binary = shutil.which(cmd="vips") if use_vips else None
    if use_vips and vips_binary is None:
        pytest.skip("vips is not installed")
    processor = partial(script.resize_with_vips, vips_binary=vips_binary) if vips_binary else script.resize_with_pillow
    directory = tmp_path / "images"
    directory.mkdir()
    with Image.new(mode="RGBA", size=(20, 20), color=(255, 0, 0, 0)) as image:
        image.save(fp=directory / "transparent.gif")
    with Image.new(mode="RGB", size=(20, 20), color="red") as image:
        image.save(fp=directory / "red.bmp")
    output = tmp_path / "sheet.jpg"
    script.create_sheet(input_path=directory, output_path=output, width=100, target_height=20, gap=2, bg_color=bg_color, processor=processor)
    with rendered_sheets[0] as sheet:
        assert sheet.getpixel(xy=(0, 0)) == expected
        assert sheet.getpixel(xy=(22, 5)) == expected
        assert sheet.getpixel(xy=(25, 5)) == expected
        assert sheet.getpixel(xy=(99, 5)) == expected
        assert sheet.getpixel(xy=(3, 5)) == (255, 0, 0)


@pytest.mark.parametrize("extension", ["jpg"])
@pytest.mark.parametrize("orientation", [1, 6, 8])
def test_metadata_dimensions_without_pixel_decoding(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
    extension: str, orientation: int,
) -> None:
    directory = tmp_path / "images"
    directory.mkdir()
    path = directory / f"image.{extension}"
    with Image.new(mode="RGB", size=(40, 20), color="red") as image:
        exif = Image.Exif()
        exif[274] = orientation
        image.save(fp=path, exif=exif)
    archive_path = tmp_path / "images.zip"
    with ZipFile(file=archive_path, mode="w") as archive:
        archive.write(filename=path, arcname=path.name)

    def reject_decoding(image: Image.Image, *arguments: object, **keywords: object) -> None:
        pytest.fail("Metadata discovery must not decode image pixels")

    Image.init()
    monkeypatch.setattr(Image.Image, "load", reject_decoding)
    expected = (40, 20) if orientation == 1 else (20, 40)
    assert script.read_dims(image_path=path) == expected
    sources = script.discover_sources(input_path=directory, output_path=tmp_path / "sheet.jpg", archive=None)
    assert sources[0].ratio == expected[0] / expected[1]
    with ZipFile(file=archive_path) as archive:
        with archive.open(name=path.name) as stream:
            assert script.read_dims(image_path=stream) == expected
        sources = script.discover_sources(input_path=archive_path, output_path=tmp_path / "sheet.jpg", archive=archive)
        assert sources[0].ratio == expected[0] / expected[1]


def test_rendering_skips_pixel_corruption(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
    directory = tmp_path / "images"
    directory.mkdir()
    valid = directory / "1.bmp"
    broken = directory / "2.jpg"
    with Image.new(mode="RGB", size=(20, 20), color="red") as image:
        image.save(fp=valid)
        image.save(fp=broken)
    broken.write_bytes(data=broken.read_bytes()[:-10])
    assert script.read_dims(image_path=broken) == (20, 20)
    output = tmp_path / "sheet.jpg"
    script.create_sheet(input_path=directory, output_path=output, width=100, target_height=20, gap=0)
    assert "Skipping 2.jpg during rendering" in caplog.text
    with Image.open(fp=output) as sheet:
        assert sheet.getpixel(xy=(10, 10)) == pytest.approx((255, 0, 0), abs=5)
    valid.unlink()
    with pytest.raises(ValueError, match="No usable images could be decoded"):
        script.create_sheet(input_path=directory, output_path=tmp_path / "failed.jpg", width=100, target_height=20, gap=0)
    assert not (tmp_path / "failed.jpg").exists()


@pytest.mark.parametrize("extension,expected_format", [("jpg", "JPEG"), ("jpeg", "JPEG")])
def test_output_format_and_exclusion(tmp_path: Path, extension: str, expected_format: str) -> None:
    with Image.new(mode="RGB", size=(20, 20), color="blue") as image:
        image.save(fp=tmp_path / "input.bmp")
    output = tmp_path / f"sheet.{extension}"
    script.create_sheet(input_path=tmp_path, output_path=output, width=100, target_height=20, gap=0)
    sources = script.discover_sources(input_path=tmp_path, output_path=output, archive=None)
    assert [source.name for source in sources] == ["input.bmp"]
    with Image.open(fp=output) as sheet:
        assert sheet.format == expected_format
        assert sheet.size == (100, 20)


def test_first_animation_frame(tmp_path: Path) -> None:
    path = tmp_path / "animated.gif"
    with Image.new(mode="RGB", size=(10, 10), color="red") as first, Image.new(mode="RGB", size=(10, 10), color="blue") as second:
        first.save(fp=path, save_all=True, append_images=[second])
    with script.read_image(source=script.Source(name=path.name, path=path), archive=None) as image:
        assert image.getpixel(xy=(0, 0)) == (255, 0, 0, 255)


@pytest.mark.parametrize("input_kind", ["directory", "zip"])
def test_cli_process_pool_matches_single_worker(tmp_path: Path, input_kind: str) -> None:
    directory = tmp_path / "images"
    directory.mkdir()
    for index in [10, 2, 1, 3, 4, 5, 6]:
        with Image.new(mode="RGBA", size=(20 + index * 5, 40), color=(index * 20, 100, 50, 128)) as image:
            exif = Image.Exif()
            exif[274] = 6 if index == 2 else 1
            image.convert(mode="RGB").save(fp=directory / f"image_{index}.jpg", exif=exif)
    input_path = directory
    if input_kind == "zip":
        input_path = tmp_path / "images.zip"
        with ZipFile(file=input_path, mode="w", compression=ZIP_DEFLATED) as archive:
            for path in directory.iterdir():
                archive.write(filename=path, arcname=f"nested/{path.name}")
    outputs = []
    for workers in [1, 2]:
        output = tmp_path / f"workers_{workers}.jpg"
        result = subprocess.run(
            args=[sys.executable, str(Path(script.__file__).resolve()), str(input_path),
                  "--output", str(output), "--workers", str(workers), "--width", "180",
                  "--target-height", "40", "--gap", "3", "--bg-color", "#123456"],
            check=True, capture_output=True, text=True, timeout=30,
        )
        assert "Progress 7/7 (100%)" in result.stderr
        with Image.open(fp=output) as image:
            outputs.append((image.size, image.tobytes()))
    assert outputs[0] == outputs[1]


@pytest.mark.parametrize("options", [
    ["--width", "0"], ["--target-height", "0"], ["--gap", "-1"],
    ["--width", "16", "--gap", "8"],
    ["--bg-color", "invalid-color"],
    ["--workers", "0"], ["--workers", "-1"],
])
def test_invalid_arguments(monkeypatch: pytest.MonkeyPatch, options: list[str]) -> None:
    monkeypatch.setattr(sys, "argv", ["imageset_thumbnails.py", "images", *options])
    with pytest.raises(SystemExit) as error:
        script.parse_args()
    assert error.value.code == 2


def test_cli_names(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(sys, "argv", ["imageset_thumbnails.py", "images", "--target-height", "180", "--gap", "4", "--bg-color", "#123456", "--workers", "3"])
    arguments = script.parse_args()
    assert arguments.target_height == 180
    assert arguments.gap == 4
    assert arguments.bg_color == "#123456"
    assert arguments.workers == 3
    assert arguments.inputs == [Path("images")]


@pytest.mark.parametrize("bad_input_kind,expected_status", [("missing", 0), ("invalid", 1)])
def test_cli_multiple_inputs(tmp_path: Path, bad_input_kind: str, expected_status: int) -> None:
    directory = tmp_path / "photos"
    directory.mkdir()
    with Image.new(mode="RGB", size=(20, 20), color="red") as image:
        image.save(fp=directory / "image.bmp")
    archive_path = tmp_path / "archive.zip"
    with ZipFile(file=archive_path, mode="w", compression=ZIP_DEFLATED) as archive:
        archive.write(filename=directory / "image.bmp", arcname="image.bmp")
    bad_input = tmp_path / "bad_input"
    if bad_input_kind == "invalid":
        bad_input.write_text(data="not a ZIP")
    result = subprocess.run(
        args=[sys.executable, str(Path(script.__file__).resolve()), str(bad_input),
              str(directory), str(archive_path), "--workers", "2"],
        check=False, capture_output=True, text=True, timeout=30,
    )
    assert result.returncode == expected_status
    if bad_input_kind == "missing":
        assert "WARNING: Skipping non-existing input" in result.stderr
    else:
        assert "Could not create contact sheet for" in result.stderr
    for output in [tmp_path / "photos_thumbs.jpg", tmp_path / "archive_thumbs.jpg"]:
        with Image.open(fp=output) as sheet:
            assert sheet.size == (2500, 316)


def test_cli_all_inputs_missing(tmp_path: Path) -> None:
    result = subprocess.run(
        args=[sys.executable, str(Path(script.__file__).resolve()),
              str(tmp_path / "missing.zip"), str(tmp_path / "missing_directory")],
        check=False, capture_output=True, text=True, timeout=30,
    )
    assert result.returncode == 1
    assert result.stderr.count("Skipping non-existing input") == 2
    assert "No existing input paths" in result.stderr
    assert list(tmp_path.iterdir()) == []


def test_cli_output_requires_single_existing_input(tmp_path: Path) -> None:
    first = tmp_path / "first"
    second = tmp_path / "second"
    first.mkdir()
    second.mkdir()
    output = tmp_path / "sheet.jpg"
    result = subprocess.run(
        args=[sys.executable, str(Path(script.__file__).resolve()), str(first),
              str(second), "--output", str(output)],
        check=False, capture_output=True, text=True, timeout=30,
    )
    assert result.returncode == 2
    assert "--output requires a single existing input" in result.stderr
    assert not output.exists()
    with Image.new(mode="RGB", size=(20, 20), color="blue") as image:
        image.save(fp=first / "image.bmp")
    result = subprocess.run(
        args=[sys.executable, str(Path(script.__file__).resolve()), str(first),
              str(tmp_path / "missing"), "--output", str(output)],
        check=True, capture_output=True, text=True, timeout=30,
    )
    assert "Skipping non-existing input" in result.stderr
    assert output.exists()


@pytest.mark.parametrize("use_vips", [False, True])
def test_main_selects_processor_once(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, use_vips: bool,
) -> None:
    vips_binary = shutil.which(cmd="vips") if use_vips else None
    if use_vips and vips_binary is None:
        pytest.skip("vips is not installed")
    directories = [tmp_path / "first", tmp_path / "second"]
    for directory in directories:
        directory.mkdir()
        with Image.new(mode="RGB", size=(40, 20), color="red") as image:
            image.save(fp=directory / "image_1.bmp")
            image.save(fp=directory / "image_2.bmp")
    lookups = []

    def find_binary(cmd: str) -> str | None:
        lookups.append(cmd)
        return vips_binary

    monkeypatch.setattr(script.shutil, "which", find_binary)
    monkeypatch.setattr(sys, "argv", ["imageset_thumbnails.py", *map(str, directories),
                                     "--width", "100", "--target-height", "20", "--workers", "2"])
    script.main()
    assert lookups == ["vips"]
    for name in ["first", "second"]:
        with Image.open(fp=tmp_path / f"{name}_thumbs.jpg") as sheet:
            assert sheet.width == 100


@pytest.mark.parametrize("input_kind", ["directory", "zip"])
def test_vips_resizes_with_legacy_options_before_pillow_reads_pixels(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, input_kind: str,
) -> None:
    vips_binary = shutil.which(cmd="vips")
    if vips_binary is None:
        pytest.skip("vips is not installed")
    path = tmp_path / "image.jpg"
    with Image.new(mode="RGB", size=(400, 200), color="red") as image:
        image.paste(im="blue", box=(200, 0, 400, 200))
        exif = Image.Exif()
        exif[274] = 6
        image.save(fp=path, exif=exif)
    cell = script.Cell(source=script.Source(name=path.name, path=path), position=(0, 0), size=(10, 20))
    archive_path = tmp_path / "images.zip"
    with ZipFile(file=archive_path, mode="w") as archive:
        archive.write(filename=path, arcname=path.name)
    original_open = Image.open
    original_run = subprocess.run
    reads = []

    def run_legacy_vips(args: list[str], **keywords: object) -> subprocess.CompletedProcess[str]:
        assert ".png" in args
        if input_kind == "zip":
            assert "thumbnail_source" in args
            assert keywords["input"] == path.read_bytes()
        else:
            assert "thumbnail" in args
            assert keywords["input"] is None
        if "--fail-on" in args:
            raise subprocess.CalledProcessError(returncode=1, cmd=args, stderr=b"Unknown option --fail-on")
        return original_run(args=args, **keywords)

    def read_thumbnail(fp: BytesIO) -> Image.Image:
        assert isinstance(fp, BytesIO)
        image = original_open(fp=fp)
        reads.append(image.size)
        return image

    monkeypatch.setattr(script.Image, "open", read_thumbnail)
    monkeypatch.setattr(script.subprocess, "run", run_legacy_vips)
    with ZipFile(file=archive_path) as archive:
        if input_kind == "zip":
            cell.source = script.Source(name=path.name, member=archive.getinfo(name=path.name))
            monkeypatch.setattr(script, "worker_archive", archive)
        pixels = script.resize_with_vips(cell=cell, vips_binary=vips_binary)
    assert reads == [cell.size]
    with Image.frombytes(mode="RGBA", size=cell.size, data=pixels) as thumbnail:
        assert thumbnail.getpixel(xy=(5, 2)) == pytest.approx((255, 0, 0, 255), abs=5)
        assert thumbnail.getpixel(xy=(5, 17)) == pytest.approx((0, 0, 255, 255), abs=5)


def test_vips_failure_falls_back_to_pillow(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture,
) -> None:
    path = tmp_path / "image.gif"
    with Image.new(mode="RGBA", size=(40, 20), color=(255, 0, 0, 128)) as image:
        image.save(fp=path)
    cell = script.Cell(source=script.Source(name=path.name, path=path), position=(0, 0), size=(10, 5))
    expected = script.resize_with_pillow(cell=cell)

    def fail_command(args: list[str], **keywords: object) -> subprocess.CompletedProcess[str]:
        raise subprocess.CalledProcessError(returncode=1, cmd=args, stderr=b"unsupported format")

    monkeypatch.setattr(script.subprocess, "run", fail_command)
    assert script.resize_with_vips(cell=cell, vips_binary="vips") == expected
    assert "using Pillow: unsupported format" in caplog.text


@pytest.mark.parametrize("input_kind", ["directory", "zip"])
@pytest.mark.parametrize("image_format,format_name", [("PNG", "PNG"), ("WEBP", "WebP"), ("TIFF", "TIFF")])
def test_unsupported_inputs_are_skipped(tmp_path: Path, input_kind: str, image_format: str, format_name: str) -> None:
    directory = tmp_path / "images"
    directory.mkdir()
    with Image.new(mode="RGB", size=(20, 20), color="red") as image:
        image.save(fp=directory / "supported.jpg")
        image.save(fp=directory / f"ignored.{image_format}", format=image_format)
        image.save(fp=directory / "disguised.jpg", format=image_format)
    output = tmp_path / "sheet.jpg"
    with pytest.raises(ValueError, match=f"{format_name} images are not supported"):
        script.read_dims(image_path=directory / f"ignored.{image_format}")
    if input_kind == "zip":
        archive_path = tmp_path / "images.zip"
        with ZipFile(file=archive_path, mode="w") as archive:
            for path in directory.iterdir():
                archive.write(filename=path, arcname=path.name)
        with ZipFile(file=archive_path) as archive:
            sources = script.discover_sources(input_path=archive_path, output_path=output, archive=archive)
    else:
        sources = script.discover_sources(input_path=directory, output_path=output, archive=None)
    assert [source.name for source in sources] == ["supported.jpg"]


@pytest.mark.parametrize("extension", ["png", "webp", "tif", "tiff"])
def test_unsupported_output_is_rejected(tmp_path: Path, extension: str) -> None:
    output = tmp_path / f"sheet.{extension}"
    with pytest.raises(ValueError, match="Output must have a .jpg or .jpeg extension"):
        script.create_sheet(input_path=tmp_path, output_path=output, width=100, target_height=20, gap=0)
    assert not output.exists()


def test_empty_and_invalid_input(tmp_path: Path) -> None:
    output = tmp_path / "sheet.jpg"
    with pytest.raises(ValueError, match="No usable images"):
        script.create_sheet(input_path=tmp_path, output_path=output, width=100, target_height=20, gap=0)
    with pytest.raises(ValueError, match="Expected a directory or ZIP"):
        script.create_sheet(input_path=tmp_path / "missing", output_path=output, width=100, target_height=20, gap=0)
    assert not output.exists()


if __name__ == "__main__":
    sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))