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

import argparse
import atexit
import logging
import os
import re
import shutil
import subprocess
import sys
from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
from contextlib import ExitStack
from dataclasses import dataclass
from functools import partial
from io import BytesIO
from multiprocessing import get_context
from pathlib import Path
from time import perf_counter
from typing import BinaryIO, Callable
from zipfile import BadZipFile, ZipFile, ZipInfo, is_zipfile

from PIL import Image, ImageColor, ImageOps

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
worker_archive: ZipFile | None = None


@dataclass
class Source:
    name: str
    path: Path | None = None
    member: ZipInfo | None = None
    ratio: float = 1.0


@dataclass
class Row:
    sources: list[Source]
    height: int
    widths: list[int]


@dataclass
class Cell:
    source: Source
    position: tuple[int, int]
    size: tuple[int, int]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create a separate justified image contact sheet for each directory or ZIP.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("inputs", nargs="+", type=Path, help="Image directories or ZIP archives; missing paths are skipped")
    parser.add_argument("--output", type=Path, help="JPEG output for a single existing input; defaults to <input-name>_thumbs.jpg beside each input")
    parser.add_argument("--width", type=int, default=2500, help="Sheet width in pixels")
    parser.add_argument("--target-height", type=int, default=300, help="Preferred row height in pixels; rows may be taller or shorter")
    parser.add_argument("--gap", type=int, default=8, help="Gap between images and at sheet edges in pixels")
    parser.add_argument("--bg-color", default="white", help="Background color as a color name or hex value")
    parser.add_argument("--workers", type=int, default=os.cpu_count() or 1, help="Number of resize worker processes")
    arguments = parser.parse_args()
    if arguments.target_height <= 0 or arguments.width <= 0 or arguments.gap < 0:
        parser.error("width and target height must be positive; gap must be nonnegative")
    if arguments.width <= 2 * arguments.gap:
        parser.error("width must exceed twice the gap")
    if arguments.workers <= 0:
        parser.error("workers must be positive")
    try:
        ImageColor.getcolor(color=arguments.bg_color, mode="RGB")
    except ValueError as error:
        parser.error(f"Invalid background color: {error}")
    return arguments


def read_image(source: Source, archive: ZipFile | None) -> Image.Image:
    with ExitStack() as stack:
        if source.member is not None:
            if archive is None:
                raise ValueError("ZIP source requires an open archive")
            stream = stack.enter_context(archive.open(name=source.member))
            original = stack.enter_context(Image.open(fp=stream))
        else:
            original = stack.enter_context(Image.open(fp=source.path))
        oriented = ImageOps.exif_transpose(image=original)
        result = oriented.convert(mode="RGBA")
        if oriented is not original:
            oriented.close()
        return result


def read_dims_jpeg(stream: BinaryIO) -> tuple[int, int]:
    EXIF_ORIENTATION_TAG = 274
    ORIENTATIONS_SWAPPING_WIDTH_HEIGHT = {5, 6, 7, 8}
    with Image.open(fp=stream) as image:
        width, height = image.size
        if image.getexif().get(EXIF_ORIENTATION_TAG) in ORIENTATIONS_SWAPPING_WIDTH_HEIGHT:
            return height, width
        return width, height


def read_dims(image_path: Path | BinaryIO) -> tuple[int, int]:
    EXIF_ORIENTATION_TAG = 274
    ORIENTATIONS_SWAPPING_WIDTH_HEIGHT = {5, 6, 7, 8}
    with ExitStack() as stack:
        stream = stack.enter_context(image_path.open(mode="rb")) if isinstance(image_path, Path) else image_path
        signature = stream.read(12)
        stream.seek(0)
        if signature.startswith(b"\xff\xd8"):
            return read_dims_jpeg(stream=stream)
        if signature.startswith(b"\x89PNG\r\n\x1a\n"):
            raise ValueError("PNG images are not supported")
        if signature[:4] == b"RIFF" and signature[8:12] == b"WEBP":
            raise ValueError("WebP images are not supported")
        with Image.open(fp=stream) as image:
            if image.format == "TIFF":
                raise ValueError("TIFF images are not supported")
            width, height = image.size
            if image.getexif().get(EXIF_ORIENTATION_TAG) in ORIENTATIONS_SWAPPING_WIDTH_HEIGHT:
                return height, width
            return width, height


def natural_sort_key(name: str) -> tuple[str | int, ...]:
    return tuple(int(part) if index % 2 else part.casefold() for index, part in enumerate(re.split(pattern=r"([0-9]+)", string=name)))


def discover_sources(input_path: Path, output_path: Path, archive: ZipFile | None) -> list[Source]:
    Image.init()
    extensions = {extension for extension, image_format in Image.registered_extensions().items() if image_format not in {"PNG", "WEBP", "TIFF"}}
    if archive is not None:
        candidates = [
            Source(name=member.filename, member=member)
            for member in archive.infolist()
            if not member.is_dir() and Path(member.filename).suffix.lower() in extensions
        ]
    else:
        candidates = [
            Source(name=path.relative_to(input_path).as_posix(), path=path)
            for path in input_path.rglob("*")
            if path.is_file() and path.suffix.lower() in extensions and path.resolve() != output_path.resolve()
        ]
    sources = []
    logging.info(f"Reading metadata for {len(candidates)} images")
    for index, source in enumerate(sorted(candidates, key=lambda source: (natural_sort_key(name=source.name), source.name)), start=1):
        try:
            if source.member is not None and archive is not None:
                with archive.open(name=source.member) as stream:
                    width, height = read_dims(image_path=stream)
            elif source.path is not None:
                width, height = read_dims(image_path=source.path)
            else:
                raise ValueError("Image source has no file or ZIP member")
            source.ratio = width / height
            sources.append(source)
        except (OSError, ValueError, SyntaxError, RuntimeError, BadZipFile, Image.DecompressionBombError) as error:
            logging.warning(f"Skipping {source.name}: {error}")
        if index % 100 == 0 or index == len(candidates):
            logging.info(f"Metadata {index}/{len(candidates)} ({index / len(candidates):.0%}); {len(sources)} usable images")
    if not sources:
        raise ValueError("No usable images found")
    return sources


def fitted_height(sources: list[Source], width: int, gap: int) -> float:
    available = width - gap * (len(sources) + 1)
    return available / sum(source.ratio for source in sources)


def make_row(sources: list[Source], width: int, gap: int, target_height: int, justify: bool) -> Row:
    available = width - gap * (len(sources) + 1)
    ideal_height = fitted_height(sources=sources, width=width, gap=gap)
    height = max(1, round(ideal_height if justify else min(target_height, ideal_height)))
    ratio_sum = sum(source.ratio for source in sources)
    total_width = available if justify else min(available, max(len(sources), round(ratio_sum * height)))
    widths = []
    cumulative = 0.0
    previous = 0
    for index, source in enumerate(sources):
        cumulative += source.ratio
        boundary = min(
            total_width - (len(sources) - index - 1),
            max(previous + 1, round(total_width * cumulative / ratio_sum)),
        )
        widths.append(boundary - previous)
        previous = boundary
    return Row(sources=sources, height=height, widths=widths)


def layout_rows(sources: list[Source], width: int, gap: int, target_height: int) -> list[Row]:
    rows = []
    start = 0
    while start < len(sources):
        end = start + 1
        while end < len(sources) and fitted_height(sources=sources[start:end], width=width, gap=gap) > target_height:
            candidate = sources[start : end + 1]
            if width - gap * (len(candidate) + 1) < len(candidate):
                break
            candidate_height = fitted_height(sources=candidate, width=width, gap=gap)
            current_height = fitted_height(sources=sources[start:end], width=width, gap=gap)
            if candidate_height <= target_height and abs(current_height - target_height) <= abs(candidate_height - target_height):
                break
            end += 1
        group = sources[start:end]
        justify = end < len(sources) or fitted_height(sources=group, width=width, gap=gap) <= target_height
        rows.append(make_row(sources=group, width=width, gap=gap, target_height=target_height, justify=justify))
        start = end
    return rows


def initialize_worker(archive_path: Path | None) -> None:
    global worker_archive
    if archive_path is not None:
        worker_archive = ZipFile(file=archive_path)
        atexit.register(worker_archive.close)


def resize_with_pillow(cell: Cell) -> bytes:
    with read_image(source=cell.source, archive=worker_archive) as image:
        with image.resize(size=cell.size, resample=Image.Resampling.LANCZOS) as thumbnail:
            return thumbnail.tobytes()


def resize_with_vips(cell: Cell, vips_binary: str) -> bytes:
    input_bytes = None
    if cell.source.member is not None:
        if worker_archive is None:
            raise ValueError("ZIP source requires an open archive")
        input_bytes = worker_archive.read(name=cell.source.member)
        operation, input_name = "thumbnail_source", "[descriptor=0]"
    elif cell.source.path is not None:
        operation, input_name = "thumbnail", str(cell.source.path)
    else:
        raise ValueError("Image source has no file or ZIP member")
    try:
        result = subprocess.run(
            args=[
                vips_binary,
                "--vips-concurrency=1",
                operation,
                input_name,
                ".png",
                str(cell.size[0]),
                "--height",
                str(cell.size[1]),
                "--size",
                "both",
            ],
            input=input_bytes,
            check=True,
            capture_output=True,
        )
    except (OSError, subprocess.CalledProcessError) as error:
        detail = error.stderr.decode(errors="replace").strip() if isinstance(error, subprocess.CalledProcessError) and error.stderr else str(error)
        logging.warning(f"vips failed for {cell.source.name}; using Pillow: {detail}")
        return resize_with_pillow(cell=cell)
    with Image.open(fp=BytesIO(result.stdout)) as image:
        with image.convert(mode="RGBA") as thumbnail:
            if thumbnail.size == cell.size:
                return thumbnail.tobytes()
            with thumbnail.resize(size=cell.size, resample=Image.Resampling.LANCZOS) as exact_thumbnail:
                return exact_thumbnail.tobytes()


def resize_thumbnail(cell: Cell, processor: Callable[..., bytes]) -> tuple[bytes, float]:
    started = perf_counter()
    logging.info(f"Resizing {cell.source.name} to {cell.size[0]}x{cell.size[1]}")
    pixels = processor(cell=cell)
    return pixels, perf_counter() - started


def render_rows(
    sheet: Image.Image,
    rows: list[Row],
    input_path: Path,
    workers: int,
    gap: int,
    processor: Callable[..., bytes],
) -> int:
    cells = []
    top = gap
    for row in rows:
        left = gap
        for source, cell_width in zip(row.sources, row.widths):
            cells.append(Cell(source=source, position=(left, top), size=(cell_width, row.height)))
            left += cell_width + gap
        top += row.height + gap
    remaining = iter(cells)
    pending: dict[Future[tuple[bytes, float]], Cell] = {}
    processed_count = rendered_count = 0
    archive_path = None if input_path.is_dir() else input_path
    with ProcessPoolExecutor(
        max_workers=workers,
        mp_context=get_context(method="spawn"),
        initializer=initialize_worker,
        initargs=(archive_path,),
    ) as pool:
        while True:
            while len(pending) < workers * 2:
                cell = next(remaining, None)
                if cell is None:
                    break
                future = pool.submit(resize_thumbnail, cell=cell, processor=processor)
                pending[future] = cell
            if not pending:
                break
            completed, _ = wait(fs=pending, return_when=FIRST_COMPLETED)
            for future in completed:
                cell = pending.pop(future)
                processed_count += 1
                try:
                    pixels, duration = future.result()
                    with Image.frombytes(mode="RGBA", size=cell.size, data=pixels) as thumbnail:
                        sheet.paste(im=thumbnail, box=cell.position, mask=thumbnail)
                    rendered_count += 1
                    logging.info(f"Placed {cell.source.name} (resized in {duration:.1f}s)")
                except (OSError, ValueError, SyntaxError, RuntimeError, BadZipFile, Image.DecompressionBombError) as error:
                    logging.warning(f"Skipping {cell.source.name} during rendering: {error}")
                logging.info(f"Progress {processed_count}/{len(cells)} ({processed_count / len(cells):.0%})")
    return rendered_count


def create_sheet(
    input_path: Path,
    output_path: Path,
    width: int,
    target_height: int,
    gap: int,
    bg_color: str = "white",
    workers: int | None = None,
    processor: Callable[..., bytes] = resize_with_pillow,
) -> None:
    if width <= 0 or target_height <= 0 or gap < 0 or width <= 2 * gap:
        raise ValueError("Invalid sheet dimensions or gap")
    if output_path.suffix.lower() not in {".jpg", ".jpeg"}:
        raise ValueError("Output must have a .jpg or .jpeg extension")
    if not input_path.is_dir() and not (input_path.is_file() and is_zipfile(filename=input_path)):
        raise ValueError(f"Expected a directory or ZIP archive: {input_path}")
    if input_path.resolve() == output_path.resolve():
        raise ValueError("Output must differ from input")
    if workers is not None and workers <= 0:
        raise ValueError("Workers must be positive")
    background = ImageColor.getcolor(color=bg_color, mode="RGB")
    with ExitStack() as stack:
        archive = None if input_path.is_dir() else stack.enter_context(ZipFile(file=input_path))
        sources = discover_sources(input_path=input_path, output_path=output_path, archive=archive)
        rows = layout_rows(sources=sources, width=width, gap=gap, target_height=target_height)
        height = sum(row.height for row in rows) + gap * (len(rows) + 1)
        worker_count = min(workers if workers is not None else os.cpu_count() or 1, len(sources))
        logging.info(f"Rendering {len(sources)} images in {len(rows)} rows on a {width}x{height} sheet using {worker_count} workers")
        sheet = stack.enter_context(Image.new(mode="RGB", size=(width, height), color=background))
        started = perf_counter()
        rendered_count = render_rows(
            sheet=sheet,
            rows=rows,
            input_path=input_path,
            workers=worker_count,
            gap=gap,
            processor=processor,
        )
        if not rendered_count:
            raise ValueError("No usable images could be decoded")
        logging.info(f"Rendering finished in {perf_counter() - started:.1f}s; saving {output_path}")
        output_path.parent.mkdir(parents=True, exist_ok=True)
        sheet.save(fp=output_path, quality=95)
    logging.info(f"Wrote {output_path} ({rendered_count} images, {width}x{height})")


def main() -> None:
    arguments = parse_args()
    input_paths = []
    for path in arguments.inputs:
        if path.exists():
            input_paths.append(path.resolve())
        else:
            logging.warning(f"Skipping non-existing input: {path}")
    if not input_paths:
        logging.error("No existing input paths")
        sys.exit(1)
    if arguments.output is not None and len(input_paths) > 1:
        logging.error("--output requires a single existing input")
        sys.exit(2)
    vips_binary = shutil.which(cmd="vips")
    processor = partial(resize_with_vips, vips_binary=vips_binary) if vips_binary else resize_with_pillow
    logging.info(f"Resize backend: {vips_binary or 'Pillow'}")
    failed_count = 0
    for index, input_path in enumerate(input_paths, start=1):
        logging.info(f"Processing input {index}/{len(input_paths)}: {input_path}")
        input_name = input_path.name if input_path.is_dir() else input_path.stem
        output_path = arguments.output or input_path.with_name(f"{input_name}_thumbs.jpg")
        try:
            create_sheet(
                input_path=input_path,
                output_path=output_path,
                width=arguments.width,
                target_height=arguments.target_height,
                gap=arguments.gap,
                bg_color=arguments.bg_color,
                workers=arguments.workers,
                processor=processor,
            )
        except (OSError, ValueError, RuntimeError, BadZipFile) as error:
            failed_count += 1
            logging.error(f"Could not create contact sheet for {input_path}: {error}")
    if failed_count:
        sys.exit(1)


if __name__ == "__main__":
    main()