# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""vLLM support for microsoft/Phi-4-reasoning-vision-15B.

Architecture: Siglip2 vision tower + MLP projector + Phi3 language model.
"""

import math
from collections.abc import Iterable, Mapping, Sequence
from typing import Annotated, Any, Literal

import torch
import torch.nn as nn
from transformers import BatchFeature, PretrainedConfig, Siglip2VisionConfig

from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.inputs import MultiModalDataDict
from vllm.logger import init_logger
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
from vllm.multimodal.parse import (
    ImageSize,
    MultiModalDataItems,
)
from vllm.multimodal.processing import (
    BaseDummyInputsBuilder,
    PromptReplacement,
    PromptUpdate,
)
from vllm.multimodal.processing.processor import (
    BaseMultiModalProcessor,
    BaseProcessingInfo,
)
from vllm.sequence import IntermediateTensors
from vllm.utils.tensor_schema import TensorSchema, TensorShape

from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from .lfm2_siglip2 import Siglip2Model
from .llava import LlavaMultiModalProjector
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    init_vllm_registered_model,
    maybe_prefix,
)

logger = init_logger(__name__)

IMAGE_TOKEN_INDEX = -200
DEFAULT_IMAGE_TOKEN = "<image>"

# The HF processor replaces "<image>" with IMAGE_TOKEN_INDEX (-200) in input_ids.
# Negative token IDs cause OverflowError during decoding, so we remap to a real
# in-vocabulary token.  The Phi-4-reasoning-vision tokenizer ships with reserved
# dummy tokens (<|dummy_0|> … <|dummy_83|>); we reuse the first one as the
# image placeholder.  This mirrors how Phi-3-vision uses its dedicated <|image|>
# token (ID 32044).
_IMAGE_TOKEN_ID = 100256  # <|dummy_0|> in the Phi-4 tokenizer


# ---------------------------------------------------------------------------
# Processing
# ---------------------------------------------------------------------------


class Phi4SiglipProcessingInfo(BaseProcessingInfo):
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
        return {"image": None}

    def _get_vision_config(self) -> dict:
        return self.get_hf_config().vision_config  # type: ignore[attr-defined]

    def _get_patch_size(self) -> int:
        vc = self._get_vision_config()
        if isinstance(vc, dict):
            return vc.get("patch_size", 16)
        return getattr(vc, "patch_size", 16)

    def _get_max_num_patches(self) -> int:
        return getattr(self.get_hf_config(), "max_num_patches", 3600)

    def _get_min_num_patches(self) -> int:
        return getattr(self.get_hf_config(), "min_num_patches", 256)

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
        patch_size = self._get_patch_size()
        min_patches = self._get_min_num_patches()
        max_patches = self._get_max_num_patches()

        num_patches_h = image_height // patch_size
        num_patches_w = image_width // patch_size
        num_patches = max(num_patches_h * num_patches_w, 1)
        num_patches = max(min(num_patches, max_patches), min_patches)
        return num_patches

    def get_image_size_with_most_features(self) -> ImageSize:
        patch_size = self._get_patch_size()
        max_patches = self._get_max_num_patches()
        side = int(math.sqrt(max_patches)) * patch_size
        return ImageSize(width=side, height=side)

    def get_mm_max_tokens_per_item(
        self, seq_len: int, mm_counts: Mapping[str, int]
    ) -> Mapping[str, int]:
        return {"image": self._get_max_num_patches()}


class Phi4SiglipDummyInputsBuilder(
    BaseDummyInputsBuilder[Phi4SiglipProcessingInfo],
):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_images = mm_counts.get("image", 0)
        return DEFAULT_IMAGE_TOKEN * num_images

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
        mm_options: Mapping[str, BaseDummyOptions],
    ) -> MultiModalDataDict:
        num_images = mm_counts.get("image", 0)
        size = self.info.get_image_size_with_most_features()
        return {
            "image": self._get_dummy_images(
                width=size.width,
                height=size.height,
                num_images=num_images,
                overrides=mm_options.get("image"),
            ),
        }


class Phi4SiglipMultiModalProcessor(
    BaseMultiModalProcessor[Phi4SiglipProcessingInfo],
):
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        processed = super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
            tok_kwargs=tok_kwargs,
        )

        # The HF processor's tokenizer_image_token() replaces the "<image>"
        # string with IMAGE_TOKEN_INDEX (-200) in input_ids.  This breaks
        # vLLM's prompt-replacement pipeline which needs to find "<image>"
        # as normal sub-tokens.  Re-tokenize with the plain tokenizer so
        # that "<image>" stays as sub-tokens and can be located by
        # PromptReplacement.
        # NOTE: tokenizer.__call__() (not .encode()) must be used so that
        # added/special tokens like <|user|>, <|end|> are kept as single IDs.
        tokenizer = self.info.get_tokenizer()
        new_ids = tokenizer(prompt).input_ids
        processed["input_ids"] = torch.tensor([new_ids])

        return processed

    def _hf_processor_applies_updates(
        self,
        prompt_text: str,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        tokenization_kwargs: Mapping[str, object],
    ) -> bool:
        # The HF processor replaces "<image>" with a single -200 placeholder
        # but does NOT expand it into N vision-encoder tokens.  Since we also
        # re-tokenize the prompt (see _call_hf_processor), prompt updates are
        # never applied by the HF processor — vLLM handles the expansion via
        # _apply_prompt_updates.
        return False

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            pixel_attention_mask=MultiModalFieldConfig.batched("image"),
            spatial_shapes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
        )

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, Any],
        out_mm_kwargs: MultiModalKwargsItems,
    ) -> Sequence[PromptUpdate]:
        def get_replacement(item_idx: int):
            # Read the actual patch grid from the NaFlex processor's
            # spatial_shapes output (same pattern as LFM2-VL).  This avoids
            # predicting from raw image dimensions, which can diverge from
            # the NaFlex resize/tile logic.
            out_item = out_mm_kwargs["image"][item_idx]
            spatial_shapes = out_item["spatial_shapes"].data
            assert isinstance(spatial_shapes, torch.Tensor)
            num_tokens = int(spatial_shapes.prod().item())
            return [_IMAGE_TOKEN_ID] * num_tokens

        return [
            PromptReplacement(
                modality="image",
                target=DEFAULT_IMAGE_TOKEN,
                replacement=get_replacement,
            ),
        ]


# ---------------------------------------------------------------------------
# Input schemas
# ---------------------------------------------------------------------------


class Phi4SiglipImagePixelInputs(TensorSchema):
    """
    Dimensions:
        - bn: Batch size * number of images
        - d: Max number of patches (padded across images in the batch)
        - fd: Features per patch (patch_size * patch_size * channels)
    """

    type: Literal["pixel_values"] = "pixel_values"
    pixel_values: Annotated[torch.Tensor, TensorShape("bn", "d", "fd")]
    pixel_attention_mask: Annotated[torch.Tensor, TensorShape("bn", "d")]
    spatial_shapes: Annotated[torch.Tensor, TensorShape("bn", 2)]


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------


@MULTIMODAL_REGISTRY.register_processor(
    Phi4SiglipMultiModalProcessor,
    info=Phi4SiglipProcessingInfo,
    dummy_inputs=Phi4SiglipDummyInputsBuilder,
)
class Phi4ForCausalLMV(nn.Module, SupportsMultiModal, SupportsPP):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "model.vision_tower.vision_tower.vision_model.head.": None,
            "model.vision_tower.vision_tower.": "vision_tower.",
            "model.mm_projector.0.": "multi_modal_projector.linear_1.",
            "model.mm_projector.2.": "multi_modal_projector.linear_2.",
            "lm_head.": "language_model.lm_head.",
            "model.": "language_model.model.",
        },
    )

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return DEFAULT_IMAGE_TOKEN
        raise ValueError("Only image modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()

        config: PretrainedConfig = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config

        vision_config_dict: dict = getattr(config, "vision_config", {})
        if isinstance(vision_config_dict, dict):
            if "patch_size" not in vision_config_dict:
                vision_config_dict["patch_size"] = 16
            siglip2_config = Siglip2VisionConfig(**vision_config_dict)
        else:
            siglip2_config = vision_config_dict

        vision_hidden_size: int = config.mm_hidden_size  # type: ignore[attr-defined]
        text_hidden_size: int = config.hidden_size  # type: ignore[attr-defined]

        with self._mark_tower_model(vllm_config, "image"):
            layer_idx = -2
            num_hidden_layers = siglip2_config.num_hidden_layers + layer_idx + 1

            self.vision_tower = Siglip2Model(
                siglip2_config,
                quant_config=quant_config,
                num_hidden_layers_override=num_hidden_layers,
                require_post_norm=False,
                prefix=maybe_prefix(prefix, "vision_tower"),
            )
            self.multi_modal_projector = LlavaMultiModalProjector(
                vision_hidden_size=vision_hidden_size,
                text_hidden_size=text_hidden_size,
                projector_hidden_act="gelu",
                multimodal_projector_bias=True,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "multi_modal_projector"),
            )

        with self._mark_language_model(vllm_config):
            self.language_model = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config,
                prefix=maybe_prefix(prefix, "language_model"),
                architectures=["Phi3ForCausalLM"],
            )

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

        self.configure_mm_token_handling(
            vocab_size=config.vocab_size,  # type: ignore[attr-defined]
            mm_token_ids=[_IMAGE_TOKEN_ID],
        )

    def _packed_from_padded(
        self,
        pixel_values: torch.Tensor,
        pixel_attention_mask: torch.Tensor,
        spatial_shapes: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Convert padded NaFlex tensors to packed format for Siglip2Model."""
        valid_counts = pixel_attention_mask.sum(dim=1).to(torch.int32)
        pixel_values_packed = pixel_values[pixel_attention_mask.bool()]
        cu_seqlens = torch.zeros(
            len(valid_counts) + 1,
            dtype=torch.int32,
            device=pixel_values.device,
        )
        cu_seqlens[1:] = valid_counts.cumsum(0)
        max_seqlen = valid_counts.max()
        return (
            pixel_values_packed,
            spatial_shapes,
            cu_seqlens,
            max_seqlen,
        )

    def _parse_and_validate_image_input(
        self, **kwargs: object
    ) -> Phi4SiglipImagePixelInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        pixel_attention_mask = kwargs.pop("pixel_attention_mask", None)
        spatial_shapes = kwargs.pop("spatial_shapes", None)
        if pixel_values is None:
            return None

        return Phi4SiglipImagePixelInputs(
            type="pixel_values",
            pixel_values=pixel_values,
            pixel_attention_mask=pixel_attention_mask,
            spatial_shapes=spatial_shapes,
        )

    def _process_image_input(
        self, image_input: Phi4SiglipImagePixelInputs
    ) -> MultiModalEmbeddings:
        pixel_values = image_input["pixel_values"]
        pixel_attention_mask = image_input["pixel_attention_mask"]
        spatial_shapes = image_input["spatial_shapes"]

        (
            pixel_values_packed,
            spatial_shapes_packed,
            cu_seqlens,
            max_seqlen,
        ) = self._packed_from_padded(pixel_values, pixel_attention_mask, spatial_shapes)

        vision_features = self.vision_tower(
            pixel_values_packed=pixel_values_packed,
            spatial_shapes=spatial_shapes_packed,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
            select_layers=[-2],
        )

        if vision_features.dim() == 3:
            vision_features = vision_features.squeeze(0)

        image_features = self.multi_modal_projector(vision_features)

        valid_counts = pixel_attention_mask.sum(dim=1).tolist()
        return torch.split(image_features, [int(c) for c in valid_counts])

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
            return []

        return self._process_image_input(image_input)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: object,
    ) -> torch.Tensor | IntermediateTensors:
        if intermediate_tensors is not None:
            inputs_embeds = None

        hidden_states = self.language_model.model(
            input_ids,
            positions,
            intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        return self.language_model.compute_logits(hidden_states)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)