"""Direct prompt-only Qwen3-VL-32B conditioning for MiniMax H3.

The MiniMax checkpoint contains the first 50 Qwen3-VL decoder layers.  H3
uses their unnormalized final output, not the language-model head or final
RMSNorm.  This module intentionally has no ComfyUI imports and does not load
the Qwen vision encoder; image/video prompt construction remains a separate
feature.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import nn

from .checkpoint import H3Checkpoint


@dataclass(frozen=True)
class Qwen3VL32BTextConfig:
    vocab_size: int = 151936
    hidden_size: int = 5120
    intermediate_size: int = 25600
    num_layers: int = 50
    num_attention_heads: int = 64
    num_key_value_heads: int = 8
    head_dim: int = 128
    rms_norm_eps: float = 1e-6
    rope_theta: float = 5_000_000.0


class _RMSNorm(nn.Module):
    def __init__(self, weight: torch.Tensor, eps: float):
        super().__init__()
        self.eps = eps
        self.register_buffer("weight", weight, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.rms_norm(x, self.weight.shape, weight=self.weight.to(x), eps=self.eps)


def _apply_rope(query: torch.Tensor, key: torch.Tensor, freqs) -> tuple[torch.Tensor, torch.Tensor]:
    """Direct PyTorch port of Comfy's `apply_rope` fed `freqs=(cos, sin, neg_sin)`.

    ``query``/``key`` are ``[batch, heads, seq, head_dim]``; ``freqs`` broadcast over
    that shape. This is the exact split-half kernel the reference uses for both the
    plain text rope and the Qwen3-VL interleaved mrope.
    """
    cosine, sine, negative_sine = freqs
    split = query.shape[-1] // 2
    q = query * cosine
    q[..., :split].addcmul_(query[..., split:], negative_sine)
    q[..., split:].addcmul_(query[..., :split], sine)
    k = key * cosine
    k[..., :split].addcmul_(key[..., split:], negative_sine)
    k[..., split:].addcmul_(key[..., :split], sine)
    return q.to(query.dtype), k.to(key.dtype)


def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.Tensor, torch.Tensor]:
    """Direct PyTorch port of Comfy's text-only `precompute_freqs_cis` / `apply_rope`."""
    sequence, head_dim = query.shape[-2:]
    position_ids = torch.arange(sequence, device=query.device).unsqueeze(0)
    theta_numerator = torch.arange(0, head_dim, 2, device=query.device).float()
    inv_freq = 1.0 / (theta ** (theta_numerator / head_dim))
    frequencies = (inv_freq[None, :, None].expand(1, -1, 1).float() @ position_ids[:, None, :].float()).transpose(1, 2)
    embedding = torch.cat((frequencies, frequencies), dim=-1)
    cosine = embedding.cos().unsqueeze(1)
    sine = embedding.sin().unsqueeze(1)
    negative_sine = -sine[..., sine.shape[-1] // 2 :]
    sine = sine[..., : sine.shape[-1] // 2]
    return _apply_rope(query, key, (cosine, sine, negative_sine))


def _qwen_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Match Comfy's small-input SDPA wrapper for masked Qwen GQA."""
    kwargs = {"attn_mask": mask.unsqueeze(0).unsqueeze(0), "dropout_p": 0.0, "is_causal": False, "enable_gqa": True}
    if query.numel() >= 1024 * 128:
        from torch.nn.attention import SDPBackend, sdpa_kernel

        priority = [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
        params = torch.backends.cuda.SDPAParams(query, key, value, kwargs["attn_mask"], 0.0, False, True)
        supports_native_gqa = (
            torch.backends.cuda.can_use_flash_attention(params)
            or torch.backends.cuda.can_use_cudnn_attention(params)
            or torch.backends.cuda.can_use_efficient_attention(params)
        )
        if not supports_native_gqa:
            repeats = query.shape[-3] // key.shape[-3]
            key = key.repeat_interleave(repeats, dim=-3)
            value = value.repeat_interleave(repeats, dim=-3)
            kwargs["enable_gqa"] = False
        with sdpa_kernel(priority, set_priority=True):
            return F.scaled_dot_product_attention(query, key, value, **kwargs)
    return F.scaled_dot_product_attention(query, key, value, **kwargs)


class _Qwen3VLBlock(nn.Module):
    def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
        super().__init__()
        self.config = config
        self.input_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.input_layernorm.weight", dtype=dtype), config.rms_norm_eps)
        self.post_attention_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.post_attention_layernorm.weight", dtype=dtype), config.rms_norm_eps)
        self.q_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.q_norm.weight", dtype=dtype), config.rms_norm_eps)
        self.k_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.k_norm.weight", dtype=dtype), config.rms_norm_eps)
        self.q_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.q_proj", output_dtype=dtype)
        self.k_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.k_proj", output_dtype=dtype)
        self.v_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.v_proj", output_dtype=dtype)
        self.o_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.o_proj", output_dtype=dtype)
        self.gate_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.gate_proj", output_dtype=dtype)
        self.up_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.up_proj", output_dtype=dtype)
        self.down_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.down_proj", output_dtype=dtype)

    def forward(self, hidden_states: torch.Tensor, freqs_cis=None) -> torch.Tensor:
        residual = hidden_states
        x = self.input_layernorm(hidden_states)
        batch, sequence, _ = x.shape
        query = self.q_proj(x).view(batch, sequence, self.config.num_attention_heads, self.config.head_dim).transpose(1, 2)
        key = self.k_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
        value = self.v_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
        query = self.q_norm(query)
        key = self.k_norm(key)
        if freqs_cis is None:
            query, key = _rope(query, key, self.config.rope_theta)
        else:
            query, key = _apply_rope(query, key, freqs_cis)
        # Comfy selects its small-input SDPA path for Qwen, with an explicit causal mask.
        causal_mask = torch.full(
            (sequence, sequence),
            torch.finfo(query.dtype).min / 4,
            dtype=query.dtype,
            device=query.device,
        ).triu_(1)
        attention = _qwen_attention(query, key, value, causal_mask)
        hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
        residual = hidden_states
        x = self.post_attention_layernorm(hidden_states)
        return residual + self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Qwen3VL32BTextEncoder(nn.Module):
    """Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states."""
    config = Qwen3VL32BTextConfig()

    def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
        super().__init__()
        self.checkpoint_path = str(checkpoint_path)
        self.device_name = str(device)
        self.dtype = dtype
        checkpoint = H3Checkpoint(checkpoint_path, device=device)
        self._validate_checkpoint(checkpoint)
        self.register_buffer("embed_tokens", checkpoint.tensor("model.embed_tokens.weight"), persistent=False)
        self.register_buffer("embed_scale", checkpoint.tensor("model.embed_tokens.weight_scale", dtype=torch.float32), persistent=False)
        self.layers = nn.ModuleList(
            _Qwen3VLBlock(checkpoint, f"model.layers.{index}", self.config, dtype)
            for index in range(self.config.num_layers)
        )

    @classmethod
    def _validate_checkpoint(cls, checkpoint: H3Checkpoint) -> None:
        from safetensors import safe_open

        required = {"model.embed_tokens.weight"}
        linear_names = (
            "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
            "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
        )
        for index in range(cls.config.num_layers):
            prefix = f"model.layers.{index}"
            required.update({
                f"{prefix}.input_layernorm.weight",
                f"{prefix}.post_attention_layernorm.weight",
                f"{prefix}.self_attn.q_norm.weight",
                f"{prefix}.self_attn.k_norm.weight",
            })
            for linear in linear_names:
                required.update(
                    f"{prefix}.{linear}.{suffix}"
                    for suffix in ("comfy_quant", "weight", "weight_scale", "weight_scale_2")
                )
        with safe_open(checkpoint.path, framework="pt", device="cpu") as file:
            names = set(file.keys())
        missing = sorted(required - names)
        if missing:
            raise ValueError(
                "Not a supported MiniMax Qwen3-VL-32B NVFP4 checkpoint; missing "
                + ", ".join(missing)
            )

    def _embed_rows(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Scaled token embeds ``[batch, tokens, 5120]`` in fp32 (pre-decoder)."""
        token_rows = F.embedding(input_ids.to(self.embed_tokens.device), self.embed_tokens).to(torch.float32)
        token_scales = F.embedding(input_ids.to(self.embed_scale.device), self.embed_scale)
        return (token_rows * token_scales).to(torch.float32)

    def _run_layers(
        self,
        hidden_states: torch.Tensor,
        *,
        position_ids: torch.Tensor | None = None,
        visual_pos_masks: torch.Tensor | None = None,
        deepstack_embeds: list[torch.Tensor] | None = None,
    ) -> torch.Tensor:
        """Run all 50 decoder blocks over an fp32 ``[batch, tokens, 5120]`` state.

        With ``position_ids`` (``[3, seq]``) the Qwen3-VL interleaved mrope is used
        instead of the plain 1D text rope; with ``visual_pos_masks`` +
        ``deepstack_embeds`` the three DeepStack features are added at the visual
        positions of the first three decoder layers (Comfy ``Llama2_.forward``).
        """
        hidden_states = hidden_states.to(self.dtype)
        freqs_cis = None
        if position_ids is not None:
            from .qwen3vl_vision import mrope_freqs_cis

            freqs_cis = mrope_freqs_cis(position_ids, device=hidden_states.device)
        for layer_index, layer in enumerate(self.layers):
            hidden_states = layer(hidden_states, freqs_cis)
            if (
                deepstack_embeds is not None
                and visual_pos_masks is not None
                and layer_index < len(deepstack_embeds)
            ):
                mask = visual_pos_masks.to(hidden_states.device)
                hidden_states[mask] = hidden_states[mask] + deepstack_embeds[layer_index].to(hidden_states)
        return hidden_states

    @torch.inference_mode()
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Return unnormalized `[batch, tokens, 5120]` output after decoder layer 50."""
        if input_ids.ndim != 2:
            raise ValueError(f"input_ids must have shape [batch, tokens], got {tuple(input_ids.shape)}")
        if input_ids.numel() == 0:
            raise ValueError("input_ids must contain at least one token")
        return self._run_layers(self._embed_rows(input_ids))


class Qwen3VLPromptConditioner:
    """Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning."""
    def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
        from .conditioning import H3PromptTokenizer

        tokenizer_dir = tokenizer_dir or Path(__file__).with_name("qwen25_tokenizer")
        self.tokenizer = H3PromptTokenizer(tokenizer_dir)
        self.encoder = Qwen3VL32BTextEncoder(checkpoint_path, device=device, dtype=dtype)
        self.device = device

    def __call__(self, prompt: str) -> torch.Tensor:
        return self.encoder(self.tokenizer(prompt, device=self.device))