"""Image I/O helpers built on Pillow + NumPy.

Pillow is used for decoding (broad format support: PNG/JPG/JPEG/TIFF/BMP)
while all pixel manipulation downstream is done with NumPy/OpenCV arrays in
RGB(A) uint8 layout for consistency across the pipeline.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
from PIL import Image

SUPPORTED_IMAGE_EXTENSIONS: tuple[str, ...] = (
    ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp",
)


def load_image_rgb(path: "Path | str") -> np.ndarray:
    """Load an image file as an ``(H, W, 3)`` uint8 RGB array."""
    with Image.open(path) as im:
        return np.array(im.convert("RGB"))


def load_image_any(path: "Path | str") -> np.ndarray:
    """Load an image preserving its native mode (grayscale, RGB, or RGBA) as uint8.

    Used for signature source images, which may be grayscale scans, palette
    images, or already have an alpha channel.
    """
    with Image.open(path) as im:
        if im.mode in ("L", "RGB", "RGBA"):
            return np.array(im)
        if im.mode == "LA":
            return np.array(im.convert("RGBA"))
        # Palette / CMYK / other -> normalize to RGB
        return np.array(im.convert("RGB"))


def save_image_rgb(path: "Path | str", array: np.ndarray) -> None:
    """Save an ``(H, W, 3)`` or ``(H, W, 4)`` uint8 array as PNG/JPEG depending on extension."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    mode = "RGBA" if array.ndim == 3 and array.shape[2] == 4 else "RGB"
    Image.fromarray(array.astype(np.uint8), mode=mode).save(path)


def save_image_gray(path: "Path | str", array: np.ndarray) -> None:
    """Save a single-channel uint8 array as a grayscale image."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(array.astype(np.uint8), mode="L").save(path)


def discover_images(root: "Path | str", extensions: tuple[str, ...] = SUPPORTED_IMAGE_EXTENSIONS) -> list[Path]:
    """Recursively discover image files under ``root``, sorted deterministically."""
    root = Path(root)
    if not root.exists():
        return []
    found = [p for p in root.rglob("*") if p.is_file() and p.suffix.lower() in extensions]
    return sorted(found, key=lambda p: p.as_posix().lower())