"""OXR model implementation for vLLM 0.25.1."""

from collections.abc import Iterable, Mapping
from functools import partial
from typing import Annotated, Literal

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

from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.inputs import MultiModalDataDict
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.models.interfaces import (
    MultiModalEmbeddings,
    SupportsMultiModal,
    SupportsPP,
)
from vllm.model_executor.models.siglip2navit import Siglip2NavitModel
from vllm.model_executor.models.utils import (
    AutoWeightsLoader,
    flatten_bn,
    init_vllm_registered_model,
    maybe_prefix,
)
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,
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
)
from vllm.renderers import TokenizeParams
from vllm.sequence import IntermediateTensors
from vllm.utils.tensor_schema import TensorSchema, TensorShape

from .oxr_processor import OXRProcessor

IMAGE_TOKEN = "<image>"
VIDEO_TOKEN = "<video>"
IMAGE_PLACEHOLDER_ID = -200
VIDEO_PLACEHOLDER_ID = -201
VISUAL_ATOM_ID = -300
OXR_INDICATOR_IDS = [-301, -302, -303, -304]

IMAGE_PAD_TOKEN_MAP = {
    "gemma2": "<unused0>",
    "llama": "<|reserved_special_token_0|>",
    "qwen2": "<|image_pad|>",
    "qwen3": "<|image_pad|>",
}
IMAGE_PAD_TOKEN_ID_MAP = {
    "gemma2": 7,
    "llama": 128002,
    "qwen2": 151655,
    "qwen3": 151655,
}


def get_llm_config(config: PretrainedConfig) -> PretrainedConfig:
    llm_config = getattr(config, "llm_config", None)
    if llm_config is not None:
        return llm_config

    text_config = getattr(config, "text_config", None)
    if text_config is not None:
        return text_config

    get_text_config = getattr(config, "get_text_config", None)
    if get_text_config is not None:
        candidate = get_text_config()
        if candidate is not config:
            return candidate

    raise ValueError("OXR config does not expose an LLM sub-config")


def _oxr_field_config() -> Mapping[str, MultiModalFieldConfig]:
    return dict(
        pixel_values=MultiModalFieldConfig.batched("image"),
        grids=MultiModalFieldConfig.batched("image"),
        indicator_tokens=MultiModalFieldConfig.batched("image"),
        video_pixel_values=MultiModalFieldConfig.batched("video"),
        video_indicator_tokens=MultiModalFieldConfig.batched("video"),
        video_grids=MultiModalFieldConfig.batched("video"),
    )


class OXRImagePatchInputs(TensorSchema):
    type: Literal["image_patches"]
    flat_data: Annotated[torch.Tensor, TensorShape("bnp", "patch_size")]
    indicator_tokens: Annotated[torch.Tensor, TensorShape("indicators")]
    patches_per_item: Annotated[list[int], TensorShape("bn")]
    grids: Annotated[torch.Tensor, TensorShape("bn", 3)]


class OXRVideoPatchInputs(TensorSchema):
    type: Literal["video_patches"]
    flat_data: Annotated[torch.Tensor, TensorShape("bnp", "patch_size")]
    indicator_tokens: Annotated[torch.Tensor, TensorShape("indicators")]
    patches_per_item: Annotated[list[int], TensorShape("bn")]
    grids: Annotated[torch.Tensor, TensorShape("bn", 3)]


class OXREmbedding(torch.nn.Embedding):
    def forward(self, visual_tokens: torch.Tensor) -> torch.Tensor:
        if visual_tokens.dtype in (
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.long,
        ):
            return super().forward(visual_tokens)
        return torch.matmul(visual_tokens, self.weight)

    @property
    def device(self) -> torch.device:
        return self.weight.device

    @property
    def dtype(self) -> torch.dtype:
        return self.weight.dtype


class OXRVisualTokenizer(torch.nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        visual_vocab_size: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.config = config
        self.vit = self._init_backbone(
            config=config,
            quant_config=quant_config,
            prefix=f"{prefix}.vit",
        )
        head_dim = visual_vocab_size - len(OXR_INDICATOR_IDS)
        self.head = torch.nn.Sequential(
            ReplicatedLinear(
                self.config.hidden_size * self.config.hidden_stride**2,
                head_dim,
                bias=False,
                return_bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.head.0",
            ),
            torch.nn.LayerNorm(head_dim),
        )

    def _init_backbone(
        self,
        config: PretrainedConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        if config.model_type == "siglip2_navit":
            return Siglip2NavitModel(
                config=config,
                quant_config=quant_config,
                prefix=prefix,
            )
        raise ValueError(f"Unsupported visual tokenizer model_type: {config.model_type}")

    @property
    def dtype(self) -> torch.dtype:
        return next(self.head.parameters()).dtype

    @property
    def device(self) -> torch.device:
        return next(self.head.parameters()).device

    def tokenize(self, logits: torch.Tensor) -> torch.Tensor:
        return torch.softmax(logits, dim=-1, dtype=torch.float32).to(logits.dtype)

    def encode(
        self,
        pixel_values: torch.Tensor,
        grid_thws: torch.Tensor,
    ) -> torch.Tensor:
        features = self.vit(pixel_values, grid_thws)
        seq_len, _ = features.shape
        return features.reshape(seq_len // (self.config.hidden_stride**2), -1)

    def forward(
        self,
        pixel_values: torch.Tensor,
        grid_thws: torch.Tensor,
    ) -> torch.Tensor:
        features = self.encode(pixel_values, grid_thws)
        logits = self.head(features)
        tokens = self.tokenize(logits)
        return torch.nn.functional.pad(
            tokens,
            (0, len(OXR_INDICATOR_IDS)),
            mode="constant",
            value=0,
        )


class OXRProcessingInfo(BaseProcessingInfo):
    def get_hf_config(self):
        return self.ctx.get_hf_config()

    def get_hf_processor(self, **kwargs):
        vit_config = self.get_hf_config().vit_config
        return self.ctx.get_hf_processor(
            OXRProcessor,
            image_pad_token=self.get_image_pad_token(),
            patch_size=vit_config.patch_size,
            hidden_stride=vit_config.hidden_stride,
            temporal_patch_size=vit_config.temporal_patch_size,
            **kwargs,
        )

    def get_default_tok_params(self) -> TokenizeParams:
        return super().get_default_tok_params().with_kwargs(add_special_tokens=False)

    def get_image_pad_token(self) -> str:
        text_model_type = get_llm_config(self.get_hf_config()).model_type
        return IMAGE_PAD_TOKEN_MAP[text_model_type]

    def get_image_processor(self) -> BaseImageProcessor:
        return self.get_hf_processor().image_processor  # type: ignore[return-value]

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

    def get_image_size_with_most_features(self) -> ImageSize:
        return ImageSize(width=1792, height=1792)

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
        num_frames: int = 1,
    ) -> int:
        vit_config = self.get_hf_config().vit_config
        padded_num_frames = num_frames + (-num_frames % vit_config.temporal_patch_size)
        grid_t = max(padded_num_frames // vit_config.temporal_patch_size, 1)
        grid_h = image_height // vit_config.patch_size
        grid_w = image_width // vit_config.patch_size
        return grid_t * grid_h * grid_w

    def get_max_image_tokens(self) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
        return self.get_num_image_tokens(
            image_width=target_width,
            image_height=target_height,
        )

    def _get_max_video_frames(self, max_tokens: int) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
        num_frames = 0
        while True:
            next_num_frames = num_frames + 1
            next_max_tokens = self.get_num_video_tokens(
                image_width=target_width,
                image_height=target_height,
                num_frames=next_num_frames,
            )
            if next_max_tokens > max_tokens:
                break
            num_frames = next_num_frames
        return num_frames

    def get_num_frames_with_most_features(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> int:
        max_images = mm_counts.get("image", 0)
        max_videos = mm_counts.get("video", 0)
        max_image_tokens = self.get_max_image_tokens() * max_images
        max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens)
        max_frames_per_video = max_total_frames // max(max_videos, 1)
        return max(max_frames_per_video, 1)

    def get_num_video_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
        num_frames: int,
    ) -> int:
        return self.get_num_image_tokens(
            image_width=image_width,
            image_height=image_height,
            num_frames=num_frames,
        )

    def get_max_video_tokens(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
        return self.get_num_video_tokens(
            image_width=target_width,
            image_height=target_height,
            num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
        )

    def get_mm_max_tokens_per_item(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> Mapping[str, int]:
        # Return a conservative upper bound so vLLM does not profile the
        # visual tower with dozens of maximum-sized dummy images/videos.
        return {
            "image": self.get_max_image_tokens() + 2,
            "video": self.get_max_video_tokens(seq_len, mm_counts) + 2,
        }


class OXRDummyInputsBuilder(BaseDummyInputsBuilder[OXRProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return IMAGE_TOKEN * mm_counts.get("image", 0) + VIDEO_TOKEN * mm_counts.get(
            "video", 0
        )

    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)
        num_videos = mm_counts.get("video", 0)
        target_width, target_height = self.info.get_image_size_with_most_features()
        target_num_frames = self.info.get_num_frames_with_most_features(
            seq_len, mm_counts
        )

        return {
            "image": self._get_dummy_images(
                width=target_width,
                height=target_height,
                num_images=num_images,
                overrides=mm_options.get("image"),
            ),
            "video": self._get_dummy_videos(
                width=target_width,
                height=target_height,
                num_frames=target_num_frames,
                num_videos=num_videos,
                overrides=mm_options.get("video"),
            ),
        }


class OXRMultiModalProcessor(BaseMultiModalProcessor[OXRProcessingInfo]):
    def visual_indicators_to_visual_tokens(
        self,
        visual_indicators: list[int],
    ) -> list[int]:
        vocab_size = self.info.get_hf_config().visual_vocab_size
        indicator_by_id = {token_id: idx for idx, token_id in enumerate(OXR_INDICATOR_IDS)}
        return [
            vocab_size - len(OXR_INDICATOR_IDS) + indicator_by_id[token_id]
            for token_id in visual_indicators
            if token_id in indicator_by_id
        ]

    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        if not mm_data:
            tokenizer = self.info.get_tokenizer()
            prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        processed_outputs = super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
            tok_kwargs=tok_kwargs,
        )
        hf_processor = self.info.get_hf_processor()

        if "videos" in mm_data:
            visual_indicators = [
                hf_processor.construct_visual_indicators((1, 1, 1), True)
                for _ in processed_outputs["video_grids"]
            ]
            processed_outputs["video_indicator_tokens"] = torch.tensor(
                [
                    self.visual_indicators_to_visual_tokens(indicators)
                    for indicators in visual_indicators
                ],
                dtype=torch.long,
            )

        if "images" in mm_data:
            visual_indicators = [
                hf_processor.construct_visual_indicators((1, 1, 1), False)
                for _ in processed_outputs["grids"]
            ]
            processed_outputs["indicator_tokens"] = torch.tensor(
                [
                    self.visual_indicators_to_visual_tokens(indicators)
                    for indicators in visual_indicators
                ],
                dtype=torch.long,
            )

        return processed_outputs

    def _apply_hf_processor_tokens_only(self, prompt_tokens: list[int]) -> list[int]:
        return prompt_tokens

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return _oxr_field_config()

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargsItems,
    ) -> list[PromptReplacement]:
        def get_replacement(item_idx: int, modality: str) -> list[int]:
            if modality == "image":
                grid = out_mm_kwargs["image"][item_idx]["grids"].data
            elif modality == "video":
                grid = out_mm_kwargs["video"][item_idx]["video_grids"].data
            else:
                raise ValueError(f"Unsupported modality: {modality}")

            hf_processor = self.info.get_hf_processor()
            return hf_processor.construct_visual_placeholders(
                grid[0],
                is_video=(modality == "video"),
            )

        return [
            PromptReplacement(
                modality=modality,
                target=IMAGE_TOKEN if modality == "image" else VIDEO_TOKEN,
                replacement=partial(get_replacement, modality=modality),
            )
            for modality in ("image", "video")
        ]


@MULTIMODAL_REGISTRY.register_processor(
    OXRMultiModalProcessor,
    info=OXRProcessingInfo,
    dummy_inputs=OXRDummyInputsBuilder,
)
class OXR(nn.Module, SupportsMultiModal, SupportsPP):
    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return IMAGE_TOKEN
        if modality.startswith("video"):
            return VIDEO_TOKEN
        raise ValueError("Only image or video modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config

        self.config: PretrainedConfig = config

        with self._mark_language_model(vllm_config):
            self.llm = init_vllm_registered_model(
                vllm_config=vllm_config.with_hf_config(get_llm_config(config)),
                prefix=maybe_prefix(prefix, "llm"),
            )

        with self._mark_tower_model(vllm_config, {"image", "video"}):
            self.visual_tokenizer = OXRVisualTokenizer(
                config=config.vit_config,
                visual_vocab_size=config.visual_vocab_size,
                quant_config=quant_config,
                prefix=f"{prefix}.visual_tokenizer",
            )
            self.vte = OXREmbedding(config.visual_vocab_size, config.hidden_size)

        text_model_type = get_llm_config(config).model_type
        self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]

        self.make_empty_intermediate_tensors = (
            self.get_language_model().make_empty_intermediate_tensors
        )

    def _parse_and_validate_image_input(
        self,
        **kwargs: object,
    ) -> OXRImagePatchInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        indicator_tokens = kwargs.pop("indicator_tokens", None)
        grids = kwargs.pop("grids", None)
        if pixel_values is None and indicator_tokens is None:
            return None

        if pixel_values is not None and indicator_tokens is not None:
            if not isinstance(pixel_values, (torch.Tensor, list)):
                raise ValueError(f"Incorrect type of pixel values: {type(pixel_values)}")
            if not isinstance(indicator_tokens, (torch.Tensor, list)):
                raise ValueError(
                    f"Incorrect type of indicator_tokens: {type(indicator_tokens)}"
                )
            if grids is None:
                raise ValueError("Missing grids for image input")

            return OXRImagePatchInputs(
                type="image_patches",
                flat_data=flatten_bn(pixel_values, concat=True),
                patches_per_item=[
                    x.shape[0] // (self.config.vit_config.hidden_stride**2)
                    for x in pixel_values
                ],
                indicator_tokens=flatten_bn(indicator_tokens, concat=True),
                grids=flatten_bn(grids, concat=True),
            )

        raise AssertionError("This line should be unreachable.")

    def _parse_and_validate_video_input(
        self,
        **kwargs: object,
    ) -> OXRVideoPatchInputs | None:
        pixel_values = kwargs.pop("video_pixel_values", None)
        indicator_tokens = kwargs.pop("video_indicator_tokens", None)
        grids = kwargs.pop("video_grids", None)
        if pixel_values is None and indicator_tokens is None:
            return None

        if pixel_values is not None and indicator_tokens is not None:
            if not isinstance(pixel_values, (torch.Tensor, list)):
                raise ValueError(f"Incorrect type of pixel values: {type(pixel_values)}")
            if not isinstance(indicator_tokens, (torch.Tensor, list)):
                raise ValueError(
                    f"Incorrect type of indicator_tokens: {type(indicator_tokens)}"
                )
            if grids is None:
                raise ValueError("Missing grids for video input")

            return OXRVideoPatchInputs(
                type="video_patches",
                flat_data=flatten_bn(pixel_values, concat=True),
                patches_per_item=[
                    x.shape[0] // (self.config.vit_config.hidden_stride**2)
                    for x in pixel_values
                ],
                indicator_tokens=flatten_bn(indicator_tokens, concat=True),
                grids=flatten_bn(grids, concat=True),
            )

        raise AssertionError("This line should be unreachable.")

    def _process_visual_input(
        self,
        visual_input: OXRImagePatchInputs | OXRVideoPatchInputs,
    ) -> MultiModalEmbeddings:
        image_patches_flat = visual_input["flat_data"]
        patches_per_item = visual_input["patches_per_item"]
        indicator_tokens = visual_input["indicator_tokens"]
        grid_thws = visual_input["grids"]

        visual_tokens = self.visual_tokenizer(
            image_patches_flat.to(self.visual_tokenizer.dtype),
            grid_thws,
        )
        visual_embeds = self.vte(visual_tokens)
        indicator_embeds = self.vte(indicator_tokens)

        visual_embeds_per_item = visual_embeds.split(patches_per_item, dim=0)
        indicator_embeds_per_item = indicator_embeds.split(
            [2] * len(patches_per_item)
        )

        return tuple(
            torch.cat([indicator[:1], visual, indicator[1:]], dim=0)
            for indicator, visual in zip(
                indicator_embeds_per_item,
                visual_embeds_per_item,
            )
        )

    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
        modalities = {}
        for input_key in kwargs:
            if (
                input_key in ("pixel_values", "indicator_tokens", "grids")
                and "images" not in modalities
            ):
                modalities["images"] = self._parse_and_validate_image_input(**kwargs)
            if (
                input_key
                in ("video_pixel_values", "video_indicator_tokens", "video_grids")
                and "videos" not in modalities
            ):
                modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
        return modalities

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
        if not modalities:
            return []

        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
        for modality in modalities:
            if modality == "images":
                image_input = modalities["images"]
                multimodal_embeddings += tuple(self._process_visual_input(image_input))
            if modality == "videos":
                video_input = modalities["videos"]
                multimodal_embeddings += tuple(self._process_visual_input(video_input))
        return multimodal_embeddings

    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

        return self.llm(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )

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

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

    def get_language_model(self) -> torch.nn.Module:
        return self.llm