from typing import Optional, Union, cast

import numpy as np
import torch
import torch.nn.functional as F
from cv2 import INTER_CUBIC
from cv2 import resize as cv2_resize
from PIL import Image
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_utils import ImageInput, load_image
from transformers.models.auto.processing_auto import AutoProcessor
from transformers.models.qwen3_vl.processing_qwen3_vl import (
    Qwen3VLProcessor,
    Qwen3VLProcessorKwargs,
)
from transformers.processing_utils import Unpack
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.video_utils import VideoInput

from .configuration_qwen3_vl import MobilintQwen3VLConfig


class MobilintQwen3VLProcessor(Qwen3VLProcessor):
    def __call__(
        self,
        images: Optional[ImageInput] = None,
        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput], None] = None,
        videos: Optional[VideoInput] = None,
        **kwargs: Unpack[Qwen3VLProcessorKwargs],
    ) -> BatchFeature:
        assert text is not None, "text is None!"

        while isinstance(images, list):
            if len(images) > 1:
                raise NotImplementedError("Only one image input is supported")
            images = images[0]

        if isinstance(images, str):
            images = load_image(images)

        size = (224, 224)

        if isinstance(images, Image.Image):
            images = images.resize(size)
        elif isinstance(images, np.ndarray):
            if images.ndim == 2:
                images = cast(BatchFeature, cv2_resize(images, size[::-1], interpolation=INTER_CUBIC))
            elif images.ndim == 3:
                images = cast(BatchFeature, cv2_resize(images, size[::-1], interpolation=INTER_CUBIC))
            else:
                raise ValueError(f"Unsupported ndarray shape: {images.shape}")
        elif torch.is_tensor(images):
            if images.ndim == 3:
                images = images.unsqueeze(0).float()
                images = F.interpolate(images, size=size, mode="bicubic", align_corners=False)
            elif images.ndim == 2:
                images = images.unsqueeze(0).unsqueeze(0).float()
                images = F.interpolate(images, size=size, mode="bicubic", align_corners=False)
            else:
                raise ValueError(f"Unsupported tensor shape: {tuple(images.shape)}")
        elif images is not None:
            raise TypeError(f"Unsupported type of image: {type(images)}")

        if videos is not None:
            raise NotImplementedError("Video inputs are not supported")

        return super().__call__(images, text, videos, **kwargs)


AutoProcessor.register(MobilintQwen3VLConfig, MobilintQwen3VLProcessor)
