"""Layer-split forward for Qwen2.5-VL and Qwen3-VL.

Method §3.2 needs ``F = F_{>l*} o F_{<=l*}`` as two separately runnable halves so
that the multimodal branch can be cut off at ``l*`` and replaced by the
workspace. §3.1 needs the same split to read and patch activations at a chosen
depth.

This module reimplements the prologue of the native Qwen VL text-model forward
-- embedding merge, M-RoPE index, causal mask, rotary embeddings -- as a
reusable :class:`SplitContext`, then exposes the decoder-layer loop as a range
you can run piecewise.  Qwen3-VL additionally injects three ``DeepStack``
vision features after language layers 0--2; those tensors are carried in the
split context and applied at the identical layer boundaries.  It deliberately
mirrors transformers 4.57.6 rather than monkeypatching it;
``tests/test_split_equivalence.py`` asserts the composed halves reproduce the
stock forward bit-for-bit, which is what makes the mirroring safe to rely on.

Shapes use ``B`` batch, ``L`` sequence, ``d`` backbone width (3584 on the 7B),
``N_v`` visual tokens, ``N_q`` question tokens.
"""

from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Any

import torch
from transformers.cache_utils import Cache
from transformers.masking_utils import (
    create_causal_mask,
    create_sliding_window_causal_mask,
)


@dataclass
class SplitContext:
    """Per-forward state shared by every decoder layer.

    Computed once by :func:`make_split_context` so that layer ranges can be run
    independently without recomputing masks or rotary tables.
    """

    hidden_states: torch.Tensor  # [B, L, d] - mutated as layers run
    position_ids: torch.Tensor  # [3, B, L] - M-RoPE (t, h, w)
    position_embeddings: tuple[torch.Tensor, torch.Tensor]  # (cos, sin) [B, L, head_dim]
    causal_mask_mapping: dict[str, torch.Tensor | None]
    cache_position: torch.Tensor  # [L]
    text_position_ids: torch.Tensor | None  # [B, L] only when packed
    past_key_values: Cache | None
    # Original 2-D key-padding mask.  ``create_causal_mask`` is allowed to
    # return ``None`` for SDPA and delegate causality to ``is_causal``; keeping
    # this tensor lets counterfactual branches materialize the equivalent mask
    # before removing a precisely selected set of attention edges.
    attention_mask: torch.Tensor | None = None  # [B, L_kv]
    # Qwen3-VL only.  DeepStack adds one visual feature tensor after each of
    # the first three language layers.  They stay ``None`` for Qwen2.5-VL and
    # for every text-only branch.
    visual_pos_masks: torch.Tensor | None = None  # [B, L] bool
    deepstack_visual_embeds: list[torch.Tensor] | None = None

    def clone_at(self, hidden_states: torch.Tensor) -> "SplitContext":
        """Same context, different hidden states (for patched re-runs)."""
        return SplitContext(
            hidden_states=hidden_states,
            position_ids=self.position_ids,
            position_embeddings=self.position_embeddings,
            causal_mask_mapping=self.causal_mask_mapping,
            cache_position=self.cache_position,
            text_position_ids=self.text_position_ids,
            past_key_values=self.past_key_values,
            attention_mask=self.attention_mask,
            visual_pos_masks=self.visual_pos_masks,
            deepstack_visual_embeds=self.deepstack_visual_embeds,
        )


# ---------------------------------------------------------------------------
# embedding / position construction
# ---------------------------------------------------------------------------


def embed_multimodal(
    vl_model,
    input_ids: torch.LongTensor,  # [B, L]
    pixel_values: torch.Tensor | None = None,
    image_grid_thw: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    *,
    return_deepstack: bool = False,
) -> (
    tuple[torch.Tensor, torch.Tensor]
    | tuple[
        torch.Tensor,
        torch.Tensor,
        torch.Tensor | None,
        list[torch.Tensor] | None,
    ]
):
    """Token embeddings with image features scattered in, plus M-RoPE indices.

    Mirrors the prefill path of ``Qwen2_5_VLModel.forward``. Pass
    ``pixel_values=None`` to get the text-only branch used for ``Q*``.

    Args:
        vl_model: a native ``Qwen2_5_VLModel`` or ``Qwen3VLModel`` (i.e.
            ``model.model``, not the ``...ForConditionalGeneration`` wrapper).
        return_deepstack: also return Qwen3-VL's visual-position mask and
            DeepStack features.  The default two-value return keeps all
            Qwen2.5 callers backward compatible.

    Returns:
        ``(inputs_embeds [B, L, d], position_ids [3, B, L])`` and optionally
        ``(visual_pos_masks, deepstack_visual_embeds)``.
    """
    inputs_embeds = vl_model.get_input_embeddings()(input_ids)  # [B, L, d]
    model_type = str(getattr(vl_model.config, "model_type", ""))
    # A freshly wrapped model exposes the native qwen3_vl config here.
    # Reloading a fully saved CLOSE checkpoint reconstructs the nested native
    # backbone from the wrapper config, so its model view carries
    # close_qwen3_vl instead. Both use the tuple-returning Qwen3 feature API.
    is_qwen3_vl = model_type in {"qwen3_vl", "close_qwen3_vl"}
    visual_pos_masks = None
    deepstack_visual_embeds = None

    if pixel_values is not None:
        image_features = vl_model.get_image_features(pixel_values, image_grid_thw)
        if is_qwen3_vl:
            image_embeds, deepstack_visual_embeds = image_features
        else:
            image_embeds = image_features
        image_embeds = torch.cat(image_embeds, dim=0).to(
            inputs_embeds.device, inputs_embeds.dtype
        )  # [N_v_total, d]
        image_mask, _ = vl_model.get_placeholder_mask(
            input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
        )
        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
        if is_qwen3_vl:
            visual_pos_masks = image_mask[..., 0]

    if is_qwen3_vl:
        position_ids, _ = vl_model.get_rope_index(
            input_ids,
            image_grid_thw,
            None,  # video_grid_thw
            attention_mask=attention_mask,
        )
    else:
        position_ids, _ = vl_model.get_rope_index(
            input_ids,
            image_grid_thw,
            None,  # video_grid_thw
            second_per_grid_ts=None,
            attention_mask=attention_mask,
        )
    if return_deepstack:
        return (
            inputs_embeds,
            position_ids,
            visual_pos_masks,
            deepstack_visual_embeds,
        )
    return inputs_embeds, position_ids


def make_split_context(
    text_model,
    inputs_embeds: torch.Tensor,  # [B, L, d]
    position_ids: torch.Tensor,  # [3, B, L]
    attention_mask: torch.Tensor | None = None,
    past_key_values: Cache | None = None,
    cache_position: torch.Tensor | None = None,
    visual_pos_masks: torch.Tensor | None = None,
    deepstack_visual_embeds: list[torch.Tensor] | None = None,
) -> SplitContext:
    """Build masks and rotary embeddings once, as the stock forward does.

    Args:
        text_model: ``Qwen2_5_VLTextModel`` (``vl_model.language_model``).
    """
    if cache_position is None:
        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
        cache_position = torch.arange(
            past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device
        )  # [L]

    if position_ids.ndim == 2:
        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)

    model_type = str(getattr(text_model.config, "model_type", ""))

    # Packed-sequence convention: a leading text-only row makes it [4, B, L].
    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
        text_position_ids = position_ids[0]  # [B, L]
        position_ids = position_ids[1:]  # [3, B, L]
    elif model_type == "qwen3_vl_text":
        # Qwen3-VL always passes the temporal M-RoPE row to both the causal-mask
        # builder and decoder layers, even for ordinary (non-packed) inputs.
        text_position_ids = position_ids[0]
    else:
        text_position_ids = None

    mask_kwargs: dict[str, Any] = {
        "config": text_model.config,
        "input_embeds": inputs_embeds,
        "attention_mask": attention_mask,
        "cache_position": cache_position,
        "past_key_values": past_key_values,
        "position_ids": text_position_ids,
    }
    causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
    if getattr(text_model, "has_sliding_layers", False):
        causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
            **mask_kwargs
        )

    position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids)

    return SplitContext(
        hidden_states=inputs_embeds,
        position_ids=position_ids,
        position_embeddings=position_embeddings,
        causal_mask_mapping=causal_mask_mapping,
        cache_position=cache_position,
        text_position_ids=text_position_ids,
        past_key_values=past_key_values,
        attention_mask=attention_mask,
        visual_pos_masks=visual_pos_masks,
        deepstack_visual_embeds=deepstack_visual_embeds,
    )


def block_attention_edges(
    ctx: SplitContext,
    query_mask: torch.Tensor,
    key_mask: torch.Tensor,
) -> SplitContext:
    """Return ``ctx`` with selected query-to-key attention edges removed.

    ``query_mask`` and ``key_mask`` are boolean ``[B, L]`` supports in the
    current (cache-free) sequence.  Every ordinary causal/padding constraint is
    preserved; only their Cartesian product is additionally masked.  Both
    boolean SDPA masks (``True`` means visible) and additive eager masks
    (``0``/negative infinity) are supported.

    The helper deliberately rejects cached contexts.  Its intended use is a
    counterfactual recurrent layer evaluation, never autoregressive decoding,
    and silently guessing the key offset of a populated cache would invalidate
    the causal comparison.
    """
    if ctx.past_key_values is not None and ctx.past_key_values.get_seq_length() > 0:
        raise ValueError("block_attention_edges requires a cache-free context")
    if query_mask.dtype != torch.bool or key_mask.dtype != torch.bool:
        raise TypeError("query_mask and key_mask must be boolean tensors")
    if query_mask.shape != key_mask.shape or query_mask.ndim != 2:
        raise ValueError(
            "query_mask and key_mask must have the same [B, L] shape, got "
            f"{tuple(query_mask.shape)} and {tuple(key_mask.shape)}"
        )
    batch_size, seq_len = query_mask.shape
    if ctx.hidden_states.shape[:2] != (batch_size, seq_len):
        raise ValueError(
            "edge masks must match the SplitContext sequence, got "
            f"{tuple(query_mask.shape)} for {tuple(ctx.hidden_states.shape[:2])}"
        )

    blocked = query_mask[:, None, :, None] & key_mask[:, None, None, :]
    updated: dict[str, torch.Tensor] = {}
    for attention_type, base_mask in ctx.causal_mask_mapping.items():
        if base_mask is None:
            # SDPA may omit an all-valid causal mask.  Materialize exactly that
            # lower triangle, then reapply key padding before deleting edges.
            q_positions = ctx.cache_position
            if q_positions.numel() != seq_len:
                raise ValueError(
                    "cache-free context must have one cache position per row"
                )
            key_positions = torch.arange(seq_len, device=query_mask.device)
            visible = key_positions[None, :] <= q_positions[:, None]
            visible = visible[None, None].expand(batch_size, 1, -1, -1)
            if ctx.attention_mask is not None:
                if ctx.attention_mask.shape != (batch_size, seq_len):
                    raise ValueError(
                        "counterfactual edge masking expects a 2-D [B, L] "
                        "attention mask"
                    )
                visible = visible & ctx.attention_mask[:, None, None, :].bool()
            updated[attention_type] = visible & ~blocked
            continue

        if not isinstance(base_mask, torch.Tensor) or base_mask.ndim != 4:
            raise TypeError(
                "counterfactual edge masking supports tensor 4-D attention "
                f"masks, got {type(base_mask)!r}"
            )
        if base_mask.shape[0] not in (1, batch_size):
            raise ValueError("attention-mask batch dimension is incompatible")
        if base_mask.shape[-2:] != (seq_len, seq_len):
            raise ValueError(
                "counterfactual edge masking expects a square cache-free mask, "
                f"got {tuple(base_mask.shape)}"
            )
        if base_mask.dtype == torch.bool:
            updated[attention_type] = base_mask & ~blocked
        elif base_mask.is_floating_point():
            updated[attention_type] = base_mask.masked_fill(
                blocked, torch.finfo(base_mask.dtype).min
            )
        else:
            raise TypeError(
                f"unsupported attention mask dtype {base_mask.dtype}"
            )

    return replace(ctx, causal_mask_mapping=updated)


# ---------------------------------------------------------------------------
# running layer ranges
# ---------------------------------------------------------------------------


def run_layer_range(
    text_model,
    ctx: SplitContext,
    start: int,
    stop: int | None = None,
    use_cache: bool = False,
    hidden_states: torch.Tensor | None = None,
    collect: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
    """Run ``text_model.layers[start:stop]`` on ``ctx``.

    ``self.norm`` is *not* applied -- it belongs to the very top of the stack.
    Call :func:`final_norm` after the last range.

    Args:
        hidden_states: override the context's states (leave ``None`` to chain).
        collect: also return the input hidden states of every layer in the range
            plus the range output, i.e. ``stop - start + 1`` tensors.

    Returns:
        ``[B, L, d]``, or ``(output, collected)`` when ``collect``.
    """
    layers = text_model.layers
    stop = len(layers) if stop is None else stop
    h = ctx.hidden_states if hidden_states is None else hidden_states

    collected: list[torch.Tensor] = []
    for layer_index, layer in enumerate(layers[start:stop], start=start):
        if collect:
            collected.append(h)
        attention_type = getattr(layer, "attention_type", "full_attention")
        h = layer(
            h,
            attention_mask=ctx.causal_mask_mapping[attention_type],
            position_ids=ctx.text_position_ids,
            past_key_values=ctx.past_key_values,
            use_cache=use_cache,
            cache_position=ctx.cache_position,
            position_embeddings=ctx.position_embeddings,
        )
        # 4.57 decoder layers return a bare tensor; older ones returned a tuple.
        if isinstance(h, tuple):
            h = h[0]
        if (
            ctx.deepstack_visual_embeds is not None
            and layer_index < len(ctx.deepstack_visual_embeds)
        ):
            if ctx.visual_pos_masks is None:
                raise ValueError("DeepStack features require visual_pos_masks")
            h = text_model._deepstack_process(
                h,
                ctx.visual_pos_masks,
                ctx.deepstack_visual_embeds[layer_index],
            )

    if collect:
        collected.append(h)
        return h, collected
    return h


def final_norm(text_model, hidden_states: torch.Tensor) -> torch.Tensor:
    """Apply the stack's final RMSNorm. ``[B, L, d] -> [B, L, d]``."""
    return text_model.norm(hidden_states)


# ---------------------------------------------------------------------------
# token selection operators (Pi_img / Pi_q in §3.2)
# ---------------------------------------------------------------------------


def image_token_mask(input_ids: torch.LongTensor, image_token_id: int) -> torch.Tensor:
    """``Pi_img`` support: ``[B, L]`` bool, True at image placeholder positions."""
    return input_ids == image_token_id


def vision_span_mask(input_ids: torch.LongTensor, config) -> torch.Tensor:
    """``[B, L]`` bool covering ``<|vision_start|>``, image pads, ``<|vision_end|>``.

    Use this (not :func:`image_token_mask`) when *removing* the visual segment to
    build the text-only branch, so the delimiters do not survive as orphans.
    """
    ids = {
        config.vision_start_token_id,
        config.vision_end_token_id,
        config.image_token_id,
        config.video_token_id,
    }
    mask = torch.zeros_like(input_ids, dtype=torch.bool)
    for tid in ids:
        mask |= input_ids == tid
    return mask


def select_tokens(
    hidden_states: torch.Tensor,  # [B, L, d]
    mask: torch.Tensor,  # [B, L] bool
) -> torch.Tensor:
    """Gather masked positions. Requires an equal count per batch element.

    Returns ``[B, N, d]`` where ``N`` is that per-element count.
    """
    counts = mask.sum(dim=1)
    if counts.numel() > 1 and not bool((counts == counts[0]).all()):
        raise ValueError(
            f"select_tokens needs the same number of selected tokens per batch "
            f"element, got {counts.tolist()}. Bucket by visual-token count or "
            f"gather per-example instead."
        )
    n = int(counts[0])
    b, _, d = hidden_states.shape
    return hidden_states[mask].view(b, n, d)


def select_tokens_padded(
    hidden_states: torch.Tensor,  # [B, L, d]
    mask: torch.Tensor,  # [B, L] bool
) -> tuple[torch.Tensor, torch.Tensor]:
    """Gather masked positions, right-padded to the batch maximum.

    Qwen2.5-VL uses dynamic resolution, so ``N_v`` differs across a batch. §3.1's
    patching genuinely needs equal counts (it transplants position by position),
    but ``r_theta`` only cross-attends over ``V*`` -- a variable-length memory is
    exactly what a key-padding mask is for.

    Returns ``(padded [B, N_max, d], key_padding_mask [B, N_max])`` where the
    mask is ``True`` at padding, matching ``nn.MultiheadAttention``.
    """
    counts = mask.sum(dim=1)
    n_max = int(counts.max())
    b, _, d = hidden_states.shape
    out = hidden_states.new_zeros((b, n_max, d))
    pad = torch.ones((b, n_max), dtype=torch.bool, device=hidden_states.device)
    for i in range(b):
        n = int(counts[i])
        out[i, :n] = hidden_states[i][mask[i]]
        pad[i, :n] = False
    return out, pad


__all__ = [
    "SplitContext",
    "embed_multimodal",
    "make_split_context",
    "block_attention_edges",
    "run_layer_range",
    "final_norm",
    "image_token_mask",
    "vision_span_mask",
    "select_tokens",
    "select_tokens_padded",
]