"""
Image preprocessing for VLMs — converts PIL images to pixel_values tensors.

Supports:
  - CLIP-style preprocessing (LLaVA-1.5, InternVL2)
  - Dynamic-resolution preprocessing (Qwen2-VL, InternVL2-HD)

Usage::

    from dokodemo_ai.utils.image import preprocess_image, CLIP_DEFAULTS

    pixel_values = preprocess_image(pil_image, **CLIP_DEFAULTS)
    # pixel_values: [1, 3, 336, 336]  float32 tensor
"""

from __future__ import annotations

from typing import Optional, Union
import torch

# CLIP ViT-L/14@336px standard normalisation
CLIP_DEFAULTS = dict(
    size=336,
    mean=(0.48145466, 0.4578275,  0.40821073),
    std= (0.26862954, 0.26130258, 0.27577711),
)

# SigLIP / InternVL2 normalisation (often used by newer VLMs)
SIGLIP_DEFAULTS = dict(
    size=448,
    mean=(0.5, 0.5, 0.5),
    std= (0.5, 0.5, 0.5),
)


def preprocess_image(
    image,                          # PIL.Image.Image or numpy array
    size: int = 336,
    mean: tuple[float, float, float] = CLIP_DEFAULTS["mean"],
    std:  tuple[float, float, float] = CLIP_DEFAULTS["std"],
    device: Optional[torch.device] = None,
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """
    Preprocess a single PIL image into a normalised pixel_values tensor.

    Args:
        image:  PIL.Image.Image (RGB) or numpy array (H, W, 3).
        size:   Target square size (resize + centre-crop to size×size).
        mean:   Per-channel normalisation mean (RGB).
        std:    Per-channel normalisation std (RGB).
        device: Target device.
        dtype:  Target dtype (float32 by default; cast to bfloat16 later if needed).

    Returns:
        Tensor [1, 3, size, size] on device.
    """
    try:
        from PIL import Image as PILImage
        import torchvision.transforms.functional as TF
    except ImportError as e:
        raise ImportError(
            "PIL and torchvision are required for VLM image preprocessing. "
            "Install with: pip install Pillow torchvision"
        ) from e

    # Accept numpy arrays
    if not isinstance(image, PILImage.Image):
        image = PILImage.fromarray(image)

    # Ensure RGB
    if image.mode != "RGB":
        image = image.convert("RGB")

    # Resize so shorter side = size, then centre-crop to size×size
    w, h = image.size
    scale = size / min(w, h)
    new_w, new_h = max(size, round(w * scale)), max(size, round(h * scale))
    image = image.resize((new_w, new_h), PILImage.BICUBIC)

    # Centre crop
    left  = (new_w - size) // 2
    top   = (new_h - size) // 2
    image = image.crop((left, top, left + size, top + size))

    # To tensor [3, H, W], normalise
    t = TF.to_tensor(image)                          # [3, H, W]  float32  [0, 1]
    t = TF.normalize(t, mean=list(mean), std=list(std))

    pixel_values = t.unsqueeze(0).to(dtype=dtype)    # [1, 3, H, W]
    if device is not None:
        pixel_values = pixel_values.to(device=device)
    return pixel_values


def preprocess_images(
    images,          # PIL image, list of PIL images, or [batch, H, W, 3] tensor
    size: int = 336,
    mean: tuple = CLIP_DEFAULTS["mean"],
    std:  tuple = CLIP_DEFAULTS["std"],
    device: Optional[torch.device] = None,
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """
    Batch-preprocess one or more images.

    Returns:
        [batch, 3, size, size]
    """
    if isinstance(images, torch.Tensor):
        # Already a tensor — assume pre-normalised [B, 3, H, W]
        return images.to(device=device, dtype=dtype)

    try:
        from PIL import Image as PILImage
    except ImportError as e:
        raise ImportError("Pillow required: pip install Pillow") from e

    if isinstance(images, PILImage.Image):
        images = [images]

    tensors = [
        preprocess_image(img, size=size, mean=mean, std=std, dtype=dtype)
        for img in images
    ]
    result = torch.cat(tensors, dim=0)   # [B, 3, H, W]
    if device is not None:
        result = result.to(device=device)
    return result


def get_vlm_image_processor_config(config: dict) -> dict:
    """
    Extract image processor settings from a HuggingFace config dict.
    Returns kwargs suitable for preprocess_images().
    """
    proc = config.get("image_processor_type", "")
    vcfg = config.get("vision_config", {})

    size = (
        config.get("image_size")
        or vcfg.get("image_size")
        or (config.get("mm_vision_select_layer") and 336)  # LLaVA default
        or 336
    )
    if isinstance(size, dict):
        size = size.get("shortest_edge", 336)

    # CLIP normalisation is the default for most VLMs
    mean = vcfg.get("image_mean", CLIP_DEFAULTS["mean"])
    std  = vcfg.get("image_std",  CLIP_DEFAULTS["std"])

    # InternVL2 / SigLIP uses 0.5/0.5
    if "siglip" in proc.lower() or "intern" in config.get("model_type", "").lower():
        mean = (0.5, 0.5, 0.5)
        std  = (0.5, 0.5, 0.5)
        size = vcfg.get("image_size", 448)

    return dict(size=size, mean=tuple(mean), std=tuple(std))
