# SPDX-License-Identifier: Apache-2.0
"""SGLang multimodal processor for glm5v.

Subclass of the in-tree Kimi-K2.5 GPU image processor with GLM's image
placeholder: ``<|image|>`` (id from ``config.media_placeholder_token_id``,
154854) instead of Kimi's ``<|media_pad|>``. The GLM chat template
(``chat_template.jinja`` in the assembled checkpoint) wraps each image as
``<|begin_of_image|><|image|><|end_of_image|>``; only the inner ``<|image|>``
run is expanded to per-patch tokens.

All preprocessing (NaViT resize math, GPU patchify, grid_thw handling) is
inherited unchanged — it is the same MoonViT pipeline the projector was
trained against.
"""

import re

from sglang.srt.multimodal.processors.base_processor import MultimodalSpecialTokens
from sglang.srt.multimodal.processors.kimi_k25 import (
    KimiGPUProcessorWrapper,
    KimiK2_5VLImageProcessor,
)

from sglang_glm5v.glm5v import Glm5vForConditionalGeneration


class Glm5vImageProcessor(KimiK2_5VLImageProcessor):
    models = [Glm5vForConditionalGeneration]

    def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
        # Call the grandparent chain directly: the parent __init__ hardcodes
        # Kimi's <|media_pad|> token; we rebuild mm_tokens + the GPU wrapper
        # with GLM's <|image|> and keep everything else identical.
        super(KimiK2_5VLImageProcessor, self).__init__(
            hf_config, server_args, _processor, *args, **kwargs
        )
        self.mm_tokens = MultimodalSpecialTokens(
            image_token="<|image|>",
            image_token_id=hf_config.media_placeholder_token_id,
            image_token_regex=re.compile(r"(?:<\|image\|>)+"),
        ).build(_processor)

        media_proc_cfg = _processor.media_processor.media_proc_cfg
        self._processor = KimiGPUProcessorWrapper(
            _processor,
            image_token=self.mm_tokens.image_token,
            patch_size=media_proc_cfg["patch_size"],
            merge_kernel_size=media_proc_cfg["merge_kernel_size"],
            in_patch_limit=media_proc_cfg["in_patch_limit"],
            patch_limit_on_one_side=media_proc_cfg["patch_limit_on_one_side"],
            fixed_output_tokens=media_proc_cfg.get("fixed_output_tokens"),
            image_mean=media_proc_cfg["image_mean"],
            image_std=media_proc_cfg["image_std"],
        )