"""Quicksviewer image ingress: PIL bicubic RGB, CPU uint8, no normalization."""
import base64
import io
import os
from urllib.request import urlopen

import numpy as np
import torch
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor
from transformers.feature_extraction_utils import BatchFeature

SIGLIP_SIZE = 980


def read_source(source):
    if isinstance(source, os.PathLike):
        source = os.fspath(source)
    if isinstance(source, str) and source.startswith("data:"):
        header, payload = source.split(",", 1)
        if ";base64" not in header:
            raise ValueError("Media data URLs must use base64 encoding")
        return io.BytesIO(base64.b64decode(payload))
    if isinstance(source, str) and source.startswith(("https://", "http://")):
        with urlopen(source, timeout=120) as response:
            return io.BytesIO(response.read())
    if isinstance(source, str) and source.startswith("file://"):
        return source[7:]
    if isinstance(source, bytes):
        return io.BytesIO(source)
    return source


def to_pil(image):
    if isinstance(image, Image.Image):
        return image
    if isinstance(image, (str, bytes, os.PathLike)) or hasattr(image, "read"):
        with Image.open(read_source(image)) as opened:
            return opened.convert("RGB")
    if isinstance(image, torch.Tensor):
        image = image.detach().cpu().numpy()
    if not isinstance(image, np.ndarray):
        raise TypeError(f"Unsupported image type: {type(image).__name__}")
    if image.ndim == 3 and image.shape[0] in (1, 3, 4) and image.shape[-1] not in (1, 3, 4):
        image = np.moveaxis(image, 0, -1)
    if image.ndim not in (2, 3):
        raise ValueError("An image must have shape HWC or CHW")
    if image.dtype != np.uint8:
        if not np.isfinite(image).all() or image.min() < 0 or image.max() > 255:
            raise ValueError("Image values must be finite and in [0,255] (or [0,1])")
        if np.issubdtype(image.dtype, np.floating) and image.max() <= 1:
            image = image * 255
        image = np.rint(image).astype(np.uint8)
    if image.ndim == 3 and image.shape[-1] == 1:
        image = image[..., 0]
    return Image.fromarray(image)


class QuicksviewerImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(self, size=None, **kwargs):
        super().__init__(**kwargs)
        self.size = size or {"height": SIGLIP_SIZE, "width": SIGLIP_SIZE}
        if self.size != {"height": SIGLIP_SIZE, "width": SIGLIP_SIZE}:
            raise ValueError("Quicksviewer requires 980x980 images")
        self.do_normalize = False
        self.do_rescale = False

    def preprocess_one(self, image):
        image = to_pil(image).convert("RGB").resize(
            (SIGLIP_SIZE, SIGLIP_SIZE), Image.Resampling.BICUBIC
        )
        array = np.array(image, dtype=np.uint8, copy=True)
        return torch.from_numpy(array).permute(2, 0, 1).contiguous()

    def preprocess(self, images, return_tensors="pt", **kwargs):
        if kwargs:
            raise TypeError(f"Unsupported image options: {sorted(kwargs)}")
        if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4:
            images = list(images)
        elif not isinstance(images, (list, tuple)):
            images = [images]
        frames = [self.preprocess_one(image) for image in images]
        pixels = torch.stack(frames) if frames else torch.empty(
            (0, 3, SIGLIP_SIZE, SIGLIP_SIZE), dtype=torch.uint8
        )
        if return_tensors not in (None, "pt", "np"):
            raise ValueError("return_tensors must be 'pt', 'np', or None")
        if return_tensors == "np":
            pixels = pixels.numpy()
        return BatchFeature({"pixel_values": pixels})


QuicksviewerImageProcessor.register_for_auto_class("AutoImageProcessor")