# Copyright 2026 Agnes AI. 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.
"""Processor for Agnes 3.0 Flash: tokenizer + image processor + video processor."""

import numpy as np

from transformers.processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin
from transformers.utils import auto_docstring, logging


logger = logging.get_logger(__name__)


class AgnesProcessorKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "text_kwargs": {"padding": False, "return_token_type_ids": False, "return_mm_token_type_ids": True},
        "videos_kwargs": {"return_metadata": True},
    }


@auto_docstring
class AgnesProcessor(ProcessorMixin):
    valid_processor_kwargs = AgnesProcessorKwargs

    def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs):
        self.image_token = getattr(tokenizer, "image_token", "<|image_pad|>")
        self.video_token = getattr(tokenizer, "video_token", "<|video_pad|>")
        self.image_token_id = getattr(tokenizer, "image_token_id", None) or tokenizer.convert_tokens_to_ids(self.image_token)
        self.video_token_id = getattr(tokenizer, "video_token_id", None) or tokenizer.convert_tokens_to_ids(self.video_token)
        super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
        self.vision_start_token = getattr(tokenizer, "vision_start_token", "<|vision_start|>")
        self.vision_end_token = getattr(tokenizer, "vision_end_token", "<|vision_end|>")
        self.vision_start_token_id = getattr(tokenizer, "vision_start_token_id", None) or tokenizer.convert_tokens_to_ids(
            self.vision_start_token
        )
        self.vision_end_token_id = getattr(tokenizer, "vision_end_token_id", None) or tokenizer.convert_tokens_to_ids(
            self.vision_end_token
        )

    def replace_image_token(self, image_inputs: dict, image_idx: int) -> str:
        per_token = self.image_processor.merge_size**2
        n = image_inputs["image_grid_thw"][image_idx].prod() // per_token
        return self.image_token * n

    def replace_video_token(self, video_inputs: dict, video_idx: int) -> str:
        per_token = self.video_processor.merge_size**2
        thw = video_inputs["video_grid_thw"][video_idx]
        n_frames = thw[0]
        per_frame = thw[1:].prod() // per_token
        meta = video_inputs["video_metadata"][video_idx]
        if meta.fps is None:
            logger.warning_once(
                "Frame timestamps are needed to build the video prompt but the `fps` of the input video could not be "
                "inferred (no `video_metadata`, pre-sampled frames?). Defaulting to `fps=24`."
            )
        meta.fps = 24 if meta.fps is None else meta.fps
        stamps = self._frame_timestamps(meta.frames_indices, meta.fps, self.video_processor.temporal_patch_size)
        text = ""
        for f in range(n_frames):
            text += f"<{stamps[f]:.1f} seconds>"
            text += self.vision_start_token + self.video_token * per_frame + self.vision_end_token
        return text

    def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
        """Placeholder counts for inputs of the given sizes, without running the
        processors on real pixels."""
        data = {}
        if image_sizes is not None:
            ik = AgnesProcessorKwargs._defaults.get("images_kwargs", {})
            ik.update(kwargs)
            merge = ik.get("merge_size", None) or self.image_processor.merge_size
            patches = [self.image_processor.get_number_of_image_patches(*s, ik) for s in image_sizes]
            data.update({"num_image_tokens": [p // merge**2 for p in patches], "num_image_patches": patches})
        if video_sizes is not None:
            vk = AgnesProcessorKwargs._defaults.get("videos_kwargs", {})
            vk.update(kwargs)
            merge = vk.get("merge_size", None) or self.video_processor.merge_size
            patches = [self.video_processor.get_number_of_video_patches(*s, vk) for s in video_sizes]
            data["num_video_tokens"] = [p // merge**2 for p in patches]
        return MultiModalData(**data)

    def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs):
        """Decode generated ids to 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):
        return super().model_input_names + ["mm_token_type_ids"]

    @staticmethod
    def _frame_timestamps(indices: list[int] | np.ndarray, video_fps: float, merge_size: int = 2):
        """One timestamp per temporal patch: the mean of the first and last
        frame time inside the patch."""
        if not isinstance(indices, list):
            indices = indices.tolist()
        if len(indices) % merge_size != 0:
            indices.extend(indices[-1] for _ in range(merge_size - len(indices) % merge_size))
        times = [i / video_fps for i in indices]
        return [(times[i] + times[i + merge_size - 1]) / 2 for i in range(0, len(times), merge_size)]


__all__ = ["AgnesProcessor"]