"""Image preprocessing using Pillow + numpy (replaces SAM2Transforms)."""

import logging

import numpy as np
from PIL import Image

from sam2_onnx.types import ImageInput, PreprocessResult

logger = logging.getLogger(__name__)

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


def preprocess_image(
    image: ImageInput,
    resolution: int,
) -> PreprocessResult:
    """Preprocess an image for SAM 2 inference.

    Args:
        image: RGB image as HWC uint8 numpy array or PIL Image.
        resolution: Target resolution (square).

    Returns:
        Tuple of (preprocessed NCHW float32 array, original (H, W)).
    """
    # Validate resolution
    if resolution <= 0:
        raise ValueError(f"resolution must be positive, got {resolution}")

    # Validate numpy array input
    if isinstance(image, np.ndarray):
        if image.ndim != 3:
            raise ValueError(
                f"Expected HWC image with 3 dimensions, got ndim={image.ndim}"
            )
        if image.shape[2] != 3:
            raise ValueError(
                f"Expected 3-channel RGB image, got {image.shape[2]} channels"
            )
        orig_hw = (image.shape[0], image.shape[1])
        image = Image.fromarray(image)
    else:
        orig_hw = (image.height, image.width)

    # Resize to square resolution using bilinear (matches torchvision)
    image = image.resize((resolution, resolution), Image.Resampling.BILINEAR)

    # Convert to float32 [0, 1]
    img = np.array(image, dtype=np.float32) / 255.0

    # Normalize with ImageNet stats (HWC layout)
    img = (img - IMAGENET_MEAN) / IMAGENET_STD

    # HWC -> CHW -> NCHW
    img = np.transpose(img, (2, 0, 1))[np.newaxis]

    logger.debug(
        "preprocessed: %s -> %s, orig_hw=%s", type(image).__name__, img.shape, orig_hw
    )
    return img, orig_hw
