# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional, Union, List

import numpy as np

from transformers.feature_extraction_utils import BatchFeature
from transformers.image_utils import ImageInput
from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.video_utils import VideoInput


class NemotronNanoVLV2ImagesKwargs(ImagesKwargs):
    min_pixels: Optional[int]
    max_pixels: Optional[int]
    patch_size: Optional[int]
    temporal_patch_size: Optional[int]
    merge_size: Optional[int]


class NemotronNanoVLV2ProcessorKwargs(ProcessingKwargs, total=False):
    images_kwargs: NemotronNanoVLV2ImagesKwargs
    videos_kwargs: VideosKwargs
    _defaults = {
        "text_kwargs": {
            "padding": False,
        },
    }


class NemotronNanoVLV2Processor(ProcessorMixin):
    r"""
    Constructs a Nemotron Nano VL V2 processor which wraps an image processor and a tokenizer into a single processor.
    [`NemotronNanoVLV2Processor`] offers all the functionalities of the image processor and tokenizer. See the
    [`~NemotronNanoVLV2Processor.__call__`] and [`~NemotronNanoVLV2Processor.decode`] for more information.
    Args:
        image_processor ([`AutoImageProcessor`], *optional*):
            The image processor is a required input.
        tokenizer ([`AutoTokenizer`], *optional*):
            The tokenizer is a required input.
        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
            in a chat into a tokenizable string.
    """

    attributes = ["image_processor", "tokenizer"]

    image_processor_class = "AutoImageProcessor"
    video_processor_class = "AutoVideoProcessor"
    tokenizer_class = ("AutoTokenizer")

    def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
        self.image_token = "<image>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
        self.video_token = "<video>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
        self.image_start_token = "<img>" if not hasattr(tokenizer, "image_start_token") else tokenizer.image_start_token
        self.image_end_token = "</img>" if not hasattr(tokenizer, "image_end_token") else tokenizer.image_end_token
        self.image_token_id = (
            tokenizer.image_token_id
            if getattr(tokenizer, "image_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.image_token)
        )
        self.video_token_id = (
            tokenizer.video_token_id
            if getattr(tokenizer, "video_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.video_token)
        )
        super().__init__(image_processor, tokenizer, chat_template=chat_template)

    def __call__(
        self,
        images: ImageInput = None,
        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
        videos: VideoInput = None,
        **kwargs: Unpack[NemotronNanoVLV2ProcessorKwargs],
    ) -> BatchFeature:
        """
        Main method to prepare multimodal inputs (text, images, videos) for the model. This method processes text by 
        replacing image/video tokens with appropriate placeholder sequences, processes images and videos through the 
        image processor, and tokenizes the final text.

        The method performs the following key operations:
        1. Processes images using the image processor to get pixel values and patch counts
        2. Processes videos using the image processor with max_num_tiles=1 to get video pixel values  
        3. Replaces `<image>` tokens in text with `<img>` + image tokens + `</img>` sequences
        4. Replaces `<video>` tokens in text with frame-by-frame descriptions including timestamps (if metadata provided)
        5. Tokenizes the processed text and combines all outputs

        Args:
            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. Both channels-first and channels-last formats are supported.
            text (`str`, `List[str]`, *optional*):
                The sequence or batch of sequences to be encoded. Each sequence should be a string. The text can contain
                special tokens `<image>` and `<video>` that will be replaced with appropriate token sequences.
            videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):
                The video or batch of videos to be prepared. Each video should be a 4D NumPy array or PyTorch
                tensor with shape (num_frames, channels, height, width). Both channels-first and channels-last formats 
                are supported. Note: Currently only supports batch size of 1 for videos.
            images_kwargs (`Dict`, *optional*):
                Additional keyword arguments for image processing, including:
                - `min_pixels` (`int`, *optional*): Minimum number of pixels for image processing
                - `max_pixels` (`int`, *optional*): Maximum number of pixels for image processing  
                - `patch_size` (`int`, *optional*): Size of patches for image processing
                - `temporal_patch_size` (`int`, *optional*): Size of temporal patches
                - `merge_size` (`int`, *optional*): Size for merging patches
            videos_kwargs (`Dict`, *optional*):
                Additional keyword arguments for video processing, including:
                - `video_metadata` (`VideoMetadata`, *optional*): Metadata containing fps information for timestamp calculation
            text_kwargs (`Dict`, *optional*):
                Additional keyword arguments for text tokenization, including:
                - `return_tensors` (`str` or [`~utils.TensorType`], *optional*): Framework for returned tensors ('tf', 'pt', 'np', 'jax')
                - `padding` (`bool`, *optional*): Whether to pad sequences (defaults to False)

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:

            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
            - **num_patches** -- Number of patches per image. Returned when `images` is not `None`.
            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.

        Raises:
            AssertionError: If videos are provided with batch size > 1 (not currently supported).

        Note:
            - Image tokens `<image>` in text are replaced with `<img>` + repeated image tokens + `</img>`
            - Video tokens `<video>` in text are replaced with frame-by-frame descriptions
            - When video metadata with fps is provided, frame descriptions include timestamps
            - Videos are processed with max_num_tiles=1 regardless of the images setting
        """
        output_kwargs = self._merge_kwargs(
            NemotronNanoVLV2ProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )
        image_inputs = videos_inputs = {}
        if images is not None:
            image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
            image_num_patches = image_inputs["num_patches"]

        if videos is not None:
            orig_tiles = self.image_processor.max_num_tiles
            self.image_processor.max_num_tiles = 1
            videos_inputs = self.image_processor(images=videos, **output_kwargs["images_kwargs"])
            self.image_processor.max_num_tiles = orig_tiles
            video_num_patches = [sum(videos_inputs["num_patches"])]
            videos_inputs["pixel_values_videos"] = videos_inputs["pixel_values"]
            del videos_inputs["pixel_values"]

        if not isinstance(text, list):
            text = [text]

        text = text.copy()  # below lines change text in-place
        if images is not None:
            index = 0
            for i in range(len(text)):
                while self.image_token in text[i]:
                    text[i] = text[i].replace(self.image_token, self.image_start_token + "<|placeholder|>" * image_num_patches[index] * self.image_processor.num_image_token + self.image_end_token, 1)
                    index += 1
                text[i] = text[i].replace("<|placeholder|>", self.image_token)
        if videos is not None:
            assert len(text) == 1, "Video is not supported for batch size > 1"
            video_metadata = output_kwargs.get("videos_kwargs", {}).get("video_metadata", None)
            i = 0
            index = 0
            if self.video_token in text[i]:
                each_frame = self.image_start_token + "<|placeholder|>" * self.image_processor.num_image_token + self.image_end_token
                video_prompt = "This is a video:\n"
                for j in range(video_num_patches[index]):
                    if video_metadata is not None and video_metadata.fps is not None:
                        timestamp = j / video_metadata.fps
                        video_prompt += f"Frame {j+1} sampled at {timestamp:.2f} seconds: {each_frame}\n"
                    else:
                        # Fallback to original format without timestamps
                        video_prompt += f"Frame {j+1}: {each_frame}\n"
                
                text[i] = text[i].replace(self.video_token, video_prompt, 1)
            text[i] = text[i].replace("<|placeholder|>", self.video_token)

        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
        return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors)

    def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
        """
        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
        Args:
            image_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (height, width) per each image.
            video_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (num_frames, height, width) per each video.
        Returns:
            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
            input modalities, along with other useful data.
        """

        vision_data = {}
        if image_sizes is not None:
            images_kwargs = NemotronNanoVLV2ProcessorKwargs._defaults.get("images_kwargs", {})
            images_kwargs.update(kwargs)
            merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size

            num_image_patches = [
                self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
                for image_size in image_sizes
            ]
            num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
        return MultiModalData(**vision_data)

    def batch_decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
        refer to the docstring of this method for more information.
        """
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
        the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, **kwargs)

    def post_process_image_text_to_text(
        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
    ):
        """
        Post-process the output of the model to decode the text.

        Args:
            generated_outputs (`torch.Tensor` or `np.ndarray`):
                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
                or `(sequence_length,)`.
            skip_special_tokens (`bool`, *optional*, defaults to `True`):
                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
            **kwargs:
                Additional arguments to be passed to the tokenizer's `batch_decode method`.

        Returns:
            `list[str]`: The decoded text.
        """
        return self.tokenizer.batch_decode(
            generated_outputs,
            skip_special_tokens=skip_special_tokens,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
            **kwargs,
        )

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
        return names_from_processor + ["second_per_grid_ts"]


__all__ = ["NemotronNanoVLV2Processor"]