#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_molmo2.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The HuggingFace Team. 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 ...image_utils import ImageInput
from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack
from ...tokenization_utils_base import PreTokenizedInput, TextInput
from ...utils import auto_docstring, logging
from ...video_utils import VideoInput


logger = logging.get_logger(__name__)


class Molmo2ProcessorKwargs(ProcessingKwargs, total=False):
    """Molmo2 processor kwargs"""

    _defaults = {
        "text_kwargs": {
            "padding": False,
            "return_mm_token_type_ids": True,
        },
        "videos_kwargs": {"return_metadata": True},
    }


@auto_docstring
class Molmo2Processor(ProcessorMixin):
    valid_processor_kwargs = Molmo2ProcessorKwargs

    @property
    def model_input_names(self):
        return super().model_input_names + ["mm_token_type_ids"]

    def __init__(
        self,
        image_processor=None,
        video_processor=None,
        tokenizer=None,
        chat_template: str | None = None,
        image_use_col_tokens: bool | None = True,
        use_single_crop_col_tokens: bool | None = None,
        use_single_crop_start_token: bool | None = True,
        video_use_col_tokens: bool | None = False,
        use_frame_special_tokens: bool | None = True,
        **kwargs,
    ) -> None:
        r"""
        image_use_col_tokens (`bool`, *optional*, defaults to `True`):
            Whether to append column-separator tokens (`<im_col>`) after each patch row of the high-resolution image
            view.
        use_single_crop_col_tokens (`bool`, *optional*):
            Whether to append column-separator tokens after each patch row of the low-resolution (single-crop) image
            view. If `None`, falls back to `image_use_col_tokens`.
        use_single_crop_start_token (`bool`, *optional*, defaults to `True`):
            Whether to start the low-resolution image view with `<low_res_im_start>` instead of the regular
            `<im_start>`.
        video_use_col_tokens (`bool`, *optional*, defaults to `False`):
            Whether to append column-separator tokens after each patch row of video frames.
        use_frame_special_tokens (`bool`, *optional*, defaults to `True`):
            Whether to wrap each video frame with `<frame_start>` / `<frame_end>` tokens. If `False`, falls back to
            `<im_start>` / `<im_end>`.
        """
        self.image_use_col_tokens = image_use_col_tokens
        self.use_single_crop_col_tokens = use_single_crop_col_tokens
        self.use_single_crop_start_token = use_single_crop_start_token
        self.video_use_col_tokens = video_use_col_tokens
        self.use_frame_special_tokens = use_frame_special_tokens
        self.image_token = getattr(tokenizer, "image_token", "<|image|>")
        self.video_token = getattr(tokenizer, "video_token", "<|video|>")
        self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
        self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token)
        self.image_token_ids = tokenizer.convert_tokens_to_ids(
            [
                "<im_patch>",
                "<im_col>",
                "<im_start>",
                "<low_res_im_start>",
                "<frame_start>",
                "<im_end>",
                "<frame_end>",
                "<im_low>",
            ]
        )
        super().__init__(image_processor, video_processor, tokenizer, chat_template=chat_template)

    @auto_docstring
    def __call__(
        self,
        images: ImageInput | None = None,
        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None,
        videos: VideoInput | None = None,
        **kwargs: Unpack[Molmo2ProcessorKwargs],
    ):
        # TODO(molbap): remove once the allenai chat templates emit `bos_token`
        # (hub PRs allenai/Molmo2-8B#11, allenai/Molmo2-4B#4, allenai/Molmo2-O-7B#2)
        bos_token = self.tokenizer.bos_token or self.tokenizer.eos_token
        if isinstance(text, str):
            text = [text]
        if text is not None:
            text = [prompt if prompt.startswith(bos_token) else bos_token + prompt for prompt in text]
        return super().__call__(images=images, text=text, videos=videos, **kwargs)

    def get_video_string(self, video_grid, timestamps) -> str:
        start_token = "<frame_start>" if self.use_frame_special_tokens else "<im_start>"
        end_token = "<frame_end>" if self.use_frame_special_tokens else "<im_end>"

        num_frames, num_patch_rows, num_patch_cols = video_grid.tolist()
        video_string = ""
        for frame_idx, frame_time in enumerate(timestamps):
            prev_space = " " if frame_idx > 0 else ""
            video_string += prev_space + f"{frame_time:.1f} "
            tokens_per_row = ["<im_patch>"] * num_patch_cols
            if self.video_use_col_tokens:
                tokens_per_row = tokens_per_row + ["<im_col>"]
            video_string += "".join([start_token] + tokens_per_row * num_patch_rows + [end_token])

        return video_string

    def validate_inputs(
        self,
        images=None,
        text=None,
        videos=None,
        audio=None,
        **kwargs,
    ):
        super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs)
        if videos is not None and text is not None:
            for sample in text:
                if sample.count(self.video_token) > 1:
                    raise ValueError("At most one video is supported per sample.")

    def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str:
        image_grid = image_inputs["image_grids"][image_idx]
        if not isinstance(image_grid, list):
            image_grid = image_grid.tolist()
        resized_h, resized_w, height, width = image_grid

        tokens_per_row = ["<im_patch>"] * width
        if self.image_use_col_tokens:
            tokens_per_row = tokens_per_row + ["<im_col>"]
        high_res_tokens = ["<im_start>"] + tokens_per_row * height + ["<im_end>"]

        tokens_per_row = ["<im_patch>"] * resized_w
        use_single_crop_col_tokens = (
            self.image_use_col_tokens if self.use_single_crop_col_tokens is None else self.use_single_crop_col_tokens
        )
        image_start_token = "<low_res_im_start>" if self.use_single_crop_start_token else "<im_start>"
        if use_single_crop_col_tokens:
            tokens_per_row = tokens_per_row + ["<im_col>"]
        low_res_tokens = [image_start_token] + tokens_per_row * resized_h + ["<im_end>"]

        return "".join(low_res_tokens + high_res_tokens)

    def _get_num_multimodal_tokens(self, image_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.

        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 = self._merge_kwargs(self.valid_processor_kwargs, **kwargs)["images_kwargs"]

            num_image_tokens, num_image_patches = [], []
            for image_size in image_sizes:
                image_grid, num_crops = self.image_processor.get_image_grid_and_crops(*image_size, images_kwargs)
                image_string = self.replace_image_token({"image_grids": [image_grid]}, image_idx=0)
                num_image_tokens.append(len(self.tokenizer.tokenize(image_string)))
                num_image_patches.append(num_crops)
            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})

        return MultiModalData(**vision_data)

    def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str:
        video_grid = video_inputs["video_grids"][video_idx]
        metadata = video_inputs["video_metadata"][video_idx]

        frames_indices = metadata.frames_indices
        if frames_indices is None:
            frames_indices = range(int(video_grid[0].item()))
        fps = metadata.fps
        if fps is None:
            fps = self.video_processor.max_fps
            logger.warning_once(
                "Molmo2 inserts frame timestamps into video prompts, but the input video's `fps` was not "
                f"provided or could not be inferred. Defaulting to `fps={fps}`. Please provide "
                "`video_metadata` for more accurate timestamps."
            )
        timestamps = [frame_idx / fps for frame_idx in frames_indices]
        return self.get_video_string(video_grid, timestamps)


__all__ = ["Molmo2Processor"]