from __future__ import annotations

import math

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

from src.model.config import BoccaccioConfig
from src.model.layers import RotaryEmbedding, apply_rotary_pos_emb


class GroupedQueryAttention(nn.Module):

    def __init__(self, config: BoccaccioConfig) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.num_kv_heads = config.num_kv_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.num_attention_heads // self.num_kv_heads

        self.q_proj = nn.Linear(self.hidden_size, self.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_attention_heads * self.head_dim, self.hidden_size, bias=False)

        self.rotary_emb = RotaryEmbedding(
            self.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            theta=config.rope_theta,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch, seq_len, _ = hidden_states.shape

        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)

        q = q.view(batch, seq_len, self.num_attention_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(v, position_ids)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Expand KV heads to match the number of query heads.
        k = k.repeat_interleave(self.num_kv_groups, dim=1)
        v = v.repeat_interleave(self.num_kv_groups, dim=1)

        # Build the attention mask: causal + padding.
        # attention_mask from InstructionDataset is (batch, seq_len) with 1=real, 0=pad.
        # We need a 3D mask (batch, 1, seq_len, seq_len) that combines:
        #   - causal masking (lower triangular)
        #   - padding masking (ignore pad positions)
        if attention_mask is not None:
            # Create causal mask (lower triangular = 0, upper = -inf).
            causal_mask = torch.zeros(
                seq_len, seq_len, device=q.device, dtype=q.dtype
            )
            causal_mask = torch.triu(
                torch.full((seq_len, seq_len), float("-inf"), device=q.device, dtype=q.dtype),
                diagonal=1,
            )

            # Create padding mask: (batch, 1, 1, seq_len) with 0 for real, -inf for pad.
            # Use masked_fill to avoid 0 * -inf = NaN.
            pad_mask = attention_mask[:, None, None, :].to(q.dtype)  # (batch, 1, 1, seq_len)
            pad_mask = pad_mask.masked_fill(pad_mask == 0, float("-inf"))

            # Combine: additive mask (batch, 1, seq_len, seq_len).
            # causal_mask is (1, 1, seq_len, seq_len), pad_mask is (batch, 1, 1, seq_len).
            # Broadcasting: pad_mask broadcasts over query dim, causal over batch dim.
            attn_mask = causal_mask[None, None, :, :] + pad_mask

            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        else:
            # No padding mask (pre-training): use built-in causal.
            attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

        attn_output = attn_output.transpose(1, 2).contiguous().view(batch, seq_len, self.hidden_size)
        return self.o_proj(attn_output)