"""ACE-Step attention module for MLX.

Ported from:
  docs/reference/modeling_acestep_v15_turbo.py (AceStepAttention, lines 253-368)

Contains:
  - AceStepAttention: GQA with RoPE and sliding window support
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Optional

import mlx.core as mx
import mlx.nn as nn

from .config import AceStepDiTConfig

if TYPE_CHECKING:
    from .cache import CrossAttentionCache


def _make_sliding_window_mask(
    seq_len: int, window_size: int, dtype: mx.Dtype = mx.float32
) -> mx.array:
    """Create a sliding window attention mask.

    Positions outside the window of +/- window_size are masked with -inf.
    Positions inside the window (including the center) are 0.

    Args:
        seq_len: Sequence length.
        window_size: Number of positions to attend to on each side.
        dtype: Output dtype for the mask.

    Returns:
        Mask of shape [1, 1, seq_len, seq_len] suitable for additive masking.
    """
    # row_idx[i] - col_idx[j] gives the relative position
    row_idx = mx.arange(seq_len)[:, None]  # [S, 1]
    col_idx = mx.arange(seq_len)[None, :]  # [1, S]
    distance = mx.abs(row_idx - col_idx)  # [S, S]

    # Mask: 0 where within window, -inf where outside
    mask = mx.where(
        distance <= window_size,
        mx.zeros((1,), dtype=dtype),
        mx.full((1,), float("-inf"), dtype=dtype),
    )
    # Expand to [1, 1, S, S] for broadcasting over batch and heads
    return mask[None, None, :, :]


def _repeat_kv(x: mx.array, num_groups: int) -> mx.array:
    """Repeat key/value heads to match query head count for GQA.

    Args:
        x: Key or value tensor of shape [B, num_kv_heads, S, head_dim].
        num_groups: Number of times to repeat each KV head.

    Returns:
        Tensor of shape [B, num_kv_heads * num_groups, S, head_dim].
    """
    if num_groups == 1:
        return x
    B, num_kv_heads, S, head_dim = x.shape
    # Expand and reshape: [B, kv_h, 1, S, D] -> [B, kv_h, G, S, D] -> [B, kv_h*G, S, D]
    x = mx.expand_dims(x, axis=2)  # [B, kv_h, 1, S, D]
    x = mx.broadcast_to(x, (B, num_kv_heads, num_groups, S, head_dim))
    return x.reshape(B, num_kv_heads * num_groups, S, head_dim)


class AceStepAttention(nn.Module):
    """Multi-headed attention for ACE-Step with GQA, RoPE, and sliding window.

    Supports both self-attention and cross-attention modes:
      - Self-attention: Q, K, V all projected from hidden_states; RoPE applied
        to Q and K; optional sliding window mask.
      - Cross-attention: Q from hidden_states, K/V from encoder_hidden_states;
        no RoPE; no sliding window.

    Uses grouped query attention (GQA) with:
      - 16 query heads, 8 KV heads -> 2 groups (default config)

    Ported from: AceStepAttention in modeling_acestep_v15_turbo.py (lines 253-368)
    """

    def __init__(
        self,
        config: AceStepDiTConfig,
        layer_idx: int,
        is_cross_attention: bool = False,
    ):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.is_cross_attention = is_cross_attention

        self.head_dim = getattr(
            config, "head_dim", config.hidden_size // config.num_attention_heads
        )
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_kv_heads
        self.scale = self.head_dim**-0.5

        # Projections
        self.q_proj = nn.Linear(
            config.hidden_size,
            self.num_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.k_proj = nn.Linear(
            config.hidden_size,
            self.num_kv_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.v_proj = nn.Linear(
            config.hidden_size,
            self.num_kv_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.o_proj = nn.Linear(
            self.num_heads * self.head_dim,
            config.hidden_size,
            bias=config.attention_bias,
        )

        # QK normalization (per-head RMSNorm)
        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)

        # RoPE for self-attention
        if not is_cross_attention:
            self.rope = nn.RoPE(
                self.head_dim,
                traditional=False,
                base=config.rope_theta,
            )

        # Sliding window: only for self-attention on sliding_attention layers
        self.attention_type = config.layer_types[layer_idx]
        if not is_cross_attention and self.attention_type == "sliding_attention":
            self.sliding_window = config.sliding_window
        else:
            self.sliding_window = None

    def __call__(
        self,
        hidden_states: mx.array,
        encoder_hidden_states: Optional[mx.array] = None,
        mask: Optional[mx.array] = None,
        cache: Optional[CrossAttentionCache] = None,
    ) -> mx.array:
        """Forward pass.

        Args:
            hidden_states: Input tensor of shape [B, S, hidden_size].
            encoder_hidden_states: For cross-attention, encoder outputs of
                shape [B, S_enc, hidden_size]. None for self-attention.
            mask: Optional additive attention mask of shape
                [B, 1, S_q, S_kv] or broadcastable.
            cache: Optional cross-attention KV cache. When provided for
                cross-attention layers, K/V are cached on first call and
                reused on subsequent calls. Ignored for self-attention.

        Returns:
            Output tensor of shape [B, S, hidden_size].
        """
        B, S, _ = hidden_states.shape

        # --- Query projection + per-head RMSNorm ---
        # [B, S, num_heads * head_dim] -> [B, S, num_heads, head_dim] -> per-head norm
        q = self.q_proj(hidden_states)
        q = q.reshape(B, S, self.num_heads, self.head_dim)
        q = self.q_norm(q)
        # -> [B, num_heads, S, head_dim]
        q = q.transpose(0, 2, 1, 3)

        # --- Key/Value projection (with cross-attention cache) ---
        is_cross = self.is_cross_attention and encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states

        # Cross-attention cache: reuse encoder K/V across denoising steps
        if is_cross and cache is not None and cache.has_layer(self.layer_idx):
            k, v = cache.get(self.layer_idx)
        else:
            B_kv, S_kv, _ = kv_input.shape

            k = self.k_proj(kv_input)
            k = k.reshape(B_kv, S_kv, self.num_kv_heads, self.head_dim)
            k = self.k_norm(k)
            k = k.transpose(0, 2, 1, 3)  # [B, num_kv_heads, S_kv, head_dim]

            v = self.v_proj(kv_input)
            v = v.reshape(B_kv, S_kv, self.num_kv_heads, self.head_dim)
            v = v.transpose(0, 2, 1, 3)  # [B, num_kv_heads, S_kv, head_dim]

            # Store in cache for future steps
            if is_cross and cache is not None:
                cache.update(self.layer_idx, k, v)

        # --- RoPE (self-attention only, and only if module has RoPE) ---
        if not is_cross and hasattr(self, "rope"):
            q = self.rope(q)
            k = self.rope(k)

        # --- GQA: repeat KV heads to match query heads ---
        if self.num_key_value_groups > 1:
            k = _repeat_kv(k, self.num_key_value_groups)
            v = _repeat_kv(v, self.num_key_value_groups)

        # --- Sliding window mask (self-attention only) ---
        if self.sliding_window is not None and not is_cross:
            sw_mask = _make_sliding_window_mask(S, self.sliding_window, dtype=q.dtype)
            if mask is not None:
                mask = mask + sw_mask
            else:
                mask = sw_mask

        # --- Scaled dot-product attention ---
        # mx.fast.scaled_dot_product_attention expects:
        #   q: [B, num_heads, S_q, head_dim]
        #   k: [B, num_heads, S_kv, head_dim]
        #   v: [B, num_heads, S_kv, head_dim]
        #   scale: float
        #   mask: [B or 1, num_heads or 1, S_q, S_kv] (additive)
        attn_output = mx.fast.scaled_dot_product_attention(
            q, k, v, scale=self.scale, mask=mask
        )

        # --- Output projection ---
        # [B, num_heads, S, head_dim] -> [B, S, num_heads, head_dim] -> [B, S, hidden_size]
        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, S, -1)
        return self.o_proj(attn_output)
