"""Input coercion for acaua adapters.

One shared entry point (`load_images`) that every adapter calls. Accepts the
full Union documented on `Model.predict`:

    str (file path or http/https URL)
    pathlib.Path
    PIL.Image.Image
    numpy.ndarray   (HWC uint8 RGB)
    torch.Tensor    (CHW float32 in [0, 1])
    Sequence[...]   (any of the above)

Returns a list of PIL RGB images plus a `was_sequence` flag so the caller
knows whether to unwrap the single-item list. All v0.1 adapters are
transformers-backed and their `AutoImageProcessor` consumes PIL directly, so
we stop at PIL rather than producing a model-ready tensor.

BGR-smell heuristic: numpy inputs whose channel-0 mean significantly exceeds
channel-2 mean are rejected with a helpful `InputFormatError` naming the fix
(`img[..., ::-1]` or `cv2.cvtColor(img, COLOR_BGR2RGB)`). Disable via
`strict_rgb=False` if your RGB image genuinely has a warm color cast that
trips the heuristic.
"""

from __future__ import annotations

import io as _io
from collections.abc import Sequence
from pathlib import Path
from urllib.parse import urlparse

import numpy as np
import torch
from PIL import Image

from acaua.exceptions import InputFormatError, WeightDownloadError

# Union of all accepted single-input types. Runtime value (used by isinstance
# indirection via `_coerce_one`); PEP 604 `X | Y` works from 3.10+.
ImageInput = str | Path | Image.Image | np.ndarray | torch.Tensor

_BGR_CHANNEL_DIFF_THRESHOLD = 10.0


def load_images(
    img: ImageInput | Sequence[ImageInput],
    *,
    strict_rgb: bool = True,
) -> tuple[list[Image.Image], bool]:
    """Coerce one input or a sequence to a list of PIL RGB images.

    Args:
        img: single input or sequence of inputs (see module docstring).
        strict_rgb: if True (default), numpy inputs that look BGR are rejected
            with a helpful error. Set False if you have an RGB image that
            trips the heuristic.

    Returns:
        `(pil_images, was_sequence)`. Single input returns `([img], False)`;
        sequence input returns `(imgs, True)`.

    Raises:
        InputFormatError: unsupported type, wrong shape/dtype, empty sequence,
            or BGR-smelling ndarray (when `strict_rgb=True`).
    """
    if isinstance(img, Sequence) and not isinstance(img, (str, bytes)):
        items = list(img)
        if len(items) == 0:
            raise InputFormatError("input sequence is empty")
        return [_coerce_one(x, strict_rgb=strict_rgb) for x in items], True
    return [_coerce_one(img, strict_rgb=strict_rgb)], False


def _coerce_one(x: ImageInput, *, strict_rgb: bool) -> Image.Image:
    if isinstance(x, Image.Image):
        return x.convert("RGB") if x.mode != "RGB" else x

    if isinstance(x, (str, Path)):
        return _open_path_or_url(str(x))

    if isinstance(x, np.ndarray):
        if x.ndim != 3 or x.shape[-1] != 3:
            raise InputFormatError(f"numpy input must be HWC with 3 channels; got shape {x.shape}")
        if x.dtype != np.uint8:
            raise InputFormatError(f"numpy input must be uint8 (values 0-255); got dtype {x.dtype}")
        if strict_rgb and _bgr_smell(x):
            raise InputFormatError(
                "numpy input looks like BGR (OpenCV default), but acaua expects RGB. "
                "If this is cv2.imread output, convert first: img[..., ::-1] or "
                "cv2.cvtColor(img, cv2.COLOR_BGR2RGB). "
                "Pass strict_rgb=False to .predict() to disable this check."
            )
        return Image.fromarray(x, mode="RGB")

    if isinstance(x, torch.Tensor):
        if x.ndim != 3 or x.shape[0] != 3:
            raise InputFormatError(
                f"torch input must be CHW with 3 channels; got shape {tuple(x.shape)}"
            )
        if x.dtype != torch.float32:
            raise InputFormatError(f"torch input must be float32 in [0, 1]; got dtype {x.dtype}")
        arr = (x.clamp(0, 1) * 255).to(torch.uint8).permute(1, 2, 0).cpu().numpy()
        return Image.fromarray(arr, mode="RGB")

    raise InputFormatError(
        f"unsupported input type: {type(x).__name__}. "
        "Expected str (path or http/https URL), pathlib.Path, PIL.Image.Image, "
        "numpy.ndarray (HWC uint8 RGB), or torch.Tensor (CHW float32 in [0, 1])."
    )


def _open_path_or_url(s: str) -> Image.Image:
    if _is_url(s):
        import requests  # lazy: the majority of users load from disk

        try:
            resp = requests.get(s, timeout=30)
            resp.raise_for_status()
        except requests.RequestException as e:
            raise WeightDownloadError(f"failed to fetch image from URL {s!r}: {e}") from e
        return Image.open(_io.BytesIO(resp.content)).convert("RGB")
    return Image.open(s).convert("RGB")


def _is_url(s: str) -> bool:
    return urlparse(s).scheme in {"http", "https"}


def _bgr_smell(arr: np.ndarray) -> bool:
    """Channel-stats heuristic: suggests an ndarray is BGR.

    Only triggers on 3-channel uint8 HWC inputs (others are rejected before
    reaching here). Threshold chosen empirically: skin tones and natural
    images rarely have ch0 > ch2 by more than ~5; OpenCV BGR output flips
    the R/B channels and pushes the gap above 10 for most real photos.
    """
    if arr.ndim != 3 or arr.shape[-1] != 3 or arr.dtype != np.uint8:
        return False
    diff = float(arr[..., 0].mean()) - float(arr[..., 2].mean())
    return diff > _BGR_CHANNEL_DIFF_THRESHOLD
