#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.13"
# dependencies = ["pillow"]
# ///
import argparse
import enum
from functools import cache
import logging
from pathlib import Path
import struct
from typing import Self
from PIL import Image
# --- Utility Functions ---


class ImageFormat(enum.Enum):
    JPEG = "jpeg"
    PNG = "png"
    WEBP = "webp"

    @property
    def suffix(self) -> str:
        return ".jpg" if self is self.JPEG else f".{self.value}"


def sniff_image_format(image_path: Path) -> ImageFormat:
    """Identify a supported image format from its file signature."""
    with image_path.open("rb") as file:
        header = file.read(12)

    if header.startswith(b"\xff\xd8\xff"):
        return ImageFormat.JPEG
    if header.startswith(b"\x89PNG\r\n\x1a\n"):
        return ImageFormat.PNG
    if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP":
        return ImageFormat.WEBP
    raise ValueError(f"Unsupported or unrecognized image format: {image_path}")


def read_dims_jpeg(image_path: Path) -> tuple[int, int]:
    with image_path.open("rb") as file:
        if file.read(2) != b"\xff\xd8":
            raise ValueError(f"{image_path} is not a valid JPEG file")

        file.seek(0)
        try:
            img = Image.open(file)
            width, height = img.size
            return width, height
        except Exception:
            file.seek(0)
            while True:
                marker = file.read(1)
                if not marker or marker != b"\xff":
                    raise ValueError(f"Invalid JPEG format in {image_path}")
                marker_type = int.from_bytes(file.read(1), byteorder="big")
                length = int.from_bytes(file.read(2), byteorder="big") - 2
                is_sof = 0xC0 <= marker_type <= 0xCF and marker_type not in (0xC4, 0xC8, 0xCC)
                if is_sof:
                    file.seek(1, 1)
                    height = int.from_bytes(file.read(2), byteorder="big")
                    width = int.from_bytes(file.read(2), byteorder="big")
                    return width, height
                file.seek(length, 1)


def read_dims_png(image_path: Path) -> tuple[int, int]:
    """
    Reads PNG dimensions by first attempting PIL, then falling back to manual
    IHDR parsing if the library fails to initialize.
    """
    with image_path.open("rb") as file:
        signature = file.read(8)
        if signature != b"\x89PNG\r\n\x1a\n":
            raise ValueError(f"{image_path} is not a valid PNG file")

        try:
            file.seek(0)
            with Image.open(file) as img:
                return img.size
        except Exception:
            # PNG IHDR is always the first chunk.
            # Offset 8: Chunk Length (4 bytes)
            # Offset 12: Chunk Type "IHDR" (4 bytes)
            # Offset 16: Width (4 bytes)
            # Offset 20: Height (4 bytes)
            file.seek(12)
            if file.read(4) != b"IHDR":
                raise ValueError(f"IHDR chunk not found in {image_path}")

            # Using struct for clean fixed-width big-endian unpacking
            dims = file.read(8)
            if len(dims) < 8:
                raise ValueError(f"Truncated IHDR in {image_path}")

            width, height = struct.unpack(">II", dims)
            return width, height


def read_dims_webp(image_path: Path) -> tuple[int, int]:
    """
    Parses WebP dimensions directly from the RIFF container without
    external dependencies. Handles VP8, VP8L, and VP8X bitstreams.
    """
    with image_path.open("rb") as f:
        header = f.read(12)
        if len(header) < 12 or header[:4] != b"RIFF" or header[8:12] != b"WEBP":
            raise ValueError(f"Not a valid WebP file: {image_path}")

        while True:
            chunk_header = f.read(8)
            if len(chunk_header) < 8:
                break

            tag, length = struct.unpack("<4sI", chunk_header)

            if tag == b"VP8X":
                # Extended Format: Canvas width/height are 24-bit integers
                # starting at offset 4 of the chunk data.
                data = f.read(10)
                width = (int.from_bytes(data[4:7], "little") & 0xFFFFFF) + 1
                height = (int.from_bytes(data[7:10], "little") & 0xFFFFFF) + 1
                return width, height

            elif tag == b"VP8L":
                # Lossless Format: 1 byte signature (0x2f), 14 bits width-1,
                # 14 bits height-1, 1 bit alpha, 3 bits version.
                data = f.read(5)
                if data[0] != 0x2F:
                    raise ValueError("Invalid VP8L signature")

                # Unpack the 4 bytes following the signature
                bits = struct.unpack("<I", data[1:5])[0]
                width = (bits & 0x3FFF) + 1
                height = ((bits >> 14) & 0x3FFF) + 1
                return width, height

            elif tag == b"VP8 ":
                # Lossy Format: Sync code 0x9d012a starts at offset 3.
                # Width/Height are 16-bit values (14 bits used) at offset 6.
                data = f.read(10)
                if data[3:6] != b"\x9d\x01\x2a":
                    raise ValueError("Invalid VP8 sync code")

                width, height = struct.unpack("<HH", data[6:10])
                return width & 0x3FFF, height & 0x3FFF

            # Skip chunk data + padding byte if length is odd
            f.seek((length + 1) & ~1, 1)

    raise ValueError(f"Could not find dimension chunks in {image_path}")


def read_dims(image_path: Path) -> tuple[int, int]:
    image_format = sniff_image_format(image_path)
    if image_format is ImageFormat.JPEG:
        return read_dims_jpeg(image_path)
    elif image_format is ImageFormat.PNG:
        return read_dims_png(image_path)
    elif image_format is ImageFormat.WEBP:
        return read_dims_webp(image_path)
    raise AssertionError(f"Unhandled image format: {image_format}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Sort images based on size, aspect ratio, pose centering, or gaze direction.")
    parser.add_argument("image_paths", nargs="+", type=Path, help="Path(s) to JPEG image(s)")
    strategy = parser.add_mutually_exclusive_group(required=True)
    strategy.add_argument("--by-size", action="store_true", help="Sort images by size")
    strategy.add_argument("--by-aspect-ratio", action="store_true", help="Sort images by aspect ratio")

    return parser.parse_args()


class AspectRatio(enum.Enum):
    PORTRAIT = "portrait"
    WIDE = "wide"

    @classmethod
    def from_dimensions(cls, width: int, height: int) -> Self:
        return cls.PORTRAIT if height > width else cls.WIDE


class Size(enum.Enum):
    SMALL = "small"
    MEDIUM = "medium"
    LARGE = "large"

    @classmethod
    def from_dimensions(cls, width: int, height: int) -> Self:
        smaller = min(width, height)
        if smaller < 1000:
            return cls.SMALL
        elif smaller < 2000:
            return cls.MEDIUM
        return cls.LARGE


@cache
def make_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def main():
    """Main function to process images."""
    args = parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")

    for image_path in args.image_paths:
        if not image_path.is_file():
            logging.warning(f"{image_path} is not a file. Skipping.")
            continue

        try:
            image_format = sniff_image_format(image_path)
            width, height = read_dims(image_path)
            target_dir = None

            if args.by_size:
                size_category = Size.from_dimensions(width, height)
                target_dir = image_path.parent / f"_size_{size_category.value}"
            elif args.by_aspect_ratio:
                aspect_ratio_category = AspectRatio.from_dimensions(width, height)
                target_dir = image_path.parent / f"_aspect_{aspect_ratio_category.value}"
            else:
                raise ValueError("No sorting strategy specified.")

            make_dir(target_dir)

            target_path = target_dir / image_path.name
            target_path = target_path.with_suffix(image_format.suffix)

            image_path.rename(target_path)
            logging.info(f"Moved {image_path} to {target_path}")

        except KeyboardInterrupt:
            raise
        except ValueError:
            logging.exception(f"Error processing {image_path}. Skipping.")
            continue


if __name__ == "__main__":
    main()