"""
SOTA Flash Attention and Cross-Attention implementations with KV cache support.

Features:
- Full KV cache support for efficient autoregressive generation
- Flash Attention 2 integration via PyTorch SDPA
- Pre-scaled Q/K for FP16 stability (prevents overflow in Q@K^T)
- Grouped Query Attention (GQA) support
- Memory-efficient cross-attention for multimodal fusion
"""

import logging
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List
from dataclasses import dataclass

logger = logging.getLogger(__name__)


def flash_attention_available() -> bool:
    """Check if Flash Attention (via SDPA) is available."""
    try:
        from torch.nn.functional import scaled_dot_product_attention
        return True
    except ImportError:
        return False


def compute_qk_scale(head_dim: int) -> float:
    """Compute the Q/K pre-scaling factor for FP16 stability.
    
    By scaling both Q and K by head_dim^-0.25, the product Q@K^T
    is effectively scaled by head_dim^-0.5 (the standard attention scaling).
    This prevents overflow in FP16 when Q and K have large values.
    """
    return head_dim ** -0.25


class AttentionKVCache:
    """Pre-allocated KV Cache — static buffer with index-based filling.
    
    Eliminates VRAM fragmentation from torch.cat during autoregressive generation.
    Buffer is allocated once at first use and reused via slice assignment.
    """
    
    __slots__ = ('key_cache', 'value_cache', 'seen_tokens', '_max_len')
    
    def __init__(self, max_seq_len: int = 131072):
        self.key_cache: torch.Tensor = None
        self.value_cache: torch.Tensor = None
        self.seen_tokens: int = 0
        self._max_len = max_seq_len
    
    def _allocate(self, batch: int, heads: int, head_dim: int, device: torch.device, dtype: torch.dtype):
        """Allocate static buffer on first use."""
        self.key_cache = torch.zeros(batch, heads, self._max_len, head_dim, device=device, dtype=dtype)
        self.value_cache = torch.zeros(batch, heads, self._max_len, head_dim, device=device, dtype=dtype)
    
    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Update cache with new key/value states using index-based filling.
        
        Args:
            key_states: New key states [batch, num_heads, seq_len, head_dim]
            value_states: New value states [batch, num_heads, seq_len, head_dim]
            
        Returns:
            Updated key and value states including cache (views, no copy)
        """
        batch, heads, new_len, head_dim = key_states.shape
        
        if self.key_cache is None:
            self._allocate(batch, heads, head_dim, key_states.device, key_states.dtype)
            self.seen_tokens = 0
        
        # Grow buffer if needed (rare fallback)
        if self.seen_tokens + new_len > self.key_cache.shape[2]:
            new_max = max(self.key_cache.shape[2] * 2, self.seen_tokens + new_len)
            new_key = torch.zeros(batch, heads, new_max, head_dim, device=key_states.device, dtype=key_states.dtype)
            new_val = torch.zeros(batch, heads, new_max, head_dim, device=key_states.device, dtype=key_states.dtype)
            new_key[:, :, :self.seen_tokens] = self.key_cache[:, :, :self.seen_tokens]
            new_val[:, :, :self.seen_tokens] = self.value_cache[:, :, :self.seen_tokens]
            self.key_cache = new_key
            self.value_cache = new_val
        
        # Index-based fill — no allocation, no fragmentation
        self.key_cache[:, :, self.seen_tokens:self.seen_tokens + new_len] = key_states
        self.value_cache[:, :, self.seen_tokens:self.seen_tokens + new_len] = value_states
        self.seen_tokens += new_len
        
        # Return valid slice (view, no copy)
        return self.key_cache[:, :, :self.seen_tokens], self.value_cache[:, :, :self.seen_tokens]
    
    def get_seq_length(self) -> int:
        """Get current sequence length in cache."""
        return self.seen_tokens
    
    def reset(self):
        """Reset cache position without deallocating the buffer."""
        self.seen_tokens = 0


class FlashAttention(nn.Module):
    """
    SOTA Flash Attention with KV cache support and FP16-safe Q/K pre-scaling.
    
    Uses PyTorch's scaled_dot_product_attention when available,
    with fallback to standard attention. Supports:
    - KV caching for efficient generation
    - Causal masking
    - Attention dropout
    - Pre-scaled Q/K for FP16 stability
    """
    
    def __init__(
        self, 
        dropout: float = 0.0, 
        causal: bool = False,
        head_dim: int = None,
    ):
        super().__init__()
        self.dropout = dropout
        self.causal = causal
        self._flash_available = flash_attention_available()
        # Store head_dim for Q/K scaling; will be inferred from input if not provided
        self._head_dim = head_dim
        self._qk_scale = compute_qk_scale(head_dim) if head_dim else None

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_mask: torch.Tensor = None,
        is_causal: bool = None,
        past_key_value: Tuple[torch.Tensor, torch.Tensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
        """
        Forward pass with KV cache support.
        
        Args:
            query: Query tensor [batch, num_heads, seq_len, head_dim]
            key: Key tensor [batch, num_heads, seq_len, head_dim]
            value: Value tensor [batch, num_heads, seq_len, head_dim]
            attn_mask: Optional attention mask
            is_causal: Override causal setting
            past_key_value: Optional tuple of (past_key, past_value) for KV cache
            use_cache: Whether to return updated KV cache
            output_attentions: Whether to return attention weights
            
        Returns:
            Tuple of (output, present_key_value, attention_weights)
        """
        causal = is_causal if is_causal is not None else self.causal
        batch_size, num_heads, seq_len, head_dim = query.shape
        
        # Compute Q/K scaling factor (for FP16 stability)
        qk_scale = self._qk_scale if self._qk_scale else compute_qk_scale(head_dim)
        
        # Handle KV cache
        if past_key_value is not None:
            past_key, past_value = past_key_value
            key = torch.cat([past_key, key], dim=2)
            value = torch.cat([past_value, value], dim=2)
        
        # Prepare cache for return
        present_key_value = (key, value) if use_cache else None
        
        kv_seq_len = key.shape[2]
        attn_weights = None

        if self._flash_available and not output_attentions:
            # CRITICAL: Scale Q and K BEFORE SDPA to prevent FP16 overflow
            query_scaled = query * qk_scale
            key_scaled = key * qk_scale
            
            dropout_p = self.dropout if self.training else 0.0
            use_causal = causal and attn_mask is None and seq_len > 1 and seq_len == kv_seq_len
            
            output = F.scaled_dot_product_attention(
                query_scaled, key_scaled, value,
                attn_mask=attn_mask,
                dropout_p=dropout_p,
                is_causal=use_causal,
                scale=1.0,  # We already scaled Q and K
            )
        else:
            # Manual attention computation with standard scaling
            scale = 1.0 / math.sqrt(head_dim)
            attn_weights = torch.matmul(query, key.transpose(-2, -1)) * scale

            # Apply causal mask if needed
            if causal and attn_mask is None and seq_len > 1:
                causal_mask = torch.triu(
                    torch.full((seq_len, kv_seq_len), float('-inf'), device=query.device, dtype=query.dtype),
                    diagonal=kv_seq_len - seq_len + 1
                )
                attn_weights = attn_weights + causal_mask.unsqueeze(0).unsqueeze(0)

            if attn_mask is not None:
                attn_weights = attn_weights + attn_mask

            attn_weights = F.softmax(attn_weights, dim=-1, dtype=query.dtype)

            if self.training and self.dropout > 0:
                attn_weights = F.dropout(attn_weights, p=self.dropout)

            output = torch.matmul(attn_weights, value)

        return output, present_key_value, attn_weights


class MultimodalCrossAttention(nn.Module):
    """
    SOTA Cross-attention layer for multimodal fusion with KV cache support.
    
    Allows text to attend to image/video/audio features with:
    - KV caching for efficient generation
    - Gated residual connection for stable training
    - Flash Attention support with pre-scaled Q/K for FP16 stability
    - Optional attention weight output
    """
    
    def __init__(
        self,
        hidden_size: int,
        num_heads: int = 8,
        dropout: float = 0.1,
        use_flash_attention: bool = True,
        gate_init: float = 0.0,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.use_flash_attention = use_flash_attention and flash_attention_available()
        self.dropout_p = dropout
        
        # Pre-compute Q/K scaling factor for FP16 stability
        self.qk_scale = compute_qk_scale(self.head_dim)

        # Projections
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

        # Normalization and gating
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(hidden_size)
        self.gate = nn.Parameter(torch.tensor(gate_init))

    def forward(
        self,
        text_hidden: torch.Tensor,
        modality_hidden: torch.Tensor,
        modality_mask: torch.Tensor = None,
        past_key_value: Tuple[torch.Tensor, torch.Tensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
        batch_size, text_len, _ = text_hidden.shape
        
        # Compute query from text
        query = self.q_proj(text_hidden)
        query = query.view(batch_size, text_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Compute key/value from modality or use cache
        if past_key_value is not None:
            key, value = past_key_value
        else:
            modality_len = modality_hidden.shape[1]
            key = self.k_proj(modality_hidden)
            value = self.v_proj(modality_hidden)
            key = key.view(batch_size, modality_len, self.num_heads, self.head_dim).transpose(1, 2)
            value = value.view(batch_size, modality_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        present_key_value = (key, value) if use_cache else None
        attn_weights = None
        
        if self.use_flash_attention and not output_attentions:
            try:
                from xorfice.optimize import fused_cross_attention_kernel
                # Triton kernel generally expects raw tensors, let's keep scaling logic native to Triton
                attn_output = fused_cross_attention_kernel(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), self.qk_scale)
                attn_output = attn_output.transpose(1, 2)
            except ImportError:
                # Fallback purely if triton is somehow disabled
                query_scaled = query * self.qk_scale
                key_scaled = key * self.qk_scale
                dropout_p = self.dropout_p if self.training else 0.0
                attn_output = F.scaled_dot_product_attention(
                    query_scaled, key_scaled, value,
                    attn_mask=modality_mask,
                    dropout_p=dropout_p,
                    is_causal=False,
                    scale=1.0, 
                )
        else:
            # Manual attention computation with standard scaling
            scale = 1.0 / math.sqrt(self.head_dim)
            attn_weights = torch.matmul(query, key.transpose(-2, -1)) * scale
            
            if modality_mask is not None:
                attn_weights = attn_weights + modality_mask
                
            attn_weights = F.softmax(attn_weights, dim=-1, dtype=text_hidden.dtype)
            
            if self.training and self.dropout_p > 0:
                attn_weights = F.dropout(attn_weights, p=self.dropout_p)
                
            attn_output = torch.matmul(attn_weights, value)

        # Reshape and project output
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, text_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)

        # Gated residual connection
        gate = torch.sigmoid(self.gate)
        output = text_hidden + gate * self.dropout(attn_output)
        output = self.layer_norm(output)

        return output, present_key_value, attn_weights


@dataclass
class MultimodalFusionCache:
    """Cache for multimodal fusion layer KV states."""
    image_kv: Tuple[torch.Tensor, torch.Tensor] = None
    video_kv: Tuple[torch.Tensor, torch.Tensor] = None
    audio_kv: Tuple[torch.Tensor, torch.Tensor] = None


class MultimodalFusionLayer(nn.Module):
    """
    SOTA Multimodal fusion layer with cross-attention for all modalities and KV cache support.
    
    Features:
    - Separate cross-attention for each modality (image, video, audio)
    - KV caching for efficient generation
    - Gated fusion MLP
    - Flash Attention support
    """
    
    def __init__(
        self,
        hidden_size: int,
        num_heads: int = 8,
        dropout: float = 0.1,
        use_flash_attention: bool = True,
    ):
        super().__init__()
        self.hidden_size = hidden_size

        # Cross-attention for each modality
        self.image_cross_attn = MultimodalCrossAttention(
            hidden_size, num_heads, dropout, use_flash_attention
        )
        self.video_cross_attn = MultimodalCrossAttention(
            hidden_size, num_heads, dropout, use_flash_attention
        )
        self.audio_cross_attn = MultimodalCrossAttention(
            hidden_size, num_heads, dropout, use_flash_attention
        )

        # Fusion MLP with gating
        self.fusion_mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size * 4, hidden_size),
            nn.Dropout(dropout),
        )
        self.fusion_norm = nn.LayerNorm(hidden_size)

    def forward(
        self,
        text_hidden: torch.Tensor,
        image_hidden: torch.Tensor = None,
        video_hidden: torch.Tensor = None,
        audio_hidden: torch.Tensor = None,
        image_mask: torch.Tensor = None,
        video_mask: torch.Tensor = None,
        audio_mask: torch.Tensor = None,
        past_key_values: MultimodalFusionCache = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, MultimodalFusionCache]:
        """
        Fuse text with available modalities via cross-attention with KV cache support.
        
        Args:
            text_hidden: Text hidden states [batch, text_len, hidden_size]
            image_hidden: Image features [batch, image_len, hidden_size]
            video_hidden: Video features [batch, video_len, hidden_size]
            audio_hidden: Audio features [batch, audio_len, hidden_size]
            image_mask: Attention mask for image
            video_mask: Attention mask for video
            audio_mask: Attention mask for audio
            past_key_values: Cached KV states from previous forward pass
            use_cache: Whether to return updated KV cache
            
        Returns:
            Tuple of (output, present_key_values)
        """
        present_key_values = MultimodalFusionCache() if use_cache else None
        
        # Get past KV states
        past_image_kv = past_key_values.image_kv if past_key_values else None
        past_video_kv = past_key_values.video_kv if past_key_values else None
        past_audio_kv = past_key_values.audio_kv if past_key_values else None
        
        # Image cross-attention
        if self._has_content(image_hidden) or past_image_kv is not None:
            try:
                text_hidden, image_kv, _ = self.image_cross_attn(
                    text_hidden, 
                    image_hidden if image_hidden is not None else torch.zeros(text_hidden.shape[0], 1, self.hidden_size, device=text_hidden.device),
                    image_mask,
                    past_key_value=past_image_kv,
                    use_cache=use_cache,
                )
                if use_cache:
                    present_key_values.image_kv = image_kv
            except Exception as e:
                logger.debug(f"Image cross-attention skipped: {e}")

        # Video cross-attention
        if self._has_content(video_hidden) or past_video_kv is not None:
            try:
                text_hidden, video_kv, _ = self.video_cross_attn(
                    text_hidden, 
                    video_hidden if video_hidden is not None else torch.zeros(text_hidden.shape[0], 1, self.hidden_size, device=text_hidden.device),
                    video_mask,
                    past_key_value=past_video_kv,
                    use_cache=use_cache,
                )
                if use_cache:
                    present_key_values.video_kv = video_kv
            except Exception as e:
                logger.debug(f"Video cross-attention skipped: {e}")

        # Audio cross-attention
        if self._has_content(audio_hidden) or past_audio_kv is not None:
            try:
                text_hidden, audio_kv, _ = self.audio_cross_attn(
                    text_hidden, 
                    audio_hidden if audio_hidden is not None else torch.zeros(text_hidden.shape[0], 1, self.hidden_size, device=text_hidden.device),
                    audio_mask,
                    past_key_value=past_audio_kv,
                    use_cache=use_cache,
                )
                if use_cache:
                    present_key_values.audio_kv = audio_kv
            except Exception as e:
                logger.debug(f"Audio cross-attention skipped: {e}")

        # Fusion MLP
        residual = text_hidden
        text_hidden = self.fusion_mlp(text_hidden)
        text_hidden = self.fusion_norm(residual + text_hidden)

        return text_hidden, present_key_values
    
    @staticmethod
    def _has_content(tensor: torch.Tensor) -> bool:
        """Check if tensor has meaningful content."""
        if tensor is None:
            return False
        if not isinstance(tensor, torch.Tensor):
            return False
        try:
            if tensor.numel() == 0:
                return False
            return bool(tensor.any())
        except Exception:
            return False
