# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Standalone HF processor for PanoVLM: image processor + tokenizer.

Expands each ``<|image|>`` placeholder into ``image_seq_length`` repeated image
tokens wrapped by ``<|begin_of_image|>``/``<|end_of_image|>``, then tokenizes.
``image_seq_length`` is fixed for a given image resolution and is stored in the
processor config at upload time.

``apply_chat_template`` also accepts the list-structured message content used by
newer HF multimodal callers, and normalizes all message content to the string
form the model expects.
"""

from __future__ import annotations

from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin


class PanoVLMProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    # Hardcoded PanoVLM image placeholder tokens (not configurable / not serialized).
    image_token = "<|image|>"
    boi_token = "<|begin_of_image|>"
    eoi_token = "<|end_of_image|>"

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        image_seq_length=None,
        chat_template=None,
        **kwargs,
    ):
        # Set extra scalar attrs before super().__init__ (mirrors LlavaProcessor pattern).
        # ProcessorMixin.__init__ in transformers 5.x only accepts modality args
        # (image_processor, tokenizer) plus chat_template — extra scalars like
        # image_seq_length must be stored before the super() call.
        self.image_seq_length = image_seq_length
        super().__init__(image_processor, tokenizer, chat_template=chat_template)

    def _expand_image_tokens(self, text: str) -> str:
        """Replace each ``<|image|>`` with BOI + N image tokens + EOI."""
        if self.image_seq_length is None:
            raise ValueError(
                "image_seq_length is not set on this processor. Pass "
                "image_seq_length=... when constructing, or load a processor saved "
                "with it set."
            )
        block = (
            self.boi_token + (self.image_token * self.image_seq_length) + self.eoi_token
        )
        return text.replace(self.image_token, block)

    def _prompt_text(self, text: str) -> str:
        """The text to tokenize: image tokens expanded, trailing space removed.

        rstrip before tokenizing: chat templates end with a trailing space (e.g.
        "ASSISTANT: ") that SentencePiece tokenizes as a standalone U+2581 token. That
        token is not expected at this position (the space belongs to the first
        generated token), so it would break the first-token distribution.
        """
        return self._expand_image_tokens(text).rstrip()

    @classmethod
    def _normalize_content(cls, content):
        """Render message content to the exact string form the model expects.

        PanoVLM's chat template concatenates
        ``message['content']`` directly, so content must be a plain string. Callers
        that follow the newer HF multimodal convention (e.g. lmms-eval's generic
        ``huggingface`` wrapper) pass a list of ``{"type": ...}`` parts instead;
        image/video parts become ``<|image|>`` placeholders (expanded later in
        ``__call__``) and text parts are concatenated as-is.

        The resulting string is then split on ``<|image|>``, each text segment
        stripped, empty ones dropped, and everything joined with no separator:
        the model expects ``"<|image|>What is ..."``, so a separator here would
        put tokens after an image that the model does not expect there.
        """
        if not isinstance(content, str):
            pieces = []
            for part in content:
                kind = part.get("type")
                if kind in ("image", "video"):
                    pieces.append(cls.image_token)
                elif kind == "text":
                    pieces.append(part.get("text", ""))
                else:
                    raise ValueError(
                        f"PanoVLM chat content part of type {kind!r} is not "
                        "supported; expected 'text', 'image' or 'video'."
                    )
            content = "".join(pieces)
        normalized = cls.image_token.join(
            chunk.strip() for chunk in content.split(cls.image_token)
        )
        # Keep whitespace-only content untouched.
        return normalized or content

    def apply_chat_template(self, conversation, *args, **kwargs):
        """Normalize message content to the expected string form, then render.

        See :meth:`_normalize_content`. Accepts a single conversation or a batch of
        conversations, matching ``ProcessorMixin.apply_chat_template``.
        """
        is_batched = isinstance(conversation, (list, tuple)) and isinstance(
            conversation[0], (list, tuple)
        )
        conversations = conversation if is_batched else [conversation]
        # Rebuild the messages instead of mutating: the caller (and, for lmms-eval,
        # the code that extracts the images from the same message objects) still
        # needs the original content parts.
        normalized = [
            [
                {**message, "content": self._normalize_content(message["content"])}
                for message in messages
            ]
            for messages in conversations
        ]
        return super().apply_chat_template(
            normalized if is_batched else normalized[0], *args, **kwargs
        )

    def __call__(self, text=None, images=None, return_tensors=None, **kwargs):
        if text is None:
            raise ValueError("PanoVLMProcessor requires `text`.")
        if isinstance(text, str):
            text = [text]
        text = [self._prompt_text(t) for t in text]

        data = {}
        if images is not None:
            image_inputs = self.image_processor(images, return_tensors=return_tensors)
            data.update(image_inputs)
        text_inputs = self.tokenizer(text, return_tensors=return_tensors, **kwargs)
        data.update(text_inputs)
        return BatchFeature(data=data, tensor_type=return_tensors)

    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)