# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Custom processor for Moondream3 model."""

import math

import numpy as np
import torch
from PIL import Image
from transformers import AutoProcessor, BatchFeature
from transformers.image_utils import ImageInput
from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
from transformers.tokenization_utils_base import (
    PreTokenizedInput,
    PreTrainedTokenizerBase,
    TextInput,
)

from vllm.multimodal.image import convert_image_mode

__all__ = ["Moondream3Processor"]


class Moondream3ProcessorKwargs(ProcessingKwargs, total=False):  # type: ignore[call-arg]
    _defaults = {
        "text_kwargs": {
            "padding": False,
        },
        "images_kwargs": {
            "max_crops": 12,
            "overlap_margin": 4,
            "crop_size": 378,
            "patch_size": 14,
            "convert_to_rgb": True,
            "return_tensors": "pt",
        },
    }


def select_tiling(
    height: int, width: int, crop_size: int, max_crops: int
) -> tuple[int, int]:
    """Determine the optimal number of tiles to cover an image."""
    if height <= crop_size or width <= crop_size:
        return (1, 1)

    min_h = math.ceil(height / crop_size)
    min_w = math.ceil(width / crop_size)

    if min_h * min_w > max_crops:
        ratio = math.sqrt(max_crops / (min_h * min_w))
        return (max(1, math.floor(min_h * ratio)), max(1, math.floor(min_w * ratio)))

    h_tiles = math.floor(math.sqrt(max_crops * height / width))
    w_tiles = math.floor(math.sqrt(max_crops * width / height))

    h_tiles = max(h_tiles, min_h)
    w_tiles = max(w_tiles, min_w)

    if h_tiles * w_tiles > max_crops:
        if w_tiles > h_tiles:
            w_tiles = math.floor(max_crops / h_tiles)
        else:
            h_tiles = math.floor(max_crops / w_tiles)

    return (max(1, h_tiles), max(1, w_tiles))


class Moondream3Processor(ProcessorMixin):
    """
    Constructs a Moondream3 processor which handles image preprocessing
    and tokenization for the Moondream3 multimodal model.

    Args:
        tokenizer: The tokenizer to use for text processing.
        chat_template: Optional chat template string.
        crop_size: Size of each image crop.
        max_crops: Maximum number of crops per image.
        overlap_margin: Margin for overlapping crops in patches.
        patch_size: Size of each patch.
    """

    attributes = ["tokenizer"]
    valid_kwargs = [
        "chat_template",
        "crop_size",
        "max_crops",
        "overlap_margin",
        "patch_size",
    ]

    tokenizer_class = "AutoTokenizer"
    # Use separate tokenizer repo
    _tokenizer_repo = "moondream/starmie-v1"

    # Default chat template for Moondream3
    # Moondream uses special tokens for prompting:
    # - Token 0 (<|endoftext|>): BOS token (ALWAYS present at position 0)
    # - Token 1 (<|md_reserved_0|>): Start of instruction
    # - Token 2 (<|md_reserved_1|>): Separator before question
    # - Token 3 (<|md_reserved_2|>): End of question / start of answer
    #
    # Task routing based on text prefix:
    #   "caption [short|normal|long]" → describe<|md_reserved_1|>{length}
    #   "describe [short|normal|long]" → describe<|md_reserved_1|>{length}
    #   otherwise                      → query<|md_reserved_1|><text>
    #
    # Format with image:
    #   <|endoftext|><image><|md_reserved_0|>{task}<|md_reserved_1|>{q}<|md_reserved_2|>
    # Format without image:
    #   <|endoftext|><|md_reserved_0|>{task}<|md_reserved_1|>{q}<|md_reserved_2|>
    _default_chat_template = (
        "{% for message in messages %}"
        "{% if message['role'] == 'user' %}"
        "{% if message['content'] is string %}"
        # Simple string content (with image assumed) - route by prefix
        "<|endoftext|><image><|md_reserved_0|>"
        "{% if message['content'] == 'caption' %}"
        "describe<|md_reserved_1|>normal<|md_reserved_2|>"
        "{% elif message['content'].startswith('caption ') %}"
        "describe<|md_reserved_1|>{{ message['content'][8:] }}<|md_reserved_2|>"
        "{% elif message['content'] == 'describe' %}"
        "describe<|md_reserved_1|>normal<|md_reserved_2|>"
        "{% elif message['content'].startswith('describe ') %}"
        "describe<|md_reserved_1|>{{ message['content'][9:] }}<|md_reserved_2|>"
        "{% else %}"
        "query<|md_reserved_1|>{{ message['content'] }}<|md_reserved_2|>"
        "{% endif %}"
        "{% else %}"
        # List content - build Moondream's image prefix independently of
        # OpenAI-style content part order, then render the text task.
        "<|endoftext|>"
        "{% for content in message['content'] %}"
        "{% if content['type'] in ['image', 'image_url', 'input_image', 'image_pil'] %}"  # noqa: E501
        "<image>"
        "{% endif %}"
        "{% endfor %}"
        "{% for content in message['content'] %}"
        "{% if content['type'] == 'text' %}"
        "<|md_reserved_0|>"
        "{% if content['text'] == 'caption' %}"
        "describe<|md_reserved_1|>normal<|md_reserved_2|>"
        "{% elif content['text'].startswith('caption ') %}"
        "describe<|md_reserved_1|>{{ content['text'][8:] }}<|md_reserved_2|>"
        "{% elif content['text'] == 'describe' %}"
        "describe<|md_reserved_1|>normal<|md_reserved_2|>"
        "{% elif content['text'].startswith('describe ') %}"
        "describe<|md_reserved_1|>{{ content['text'][9:] }}<|md_reserved_2|>"
        "{% else %}"
        "query<|md_reserved_1|>{{ content['text'] }}<|md_reserved_2|>"
        "{% endif %}"
        "{% endif %}"
        "{% endfor %}"
        "{% endif %}"
        "{% elif message['role'] == 'assistant' %}"
        "{{ message['content'] }}"
        "{% endif %}"
        "{% endfor %}"
    )

    def __init__(
        self,
        tokenizer: PreTrainedTokenizerBase | None = None,
        chat_template: str | None = None,
        crop_size: int = 378,
        max_crops: int = 12,
        overlap_margin: int = 4,
        patch_size: int = 14,
        **kwargs,
    ):
        self.image_token = "<image>"
        self.crop_size = crop_size
        self.max_crops = max_crops
        self.overlap_margin = overlap_margin
        self.patch_size = patch_size

        # Number of patches per crop (27x27 = 729 for 378/14)
        self.patches_per_crop = (crop_size // patch_size) ** 2

        # Use default chat template if none provided
        if chat_template is None:
            chat_template = self._default_chat_template

        super().__init__(tokenizer, chat_template=chat_template)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path,
        **kwargs,
    ):
        """
        Load the processor, using a separate tokenizer repo.

        The moondream3 model uses a custom tokenizer from 'moondream/starmie-v1'
        instead of having tokenizer files in the model repo.
        """
        from transformers import AutoTokenizer, TokenizersBackend
        from transformers.utils import cached_file

        tokenizer = kwargs.pop("tokenizer", None)

        tokenizer_kwargs = {
            "trust_remote_code": kwargs.get("trust_remote_code", False),
        }
        for key in (
            "cache_dir",
            "force_download",
            "local_files_only",
            "revision",
            "subfolder",
            "token",
            "use_fast",
        ):
            if key in kwargs:
                tokenizer_kwargs[key] = kwargs[key]

        cached_file_kwargs = {
            key: tokenizer_kwargs[key]
            for key in (
                "cache_dir",
                "force_download",
                "local_files_only",
                "revision",
                "subfolder",
                "token",
            )
            if key in tokenizer_kwargs
        }

        def load_tokenizer(repo_or_path):
            try:
                return AutoTokenizer.from_pretrained(repo_or_path, **tokenizer_kwargs)
            except Exception:
                tokenizer_file = cached_file(
                    repo_or_path,
                    "tokenizer.json",
                    **cached_file_kwargs,
                )
                return TokenizersBackend(
                    tokenizer_file=tokenizer_file,
                    clean_up_tokenization_spaces=False,
                )

        if isinstance(tokenizer, str):
            tokenizer = load_tokenizer(tokenizer)

        if tokenizer is None:
            # Prefer model-local tokenizer files first. If unavailable, fall
            # back to moondream's dedicated tokenizer repository.
            try:
                tokenizer = load_tokenizer(pretrained_model_name_or_path)
            except Exception:
                tokenizer = load_tokenizer(cls._tokenizer_repo)

        # Configure special tokens for Moondream3
        # BOS and EOS are both token 0 (<|endoftext|>), matching the native
        # config (TokenizerConfig.bos_id=0, eos_id=0). This is standard for
        # GPT-2 style models where <|endoftext|> signals both start and end.
        # Token 1 (<|md_reserved_0|>) is a template delimiter, NOT the EOS.
        tokenizer.bos_token = "<|endoftext|>"
        tokenizer.bos_token_id = 0
        tokenizer.eos_token = "<|endoftext|>"
        tokenizer.eos_token_id = 0

        # Extract processor-specific kwargs
        crop_size = kwargs.pop("crop_size", 378)
        max_crops = kwargs.pop("max_crops", 12)
        overlap_margin = kwargs.pop("overlap_margin", 4)
        patch_size = kwargs.pop("patch_size", 14)
        chat_template = kwargs.pop("chat_template", None)

        # Set default chat template on tokenizer if not already set
        if chat_template is None:
            chat_template = cls._default_chat_template
        if tokenizer.chat_template is None:
            tokenizer.chat_template = chat_template

        return cls(
            tokenizer=tokenizer,
            chat_template=chat_template,
            crop_size=crop_size,
            max_crops=max_crops,
            overlap_margin=overlap_margin,
            patch_size=patch_size,
        )

    def __call__(
        self,
        images: ImageInput = None,
        text: TextInput
        | PreTokenizedInput
        | list[TextInput]
        | list[PreTokenizedInput] = None,
        **kwargs: Unpack[Moondream3ProcessorKwargs],
    ) -> BatchFeature:
        """
        Process images and text for Moondream3 model.

        Args:
            images: Input images (PIL Image, numpy array, or list thereof).
            text: Input text or list of texts.
            **kwargs: Additional processing arguments.

        Returns:
            BatchFeature with processed inputs.
        """
        output_kwargs = self._merge_kwargs(
            Moondream3ProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )

        # Process images
        image_features = {}
        if images is not None:
            processed_images = []
            tilings = []

            images_list = images if isinstance(images, list) else [images]
            for image in images_list:
                pixel_values, tiling = self.preprocess_image(
                    image, **output_kwargs["images_kwargs"]
                )
                processed_images.append(pixel_values)
                tilings.append(tiling)

            if processed_images:
                image_features["pixel_values"] = processed_images
                image_features["tilings"] = tilings

        # Process text
        if text is not None:
            if not isinstance(text, list):
                text = [text]

            # Get text kwargs, remove keys we set ourselves
            text_kwargs = output_kwargs.get("text_kwargs", {}).copy()
            text_kwargs.pop("return_tensors", None)
            text_kwargs.pop("add_special_tokens", None)

            # Tokenize text
            tokenized = self.tokenizer(
                text,
                add_special_tokens=True,
                return_tensors="pt",
                **text_kwargs,
            )

            output = BatchFeature(data=dict(tokenized))

            # Add image features
            if image_features:
                output["pixel_values"] = image_features["pixel_values"]
                output["tilings"] = image_features["tilings"]

            return output

        # If only images were provided
        return BatchFeature(data=image_features)

    @staticmethod
    def _image_array_to_uint8(array: np.ndarray) -> np.ndarray:
        if array.dtype == np.uint8:
            return np.ascontiguousarray(array)

        if array.dtype == np.bool_:
            return np.ascontiguousarray(array.astype(np.uint8) * 255)

        if np.issubdtype(array.dtype, np.floating):
            array = np.nan_to_num(array, nan=0.0, posinf=255.0, neginf=0.0)
            if array.size > 0 and array.max() <= 1.0:
                array = array * 255.0
            array = np.rint(array)

        return np.ascontiguousarray(np.clip(array, 0, 255).astype(np.uint8))

    @staticmethod
    def _to_pil_image(image: ImageInput) -> Image.Image:
        if isinstance(image, Image.Image):
            return image

        if isinstance(image, torch.Tensor):
            tensor = image.detach().cpu()
            if tensor.dtype == torch.bfloat16:
                tensor = tensor.to(torch.float32)
            image_array = tensor.numpy()
        elif isinstance(image, np.ndarray):
            image_array = image
        else:
            raise TypeError(
                "Moondream3 images must be PIL images, numpy arrays, "
                f"or torch tensors, got {type(image)!r}."
            )

        if image_array.ndim == 2:
            image_array = Moondream3Processor._image_array_to_uint8(image_array)
            return Image.fromarray(image_array)

        if image_array.ndim != 3:
            raise ValueError(
                "Moondream3 image arrays must have 2 or 3 dimensions, "
                f"got shape {image_array.shape}."
            )

        channel_dims = (1, 3, 4)
        if image_array.shape[-1] not in channel_dims:
            if image_array.shape[0] not in channel_dims:
                raise ValueError(
                    "Moondream3 image arrays must be HWC or CHW with 1, 3, "
                    f"or 4 channels, got shape {image_array.shape}."
                )
            image_array = np.transpose(image_array, (1, 2, 0))

        image_array = Moondream3Processor._image_array_to_uint8(image_array)
        if image_array.shape[-1] == 1:
            image_array = image_array[..., 0]

        return Image.fromarray(image_array)

    def preprocess_image(
        self,
        image: ImageInput,
        max_crops: int = 12,
        overlap_margin: int = 4,
        crop_size: int = 378,
        patch_size: int = 14,
        convert_to_rgb: bool = True,
        return_tensors: str = "pt",
    ) -> tuple[torch.Tensor, tuple[int, int]]:
        """
        Preprocess an image using overlap-and-resize cropping strategy.

        Args:
            image: Input PIL image, numpy array, or torch tensor.
            max_crops: Maximum number of crops.
            overlap_margin: Margin for overlapping in patches.
            crop_size: Size of each crop.
            patch_size: Size of each patch.
            convert_to_rgb: Whether to convert to RGB.
            return_tensors: Return type ("pt" for PyTorch).

        Returns:
            Tuple of (pixel_values tensor, tiling tuple).
        """
        image = self._to_pil_image(image)
        if convert_to_rgb:
            image = convert_image_mode(image, "RGB")

        # Convert to numpy array
        image_array = np.array(image)
        original_h, original_w = image_array.shape[:2]

        margin_pixels = patch_size * overlap_margin
        total_margin_pixels = margin_pixels * 2

        crop_patches = crop_size // patch_size
        crop_window_patches = crop_patches - (2 * overlap_margin)
        crop_window_size = crop_window_patches * patch_size

        tiling = select_tiling(
            original_h - total_margin_pixels,
            original_w - total_margin_pixels,
            crop_window_size,
            max_crops,
        )

        n_crops = tiling[0] * tiling[1] + 1
        crops = np.zeros((n_crops, crop_size, crop_size, 3), dtype=np.uint8)

        target_size = (
            tiling[0] * crop_window_size + total_margin_pixels,
            tiling[1] * crop_window_size + total_margin_pixels,
        )

        # Resize image
        pil_img = Image.fromarray(image_array)
        resized = pil_img.resize(
            (int(target_size[1]), int(target_size[0])),
            resample=Image.Resampling.LANCZOS,
        )
        resized_array = np.asarray(resized)

        # Create global crop
        global_pil = pil_img.resize(
            (crop_size, crop_size), resample=Image.Resampling.LANCZOS
        )
        crops[0] = np.asarray(global_pil)

        # Create local crops
        for i in range(tiling[0]):
            for j in range(tiling[1]):
                y0 = i * crop_window_size
                x0 = j * crop_window_size
                y_end = min(y0 + crop_size, resized_array.shape[0])
                x_end = min(x0 + crop_size, resized_array.shape[1])

                crop_region = resized_array[y0:y_end, x0:x_end]
                crop_idx = 1 + i * tiling[1] + j
                h_slice = slice(None, crop_region.shape[0])
                w_slice = slice(None, crop_region.shape[1])
                crops[crop_idx, h_slice, w_slice] = crop_region

        # Convert to tensor: (n_crops, H, W, C) -> (n_crops, C, H, W)
        pixel_values = np.transpose(crops, (0, 3, 1, 2))

        if return_tensors == "pt":
            # Match HF reference preprocessing exactly: convert uint8 crops to
            # bfloat16 before in-place normalization.
            pixel_values = (
                torch.from_numpy(pixel_values)
                .to(dtype=torch.bfloat16)
                .div_(255.0)
                .sub_(0.5)
                .div_(0.5)
            )
        else:
            pixel_values = pixel_values.astype(np.float32) / 255.0
            pixel_values = (pixel_values - 0.5) / 0.5

        return pixel_values, tiling

    def get_num_image_tokens(self) -> int:
        """Return the number of image tokens (729 = 27x27 patches)."""
        return self.patches_per_crop

    def batch_decode(self, *args, **kwargs):
        """Forward to tokenizer's batch_decode."""
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """Forward to tokenizer's decode."""
        return self.tokenizer.decode(*args, **kwargs)

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        return tokenizer_input_names + ["pixel_values", "tilings"]


AutoProcessor.register("Moondream3Processor", Moondream3Processor)