"""Bidirectional Qwen3 encoder for embedding models.

This is a thin wrapper around ``transformers``' ``Qwen3Model`` that turns the
decoder-only (causal) model into a bidirectional encoder: every self-attention
layer attends to the full sequence instead of only past tokens. Everything else
(weights, RoPE, RMSNorm, QK-norm) is identical to stock Qwen3, so the same
checkpoint loads without any key remapping.

Designed to work across a wide range of ``transformers`` releases (4.5x and 5.x)
by importing the bidirectional-mask helper defensively and falling back to a
local implementation when it is not available.
"""

from typing import Optional, Union

import torch

from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.qwen3.modeling_qwen3 import Qwen3Model
from transformers.utils import logging

from .configuration_qwen3_bidirectional import Qwen3BidirectionalConfig

logger = logging.get_logger(__name__)


# --- resolve a bidirectional-mask builder across transformers versions -------
# transformers >= 4.53 ships ``create_bidirectional_mask`` in ``masking_utils``.
# Its keyword names have shifted slightly across releases, so we normalise the
# call in ``_build_bidirectional_mask`` below.
try:
    from transformers.masking_utils import create_bidirectional_mask as _hf_bidir_mask
except Exception:  # pragma: no cover - very old transformers
    _hf_bidir_mask = None


def _build_bidirectional_mask(model: Qwen3Model, inputs_embeds, attention_mask):
    """Return a full (non-causal) attention mask for ``inputs_embeds``.

    Prefers transformers' native ``create_bidirectional_mask`` and adapts to the
    two keyword spellings it has used (``input_embeds`` vs ``inputs_embeds``).
    Falls back to expanding the 2-D padding mask into an additive 4-D mask.
    """
    config = model.config
    if _hf_bidir_mask is not None:
        for embeds_kw in ("inputs_embeds", "input_embeds"):
            try:
                return _hf_bidir_mask(
                    **{
                        "config": config,
                        embeds_kw: inputs_embeds,
                        "attention_mask": attention_mask,
                    }
                )
            except TypeError:
                continue

    # Fallback: additive mask from the padding mask. SDPA / eager both accept a
    # float mask of shape [batch, 1, q_len, kv_len] with 0 for keep, -inf for pad.
    bsz, seq_len = inputs_embeds.shape[:2]
    dtype = inputs_embeds.dtype
    if attention_mask is None:
        return None  # nothing to mask -> full bidirectional attention
    min_val = torch.finfo(dtype).min
    pad = (1.0 - attention_mask[:, None, None, :].to(dtype)) * min_val
    return pad.expand(bsz, 1, seq_len, seq_len)


class Qwen3BidirectionalModel(Qwen3Model):
    """Qwen3 with bidirectional (encoder-style) self-attention."""

    config_class = Qwen3BidirectionalConfig

    def __init__(self, config: Qwen3BidirectionalConfig):
        super().__init__(config)
        # Force every attention layer to be non-causal.
        for layer in self.layers:
            layer.self_attn.is_causal = False

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        bidirectional_mask = _build_bidirectional_mask(self, inputs_embeds, attention_mask)

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=bidirectional_mask,
                position_ids=position_ids,
                past_key_values=None,
                use_cache=False,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(last_hidden_state=hidden_states)