"""Vision-Exp wrapper around vLLM DeepseekV4ForCausalLM.

Loads official `vision.*` / `aligner.*` / image sentinel tensors and merges
ViT embeddings on prefill. Language decode stays on the B12X DeepSeek V4 path.
"""

from __future__ import annotations

from types import SimpleNamespace
from typing import Any

import torch
import torch.nn as nn

import torch
import torch.nn as nn

from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.models.deepseek_v4 import DeepseekV4ForCausalLM
from vllm.multimodal import MULTIMODAL_REGISTRY

from .processor import (
    DSV4VisionDummyInputsBuilder,
    DSV4VisionMultiModalProcessor,
    DSV4VisionProcessingInfo,
)
from .vision import Aligner, ViT


def _vision_args(hf: Any) -> SimpleNamespace:
    return SimpleNamespace(
        vision_n_layers=int(getattr(hf, "vision_n_layers", 0) or 0),
        vision_dim=int(getattr(hf, "vision_dim", 1024)),
        vision_n_heads=int(getattr(hf, "vision_n_heads", 16)),
        vision_inter_dim=int(getattr(hf, "vision_inter_dim", 2816)),
        vision_patch_size=int(getattr(hf, "vision_patch_size", 14)),
        vision_rope_theta=float(getattr(hf, "vision_rope_theta", 10000.0)),
        vision_downsample_ratio=int(getattr(hf, "vision_downsample_ratio", 3)),
        vision_max_n_token=int(getattr(hf, "vision_max_n_token", 384)),
        vision_min_pixels=int(getattr(hf, "vision_min_pixels", 147456)),
        vision_max_wh_ratio=int(getattr(hf, "vision_max_wh_ratio", 8)),
        dim=int(getattr(hf, "hidden_size", 4096)),
        num_hash_layers=int(getattr(hf, "num_hash_layers", 0) or 0),
    )


@MULTIMODAL_REGISTRY.register_processor(
    DSV4VisionMultiModalProcessor,
    info=DSV4VisionProcessingInfo,
    dummy_inputs=DSV4VisionDummyInputsBuilder,
)
class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
    # ModelConfig.is_multimodal_model keys off this class attribute
    # (interfaces.supports_multimodal) — without it the OpenAI layer
    # answers "not a multimodal model" for image_url requests.
    supports_multimodal = True

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        # chat_utils asks the resolved model class for the per-image prompt
        # placeholder; the tokenizer encodes it to <｜deepseek_image｜> (id 129264).
        if modality.startswith("image"):
            from .encoding_dsv4 import IMAGE_PLACEHOLDER

            return IMAGE_PLACEHOLDER
        return None

    def __init__(self, *, vllm_config, prefix: str = ""):
        super().__init__(vllm_config=vllm_config, prefix=prefix)
        hf = vllm_config.model_config.hf_config
        args = _vision_args(hf)
        self._vision_args = args
        if args.vision_n_layers <= 0:
            self.vision = None
            self.aligner = None
            return
        self.vision = ViT(args)
        self.aligner = Aligner(args)
        self.image_start = nn.Parameter(torch.empty(args.dim))
        self.image_end = nn.Parameter(torch.empty(args.dim))
        self.image_newline = nn.Parameter(torch.empty(args.dim))
        self.image_pad = nn.Parameter(torch.empty(args.dim))

    def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)

    def _merge_mm_embeddings(
        self,
        inputs_embeds: torch.Tensor,
        multimodal_embeddings,
        is_multimodal: torch.Tensor,
    ) -> torch.Tensor:
        # Mirrors interfaces._merge_multimodal_embeddings: fill masked rows in
        # order with the concatenated per-item embedding runs.
        flat = torch.cat([emb.reshape(-1) for emb in multimodal_embeddings])
        inputs_embeds.masked_scatter_(
            is_multimodal.to(inputs_embeds.device).unsqueeze(-1), flat
        )
        return inputs_embeds

    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings=None,
        *,
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if is_multimodal is not None:
            # Sentinel block ids may exceed vocab; zero them before lookup.
            input_ids = input_ids.masked_fill(
                is_multimodal.to(input_ids.device).unsqueeze(-1), 0
            )
        inputs_embeds = super().embed_input_ids(input_ids)
        if not multimodal_embeddings:
            return inputs_embeds
        return self._merge_mm_embeddings(
            inputs_embeds, multimodal_embeddings, is_multimodal
        )

    def embed_multimodal(self, **kwargs) -> tuple:
        patches = kwargs.get("image_patches")
        if patches is None or self.vision is None:
            return ()
        grids = kwargs["image_grids"].data
        types_all = kwargs["image_types"].data
        perm_all = kwargs["image_perm"].data
        t_counts = kwargs["image_type_counts"].data.tolist() if hasattr(kwargs["image_type_counts"], "tolist") else kwargs["image_type_counts"].data.tolist()
        perm_counts = kwargs["image_perm_counts"].data.tolist()
        n_items = int(grids.shape[0])

        out: list[torch.Tensor] = []
        p_off = t_off = q_off = 0
        start = self.image_start
        pad = self.image_pad
        newline = self.image_newline
        end = self.image_end
        for i in range(n_items):
            n_vit_h, n_vit_w, n_llm_h, n_llm_w = grids[i].tolist()
            p_count = n_vit_h * n_vit_w
            t_count = int(t_counts[i])
            q_count = int(perm_counts[i])
            item_types = types_all[t_off : t_off + t_count].tolist()
            item_perm = perm_all[q_off : q_off + q_count].to(torch.long)
            if q_count != sum(t == 2 for t in item_types):
                raise ValueError(
                    "DSV4 vision perm/sentinel layout mismatch "
                    f"({q_count} != {sum(t == 2 for t in item_types)})"
                )
            t_off += t_count
            q_off += q_count
            vit = self.encode_image(patches[p_off : p_off + p_count], n_vit_h, n_vit_w)
            p_off += p_count
            rows: list[torch.Tensor] = []
            vi = 0
            for t in item_types:
                if t == 0:  # IMAGE_START
                    rows.append(start)
                elif t == 1:  # IMAGE_PAD
                    rows.append(pad)
                elif t == 2:  # IMAGE
                    rows.append(vit[int(item_perm[vi])])
                    vi += 1
                elif t == 3:  # IMAGE_NEW_LINE
                    rows.append(newline)
                else:  # IMAGE_END
                    rows.append(end)
            out.append(torch.stack(rows))

        return tuple(out)

    def forward(
        self,
        input_ids,
        positions,
        intermediate_tensors=None,
        inputs_embeds=None,
        **kwargs,
    ):
        # The runner embeds via embed_input_ids (mm merge) and hands us
        # merged inputs_embeds; lm_head and get_mtp_target_hidden_states
        # stay on the parent class.
        return super().forward(
            input_ids,
            positions,
            intermediate_tensors,
            inputs_embeds,
        )

    def load_weights(self, weights):
        if self.vision is None:
            return super().load_weights(weights)
        leftover = []
        stacked = {
            "vision.": self.vision,
            "aligner.": self.aligner,
        }
        params = {
            "image_start": self.image_start,
            "image_end": self.image_end,
            "image_newline": self.image_newline,
            "image_pad": self.image_pad,
        }
        for name, tensor in weights:
            if name in params:
                params[name].data.copy_(tensor)
                continue
            routed = False
            for prefix, module in stacked.items():
                if name.startswith(prefix):
                    AutoWeightsLoader(module).load_weights(
                        [(name[len(prefix) :], tensor)]
                    )
                    routed = True
                    break
            if not routed:
                leftover.append((name, tensor))
        leftover = [
            pair
            for pair in leftover
            if not _skip_unmapped_gate(pair[0], self._vision_args)
        ]
        return super().load_weights(leftover)


def _skip_unmapped_gate(name: str, args: SimpleNamespace) -> bool:
    # Hash MoE layers register tid2eid and leave e_score_correction_bias as
    # None. The B12X mapper still rewrites ffn.gate.bias onto that missing
    # name. bias_vl has no mapper at all.
    if name.endswith("ffn.gate.bias_vl"):
        return True
    if name.endswith("ffn.gate.bias"):
        parts = name.split(".")
        if (
            len(parts) >= 5
            and parts[0] == "layers"
            and parts[1].isdigit()
            and int(parts[1]) < args.num_hash_layers
        ):
            return True
    return False