"""
Hugging Face image processor for Trendyol DinoV2.1 (GeM).

Matches training/inference preprocess used by SimilarityInferenceModel /
dino_v2_gem.Preprocessor: scale-to-max(224), pad-to-square(255), resize 224,
ImageNet normalize. (No JPEG / 332 downscale stage.)
"""

from __future__ import annotations

from typing import List, Optional, Union

import numpy as np
import torch
from PIL import Image
from torchvision import transforms
from torchvision.transforms import functional as TF
from transformers import BatchFeature, ImageProcessingMixin
from transformers.utils import TensorType


def resize_keep_ratio(img: Image.Image, size: int) -> Image.Image:
    w, h = img.size
    max_size = max(h, w)
    scale = size / max_size
    new_size = int(w * scale), int(h * scale)
    return img.resize(new_size, Image.BILINEAR)


class ScaleImage:
    def __init__(self, target_size: int):
        self.target_size = target_size

    def __call__(self, img: Image.Image) -> Image.Image:
        return resize_keep_ratio(img, self.target_size)


class PadToSquare:
    def __init__(self, color: int = 255):
        self.color = color

    def __call__(self, img: Image.Image) -> Image.Image:
        width, height = img.size
        padding = abs(width - height) // 2
        if width < height:
            return TF.pad(
                img,
                (padding, 0, padding + (height - width) % 2, 0),
                fill=self.color,
                padding_mode="constant",
            )
        if width > height:
            return TF.pad(
                img,
                (0, padding, 0, padding + (width - height) % 2),
                fill=self.color,
                padding_mode="constant",
            )
        return img


class TrendyolDinoV21ImageProcessor(ImageProcessingMixin):
    model_input_names = ["pixel_values"]

    def __init__(
        self,
        input_size: int = 224,
        pad_color: int = 255,
        do_normalize: bool = True,
        image_mean=(0.485, 0.456, 0.406),
        image_std=(0.229, 0.224, 0.225),
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.input_size = input_size
        self.pad_color = pad_color
        self.do_normalize = do_normalize
        self.image_mean = list(image_mean)
        self.image_std = list(image_std)

    def _get_preprocess_fn(self):
        steps = [
            ScaleImage(self.input_size),
            PadToSquare(self.pad_color),
            transforms.Resize((self.input_size, self.input_size)),
            transforms.ToTensor(),
        ]
        if self.do_normalize:
            steps.append(transforms.Normalize(self.image_mean, self.image_std))
        return transforms.Compose(steps)

    def __call__(
        self,
        images: Union[Image.Image, np.ndarray, List],
        return_tensors: Optional[Union[str, TensorType]] = None,
        **kwargs,
    ) -> BatchFeature:
        if not isinstance(images, list):
            images = [images]
        preprocess_fn = self._get_preprocess_fn()
        processed = []
        for image in images:
            if isinstance(image, str):
                image = Image.open(image).convert("RGB")
            elif isinstance(image, np.ndarray):
                image = Image.fromarray(image).convert("RGB")
            elif not isinstance(image, Image.Image):
                raise ValueError(f"Unsupported image type: {type(image)}")
            else:
                image = image.convert("RGB")
            processed.append(preprocess_fn(image))
        data = {"pixel_values": torch.stack(processed)}
        return BatchFeature(data=data, tensor_type=return_tensors)


TrendyolDinoV21ImageProcessor.register_for_auto_class("AutoImageProcessor")