"""vLLM FastMTP plugin for Jina-OCR-v1.

Call ``register()`` once before constructing ``vllm.LLM``. This file lives in
the Hugging Face snapshot so serving does not depend on any private training repo.
"""

from __future__ import annotations

import copy
import os
import sys
from collections.abc import Iterable
from pathlib import Path

import torch
import torch.nn as nn
from typing_extensions import override
from vllm import ModelRegistry
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
from vllm.config.speculative import SpeculativeConfig
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
from vllm.model_executor.models.deepseek_mtp import (
    DeepSeekMTP,
    DeepSeekMultiTokenPredictor,
    DeepSeekMultiTokenPredictorLayer,
    SharedHead,
)
from vllm.model_executor.models.deepseek_ocr import DeepseekOCRForCausalLM
from vllm.model_executor.models.deepseek_v2 import DeepseekV2DecoderLayer, get_spec_layer_idx_from_weight_name
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
from vllm.platforms import current_platform

IMAGE_TOKEN_INDEX = 128815
_MODULE_QUALNAME = "deepseek_ocr_mtp"
_registered = False

_STACKED_PARAMS = [
    ("gate_up_proj", "gate_proj", 0),
    ("gate_up_proj", "up_proj", 1),
    ("qkv_proj", "q_proj", "q"),
    ("qkv_proj", "k_proj", "k"),
    ("qkv_proj", "v_proj", "v"),
]


class DeepSeekOCRMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):
    """MTP layer with dense MLP (no MoE). ``n_routed_experts=None`` forces MLP."""

    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
        # Skip parent __init__ to intercept decoder layer creation
        nn.Module.__init__(self)

        config = vllm_config.speculative_config.draft_model_config.hf_config
        self.config = config
        quant_config = vllm_config.quant_config

        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
        self.device = current_platform.device_type

        # V3.2 topk handling (our model doesn't have index_topk)
        self.is_v32 = hasattr(config, "index_topk")

        self.shared_head = SharedHead(config=config, prefix=prefix, quant_config=quant_config)

        # Patch config for MTP decoder block: force dense MLP (no MoE)
        mtp_config = copy.deepcopy(config)
        mtp_config.n_routed_experts = None

        self.mtp_block = DeepseekV2DecoderLayer(
            vllm_config,
            prefix,
            config=mtp_config,
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
        spec_step_index: int = 0,
    ) -> torch.Tensor:
        """Return post-norm hidden states so FastMTP step k+1 matches training."""
        assert inputs_embeds is not None
        # masking inputs at position 0, as not needed by MTP
        inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
        inputs_embeds = self.enorm(inputs_embeds)
        previous_hidden_states = self.hnorm(previous_hidden_states)

        hidden_states = self.eh_proj(torch.cat([inputs_embeds, previous_hidden_states], dim=-1))

        hidden_states, residual = self.mtp_block(positions=positions, hidden_states=hidden_states, residual=None)
        hidden_states = residual + hidden_states
        return self.shared_head(hidden_states)


class DeepSeekOCRMultiTokenPredictor(DeepSeekMultiTokenPredictor):
    """Single shared dense MTP head; K draft steps reuse the same weights."""

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        # Skip parent __init__ to use our custom layer class
        nn.Module.__init__(self)

        config = vllm_config.model_config.hf_config
        self.mtp_start_layer_idx = config.num_hidden_layers
        # FastMTP recursive mode requires a single shared head; the K draft
        # steps are driven by speculative_config.num_speculative_tokens.
        self.recursive = getattr(config, "mtp_recursive", False)
        num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1)
        if self.recursive:
            assert num_mtp_layers == 1, (
                "FastMTP recursive mode (mtp_recursive=True) requires "
                "num_nextn_predict_layers=1 (single shared head); got "
                f"{num_mtp_layers}. Use num_speculative_tokens=K for draft depth."
            )
        self.num_mtp_layers = num_mtp_layers

        self.layers = torch.nn.ModuleDict(
            {
                str(idx): DeepSeekOCRMultiTokenPredictorLayer(vllm_config, f"{prefix}.layers.{idx}")
                for idx in range(
                    self.mtp_start_layer_idx,
                    self.mtp_start_layer_idx + self.num_mtp_layers,
                )
            }
        )
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            prefix=maybe_prefix(prefix, "embed_tokens"),
        )
        self.logits_processor = LogitsProcessor(config.vocab_size)

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
        """Skip shared_head.norm — layer forward already returned post-norm states."""
        current_step_idx = spec_step_idx % self.num_mtp_layers
        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
        logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
        return logits


# This file postpones annotations, so vLLM cannot infer dims from live
# types. Pass them explicitly or class import raises.
@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": 0,
        "hidden_states": 0,
        "intermediate_tensors": 0,
        "inputs_embeds": 0,
    }
)
class DeepSeekOCRMTP(DeepSeekMTP):
    """Dense-MLP FastMTP draft model.

    vLLM's ``method="mtp"`` proposer re-grounds every draft step on the target
    hidden state. FastMTP was trained with recursive feedback, so callers must
    set ``speculative_config.method="eagle"`` and this ``forward`` returns
    ``(output, output)`` for (logits, next-step hidden state).
    """

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        # Bypass DeepSeekMTP.__init__: it builds a MoE predictor whose
        # attention names collide with ours. The compile decorator sets
        # vllm_config / do_not_compile after this body returns.
        nn.Module.__init__(self)
        self.config = vllm_config.model_config.hf_config
        self.model = DeepSeekOCRMultiTokenPredictor(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
        self.set_moe_parameters()

    def set_moe_parameters(self) -> None:
        """Override: MTP blocks use dense MLP; skip MoE loop and n_group access.
        Pre-initialise the list attributes the parent sets before its loop, then
        delegate to extract_moe_parameters(None) for the num_* scalar fields.
        """
        self.expert_weights = []
        self.moe_layers = []
        self.moe_mlp_layers = []
        self.extract_moe_parameters(None)

    @override
    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        intermediate_tensors=None,
        inputs_embeds: torch.Tensor | None = None,
        spec_step_idx: int = 0,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        output = self.model(input_ids, positions, hidden_states, inputs_embeds, spec_step_idx)
        if self.model.recursive:
            return output, output
        return output

    def _build_weights_mapper(self) -> WeightsMapper:
        """Build a WeightsMapper that translates checkpoint names to vLLM's format.

        Insertion order matters — substr rules are applied in sequence by _map_name,
        so more-specific patterns (mtp_block., shared_head._norm) must come before
        the catch-all mtp_module.heads.{i}. entry for each head.

        Note: _rewrite_spec_layer_name() runs AFTER this mapper and re-inserts
        .mtp_block. for transformer-block weights, so we strip it here.
        """
        N = self.config.num_hidden_layers
        num_mtp = self.model.num_mtp_layers

        # Self-contained heads carry their own lm_head / norm / input embedding
        # (mtp_share_*=False, the newer format). Shared heads (older format) omit
        # those tensors and reuse the main model's lm_head / model.norm /
        # model.embed_tokens. Default to shared for backward compatibility.
        share_lm_head = getattr(self.config, "mtp_share_lm_head", True)
        share_norm = getattr(self.config, "mtp_share_norm", True)
        share_embed = getattr(self.config, "mtp_share_embedding_weights", True)

        substr: dict[str, str | None] = {}
        for i in range(num_mtp):
            L = N + i
            # Strip mtp_block. prefix (_rewrite_spec_layer_name adds it back)
            substr[f"mtp_module.heads.{i}.mtp_block."] = f"model.layers.{L}."
            # Rename _norm/_head attrs (must precede the catch-all below)
            substr[f"mtp_module.heads.{i}.shared_head._norm"] = f"model.layers.{L}.shared_head.norm"
            substr[f"mtp_module.heads.{i}.shared_head._head"] = f"model.layers.{L}.shared_head.head"
            # Self-contained head: its own lm_head is stored as shared_head.local_head
            # (must precede the catch-all below). shared_head.norm falls through to
            # the catch-all unchanged. Absent in the shared format -> harmless.
            substr[f"mtp_module.heads.{i}.shared_head.local_head"] = f"model.layers.{L}.shared_head.head"
            # Catch-all: enorm, hnorm, eh_proj, shared_head.norm, and anything else
            substr[f"mtp_module.heads.{i}."] = f"model.layers.{L}."

        # Shared weights — 1:1 while num_mtp==1; first head gets all shared params
        # (_rewrite_spec_layer_name routes embed_tokens back to model.embed_tokens.weight).
        # Only wire the main model's tensors into the head when the head shares them;
        # otherwise the head's own local_head / shared_head.norm / mtp_embed_tokens win.
        if share_lm_head:
            substr["lm_head.weight"] = f"model.layers.{N}.shared_head.head.weight"
        if share_norm:
            substr["model.norm.weight"] = f"model.layers.{N}.shared_head.norm.weight"
        if share_embed:
            substr["model.embed_tokens.weight"] = f"model.layers.{N}.embed_tokens.weight"
        else:
            substr["mtp_embed_tokens.weight"] = f"model.layers.{N}.embed_tokens.weight"

        return WeightsMapper(
            orig_to_new_substr=substr,
            # Drop mtp_module.shared_head.* (duplicates main model's norm/lm_head)
            orig_to_new_prefix={"mtp_module.shared_head.": None},
        )

    @override
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        """Load weights with name remapping from our checkpoint format.

        Cannot call super().load_weights() — parent raises if any MTP layer index
        is missing from loaded_params. Remap checkpoint names, fuse QKV/gate_up,
        then mark remaining params loaded so the tracker does not fail.
        """
        mapper = self._build_weights_mapper()
        params_dict = dict(self.named_parameters())
        loaded: set[str] = set()
        remaining: list[tuple[str, torch.Tensor]] = []

        for name, w in mapper.apply(weights):
            if "rotary_emb.inv_freq" in name:
                continue
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is None:
                continue
            name = self._rewrite_spec_layer_name(spec_layer, name)

            for fused, shard, shard_id in _STACKED_PARAMS:
                if shard not in name:
                    continue
                mapped = name.replace(shard, fused)
                if mapped not in params_dict:
                    continue
                params_dict[mapped].weight_loader(params_dict[mapped], w, shard_id)
                loaded.add(mapped)
                break
            else:
                remaining.append((name, w))

        loaded |= AutoWeightsLoader(self).load_weights(iter(remaining))
        # MTP block params (enorm, hnorm, eh_proj, layernorms, etc.) are at default
        # init — mark all as handled so vLLM's weight tracker doesn't raise ValueError.
        loaded.update(n for n, _ in self.named_parameters())
        return loaded


class DeepseekOCRForCausalLMOCR(DeepseekOCRForCausalLM):
    """Main model: drop MTP-only tensors so AutoWeightsLoader does not reject them."""

    _MTP_PREFIXES = ("mtp_module.", "mtp_embed_tokens")

    @override
    def load_weights(self, weights):
        def _drop_mtp(it):
            for name, tensor in it:
                if not name.startswith(self._MTP_PREFIXES):
                    yield name, tensor

        return super().load_weights(_drop_mtp(weights))


_OURS = {
    "DeepseekOCRForCausalLM",
    "DeepseekOCRForCausalLMOCR",
    "DeepSeekMTPModel",
    "EagleDeepSeekMTPModel",
}

_NESTED_TEXT_CONFIG_KEYS = ("text_config", "language_config")
_NESTED_TEXT_SKIP = frozenset((*_NESTED_TEXT_CONFIG_KEYS, "vision_config", "projector_config", "auto_map"))


def _flatten_nested_text_config(hf_config) -> None:
    """Promote decoder fields; drop nested text configs.

    EAGLEConfig copies ``to_dict()`` onto itself, turning nested
    text/language configs into dicts that vLLM then rejects.
    """
    for name in _NESTED_TEXT_CONFIG_KEYS:
        nested = getattr(hf_config, name, None)
        if nested is None or nested is hf_config:
            continue
        src = nested.to_dict() if hasattr(nested, "to_dict") else nested
        if isinstance(src, dict):
            for key, value in src.items():
                if key in _NESTED_TEXT_SKIP:
                    continue
                if getattr(hf_config, key, None) is None:
                    setattr(hf_config, key, value)
        try:
            delattr(hf_config, name)
        except Exception:
            setattr(hf_config, name, None)


def _ensure_importable() -> None:
    """Put this snapshot directory on sys.path / PYTHONPATH for vLLM workers."""
    root = str(Path(__file__).resolve().parent)
    if root not in sys.path:
        sys.path.insert(0, root)
    parts = [p for p in os.environ.get("PYTHONPATH", "").split(os.pathsep) if p]
    if root not in parts:
        os.environ["PYTHONPATH"] = os.pathsep.join([root, *parts])


def register() -> None:
    """Register FastMTP architectures with vLLM. Call once before ``LLM(...)``.

    Registers ``DeepseekOCRForCausalLMOCR``, ``DeepSeekMTPModel``, and
    ``EagleDeepSeekMTPModel`` (vLLM ≥ 0.21 prefixes ``Eagle`` when
    ``speculative_config.method="eagle"``). Safe to call more than once.
    """
    global _registered
    _ensure_importable()
    ModelRegistry.register_model(
        "DeepseekOCRForCausalLMOCR",
        f"{_MODULE_QUALNAME}:DeepseekOCRForCausalLMOCR",
    )
    ModelRegistry.register_model(
        "DeepSeekMTPModel",
        f"{_MODULE_QUALNAME}:DeepSeekOCRMTP",
    )
    ModelRegistry.register_model(
        "EagleDeepSeekMTPModel",
        f"{_MODULE_QUALNAME}:DeepSeekOCRMTP",
    )
    if _registered:
        return

    _orig_hf_config_override = SpeculativeConfig.hf_config_override

    @staticmethod
    def _mtp_hf_config_override(hf_config):
        archs = getattr(hf_config, "architectures", None) or []
        if not any(a in _OURS for a in archs):
            return _orig_hf_config_override(hf_config)
        hf_config.model_type = "deepseek_mtp"
        result = _orig_hf_config_override(hf_config)
        result.architectures = ["DeepSeekMTPModel"]
        _flatten_nested_text_config(result)
        if result is not hf_config:
            hf_config.architectures = ["DeepSeekMTPModel"]
            _flatten_nested_text_config(hf_config)
        return result

    SpeculativeConfig.hf_config_override = _mtp_hf_config_override
    _registered = True


def _default_repetition_detection():
    """Built-in n-gram loop stopper; works with FastMTP (custom logits processors do not)."""
    try:
        from vllm.sampling_params import RepetitionDetectionParams
    except ImportError:
        return None
    return RepetitionDetectionParams(
        max_pattern_size=35,
        min_pattern_size=35,
        min_count=10,
    )


def vllm_sampling_params(**kwargs):
    """``SamplingParams`` with OCR defaults. Pass ``repetition_detection=None`` to disable."""
    from vllm import SamplingParams

    kwargs.setdefault("temperature", 0.0)
    kwargs.setdefault("repetition_penalty", 1.05)
    if "repetition_detection" not in kwargs:
        detection = _default_repetition_detection()
        if detection is not None:
            kwargs["repetition_detection"] = detection
    try:
        return SamplingParams(**kwargs)
    except TypeError as exc:
        if "repetition_detection" in kwargs and "repetition_detection" in str(exc):
            kwargs.pop("repetition_detection")
            return SamplingParams(**kwargs)
        raise


def vllm_llm_kwargs(
    model: str,
    *,
    num_speculative_tokens: int = 3,
    mtp_heads: int = 1,
    mtp_recursive: bool = True,
    tensor_parallel_size: int = 1,
    dtype: str = "bfloat16",
    **extra,
) -> dict:
    """Keyword arguments for ``vllm.LLM`` after ``register()``.

    FastMTP (``mtp_recursive=True``) sets ``speculative_config.method="eagle"``
    so each draft step consumes the previous step's hidden state.
    Use ``vllm_sampling_params()`` for n-gram repetition detection; vLLM rejects
    custom logits processors while speculative decoding is enabled.
    """
    kwargs = {
        "model": model,
        "trust_remote_code": True,
        "dtype": dtype,
        "tensor_parallel_size": tensor_parallel_size,
        "hf_overrides": {
            "architectures": ["DeepseekOCRForCausalLMOCR"],
            "num_nextn_predict_layers": mtp_heads,
            "mtp_recursive": mtp_recursive,
            "image_token_index": IMAGE_TOKEN_INDEX,
        },
    }
    if num_speculative_tokens > 0:
        kwargs["disable_log_stats"] = False
        spec = {
            "model": model,
            "num_speculative_tokens": num_speculative_tokens,
        }
        if mtp_recursive:
            spec["method"] = "eagle"
        kwargs["speculative_config"] = spec
    kwargs.update(extra)
    return kwargs


def log_spec_stats(llm) -> None:
    """Flush vLLM's ``SpecDecoding metrics`` line after a short generate.

    The engine only prints acceptance on its stats interval (~10s). One-shot
    ``LLM.generate`` exits before that unless this is called.
    """
    engine = getattr(llm, "llm_engine", None) or getattr(llm, "engine", None)
    do_log = getattr(engine, "do_log_stats", None)
    if do_log is None:
        return
    try:
        do_log()
    except Exception as exc:
        print(f"[WARN] do_log_stats() failed: {exc}", flush=True)


DEFAULT_OCR_PROMPT = (
    "Transcribe the provided document image into a clean Markdown format, preserving the natural reading order."
)


def prepare_vllm_input(tokenizer, image, prompt: str | None = None) -> dict:
    """Build one ``llm.generate`` item: chat-templated prompt plus pixels."""
    prompt = DEFAULT_OCR_PROMPT if prompt is None else prompt
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    return {"prompt": text, "multi_modal_data": {"image": image}}