"""GreenLeaf Law Embed — model implementation.

Bidirectional transformer encoder for legal-domain text embedding.
Built on the Qwen3 architecture with causal masking removed,
enabling full-sequence attention in both directions.

This is critical for legal text where relevant context (holdings,
citations, defined terms) can appear anywhere in a document.
"""

import inspect
from typing import Callable

import torch
from transformers import Qwen3Model
from transformers.cache_utils import Cache
from transformers.masking_utils import create_causal_mask
from transformers.modeling_outputs import BaseModelOutputWithPooling
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs

from .configuration import GreenLeafEmbedConfig

# ---------------------------------------------------------------------------
# Compatibility shim: the `create_causal_mask` API has changed across
# transformers releases. We detect the correct parameter names at import
# time so this code works across versions 5.1 through 5.15+.
#
#   <= 5.1   : kwarg is `input_embeds`, `cache_position` required
#   5.2-5.5  : renamed to `inputs_embeds`, `cache_position` still required
#   5.6-5.8  : `cache_position` has a default (backward compat)
#   >= 5.9   : `cache_position` removed entirely
# ---------------------------------------------------------------------------
_mask_fn_params = inspect.signature(create_causal_mask).parameters
_embeds_param = "inputs_embeds" if "inputs_embeds" in _mask_fn_params else "input_embeds"
_has_cache_position = "cache_position" in _mask_fn_params


def _build_bidirectional_mask_fn(attn_mask: torch.Tensor | None) -> Callable:
    """Return a mask function that allows every token to attend to every
    other token, subject only to the padding mask.

    Standard causal masking restricts token i to attend only to tokens
    j <= i. For embedding models we want the opposite — full visibility —
    so that the representation of each token is informed by the entire
    input sequence.
    """

    def _mask(batch: int, head: int, q_pos: int, kv_pos: int) -> bool:
        if attn_mask is None:
            return torch.ones((), dtype=torch.bool)
        return attn_mask[batch, kv_pos].to(torch.bool)

    return _mask


class GreenLeafEmbedModel(Qwen3Model):
    """Bidirectional Qwen3 encoder for text embedding.

    Overrides the causal self-attention in Qwen3 with bidirectional
    attention so that every token representation captures full-sequence
    context. This is essential for retrieval tasks where the meaning of
    a passage depends on information that may appear before or after
    any given token.
    """

    _supports_flash_attn = True
    _supports_sdpa = True

    config_class = GreenLeafEmbedConfig

    def __init__(self, config):
        super().__init__(config)
        self.post_init()

    def post_init(self):
        super().post_init()
        # Disable causal attention in every transformer layer.
        # This works with both flash_attention_2 and sdpa backends.
        for layer in self.layers:
            layer.self_attn.is_causal = False

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        cache_position: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPooling:
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
            input_ids = None

        mask_args = {
            "config": self.config,
            _embeds_param: inputs_embeds,
            "attention_mask": attention_mask,
            "past_key_values": None,
            "position_ids": position_ids,
            "or_mask_function": _build_bidirectional_mask_fn(attention_mask),
        }
        if _has_cache_position:
            mask_args["cache_position"] = torch.arange(
                inputs_embeds.shape[1],
                device=inputs_embeds.device,
                dtype=torch.long,
            )
        attention_mask = {"full_attention": create_causal_mask(**mask_args)}

        outputs = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            cache_position=cache_position,
            **kwargs,
        )
        return outputs