from __future__ import annotations

from typing import Any

import torch
import torch.nn.functional as F
from torch import Tensor, nn
from transformers import Cache
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

from .configuration_neuron_lm import NeuronLMConfig
from .layers import RMSNorm
from .rotary import apply_rotary_pos_emb

__all__ = ["NeuronLMAttention"]


def _repeat_kv(hidden_states: Tensor, repeats: int) -> Tensor:
    if repeats == 1:
        return hidden_states
    return hidden_states.repeat_interleave(repeats, dim=1)


def eager_attention_forward(
    module: NeuronLMAttention,
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attention_mask: Tensor | None,
    *,
    scaling: float,
    dropout: float = 0.0,
    **_: Any,
) -> tuple[Tensor, Tensor]:
    """Numerically clear GQA reference used for attention-weight outputs."""

    key = _repeat_kv(key, module.num_key_value_groups)
    value = _repeat_kv(value, module.num_key_value_groups)
    attention_weights = (
        torch.matmul(
            query,
            key.transpose(-2, -1),
        )
        * scaling
    )

    fully_masked: Tensor | None = None
    if attention_mask is not None:
        if attention_mask.dtype == torch.bool:
            fully_masked = ~attention_mask.any(
                dim=-1,
                keepdim=True,
            )
            attention_weights = attention_weights.masked_fill(
                ~attention_mask,
                torch.finfo(attention_weights.dtype).min,
            )
        else:
            minimum = torch.finfo(attention_mask.dtype).min
            fully_masked = (
                torch.isneginf(attention_mask) | (attention_mask == minimum)
            ).all(dim=-1, keepdim=True)
            attention_weights = attention_weights + attention_mask

    if fully_masked is not None:
        attention_weights = attention_weights.masked_fill(
            fully_masked,
            0.0,
        )

    attention_weights = F.softmax(
        attention_weights,
        dim=-1,
        dtype=torch.float32,
    ).to(query.dtype)
    attention_weights = torch.nan_to_num(
        attention_weights,
        nan=0.0,
    )
    if fully_masked is not None:
        attention_weights = attention_weights.masked_fill(
            fully_masked,
            0.0,
        )
    attention_weights = F.dropout(
        attention_weights,
        p=dropout,
        training=module.training,
    )
    attention_output = torch.matmul(attention_weights, value)
    return (
        attention_output.transpose(1, 2).contiguous(),
        attention_weights,
    )


class NeuronLMAttention(nn.Module):
    """Fused, bias-free GQA using a checkpoint-stable ``[Q, K, V]`` layout."""

    def __init__(
        self,
        config: NeuronLMConfig,
        layer_idx: int = 0,
    ) -> None:
        super().__init__()

        if type(layer_idx) is not int or layer_idx < 0:
            raise ValueError(
                f"layer_idx must be a non-negative integer, got {layer_idx!r}"
            )

        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.head_dim = config.head_dim
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.layer_idx = layer_idx
        self.is_causal = True

        self.query_size = self.num_heads * self.head_dim
        self.key_value_size = self.num_key_value_heads * self.head_dim

        # State-dict contract: rows are Q, then K, then V.
        self.qkv_proj = nn.Linear(
            in_features=self.hidden_size,
            out_features=config.qkv_projection_size,
            bias=False,
        )
        self.out_proj = nn.Linear(
            in_features=self.query_size,
            out_features=self.hidden_size,
            bias=False,
        )

        # Per-head normalization of queries and keys before RoPE, as in
        # Qwen3 / OLMo-2 / Gemma-3. Bounds the growth of q.k during long bf16
        # runs, which depth-scaled initialization does not address: init
        # controls the residual stream at step 0, while attention logits
        # drift as the projection norms are learned.
        self.use_qk_norm = config.use_qk_norm
        if self.use_qk_norm:
            self.q_norm = RMSNorm(
                hidden_size=self.head_dim,
                eps=config.rms_norm_eps,
            )
            self.k_norm = RMSNorm(
                hidden_size=self.head_dim,
                eps=config.rms_norm_eps,
            )

    def forward(
        self,
        hidden_states: Tensor,
        position_embeddings: tuple[Tensor, Tensor],
        attention_mask: Tensor | None = None,
        past_key_values: Cache | None = None,
        output_attentions: bool = False,
        **kwargs: Any,
    ) -> Tensor | tuple[Tensor, Tensor | None]:
        if hidden_states.ndim != 3:
            raise ValueError(
                "hidden_states must have shape "
                "(batch_size, sequence_length, hidden_size), "
                f"got shape={tuple(hidden_states.shape)}"
            )

        batch_size, sequence_length, hidden_size = hidden_states.shape
        if hidden_size != self.hidden_size:
            raise ValueError(
                f"Expected hidden_size={self.hidden_size}, "
                f"got hidden_size={hidden_size}"
            )
        if sequence_length == 0:
            raise ValueError("sequence_length must be greater than zero")

        cos, sin = position_embeddings
        query_states, key_states, value_states = self._project_qkv(hidden_states)
        query_states, key_states = apply_rotary_pos_emb(
            query=query_states,
            key=key_states,
            cos=cos,
            sin=sin,
        )

        if past_key_values is not None:
            # Transformers v5 caches track their own write offset; passing
            # cache_position here was removed from the library's convention.
            key_states, value_states = past_key_values.update(
                key_states,
                value_states,
                self.layer_idx,
            )

        implementation = self.config._attn_implementation or "sdpa"
        if output_attentions:
            implementation = "eager"

        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
            implementation,
            eager_attention_forward,
        )
        attention_output, attention_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=(self.attention_dropout if self.training else 0.0),
            scaling=self.scaling,
            output_attentions=output_attentions,
            **kwargs,
        )

        attention_output = attention_output.reshape(
            batch_size,
            sequence_length,
            self.query_size,
        )
        attention_output = self.out_proj(attention_output)

        if output_attentions:
            return attention_output, attention_weights
        return attention_output

    def _project_qkv(
        self,
        hidden_states: Tensor,
    ) -> tuple[Tensor, Tensor, Tensor]:
        batch_size, sequence_length, _ = hidden_states.shape
        qkv_states = self.qkv_proj(hidden_states)
        query_states, key_states, value_states = qkv_states.split(
            (
                self.query_size,
                self.key_value_size,
                self.key_value_size,
            ),
            dim=-1,
        )

        query_states = query_states.view(
            batch_size,
            sequence_length,
            self.num_heads,
            self.head_dim,
        ).transpose(1, 2)
        key_states = key_states.view(
            batch_size,
            sequence_length,
            self.num_key_value_heads,
            self.head_dim,
        ).transpose(1, 2)
        value_states = value_states.view(
            batch_size,
            sequence_length,
            self.num_key_value_heads,
            self.head_dim,
        ).transpose(1, 2)

        # Applied before RoPE so the rotation acts on unit-scale vectors and
        # the norm never sees position-dependent structure.
        if self.use_qk_norm:
            query_states = self.q_norm(query_states)
            key_states = self.k_norm(key_states)

        return query_states, key_states, value_states

    def _load_from_state_dict(
        self,
        state_dict: dict[str, Tensor],
        prefix: str,
        local_metadata: dict[str, Any],
        strict: bool,
        missing_keys: list[str],
        unexpected_keys: list[str],
        error_msgs: list[str],
    ) -> None:
        qkv_key = f"{prefix}qkv_proj.weight"
        qkv_weight = state_dict.get(qkv_key)
        expected_shape = tuple(self.qkv_proj.weight.shape)
        if qkv_weight is not None and tuple(qkv_weight.shape) != expected_shape:
            error_msgs.append(
                f"{qkv_key} must use fused [Q, K, V] layout with shape "
                f"{expected_shape}, got {tuple(qkv_weight.shape)}"
            )

        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def extra_repr(self) -> str:
        return (
            f"hidden_size={self.hidden_size}, "
            f"num_heads={self.num_heads}, "
            f"num_key_value_heads={self.num_key_value_heads}, "
            f"head_dim={self.head_dim}, "
            f"attention_dropout={self.attention_dropout}, "
            f"use_qk_norm={self.use_qk_norm}, "
            f"layer_idx={self.layer_idx}"
        )