"""
FlashAttention-3 Integration for NVIDIA GPUs

Optimized attention implementation for H100, Blackwell, and newer NVIDIA GPUs.
Uses shared attention operations from torchbridge.attention.core.attention_ops.
"""

import logging
import warnings

import torch
import torch.nn as nn

from torchbridge.attention.core.attention_ops import (
    check_flash_attention_available,
    flash_attention_forward,
)
from torchbridge.core.config import TorchBridgeConfig

logger = logging.getLogger(__name__)


class FlashAttention3(nn.Module):
    """
    FlashAttention-3 implementation for NVIDIA H100/Blackwell.

    Provides memory-efficient attention with 3x memory reduction
    and faster execution compared to standard attention.

    This implementation uses the shared attention operations from
    torchbridge.attention.core.attention_ops for the core computation.
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        bias: bool = True,
        causal: bool = False,
        config: TorchBridgeConfig | None = None
    ):
        """
        Initialize FlashAttention-3.

        Args:
            embed_dim: Embedding dimension
            num_heads: Number of attention heads
            dropout: Dropout probability
            bias: Whether to use bias in linear layers
            causal: Whether to use causal (autoregressive) masking
            config: TorchBridge configuration
        """
        super().__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.causal = causal
        self.head_dim = embed_dim // num_heads

        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})"
            )

        self.config = config or TorchBridgeConfig()
        self.nvidia_config = self.config.hardware.nvidia

        # Check for optimal head dimensions
        optimal_div = 16 if self.nvidia_config.tensor_core_version >= 4 else 8
        if self.head_dim % optimal_div != 0:
            warnings.warn(
                f"Head dimension {self.head_dim} not divisible by {optimal_div}. "
                f"Consider using head_dim divisible by {optimal_div} for optimal Tensor Core performance.",
            stacklevel=2,
            )

        # QKV projection
        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)

        # Output projection
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        # Dropout for residual
        self.resid_dropout = nn.Dropout(dropout) if dropout > 0 else None

        # FlashAttention availability (using shared check)
        self._flash_available = check_flash_attention_available()

        # Use FlashAttention-3 if available
        self.use_flash_attention = (
            self._flash_available and
            self.nvidia_config.flash_attention_enabled and
            self.nvidia_config.flash_attention_version == "3"
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        return_attention_weights: bool = False
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Forward pass with FlashAttention-3.

        Args:
            x: Input tensor of shape (batch, seq_len, embed_dim)
            attention_mask: Optional attention mask
            return_attention_weights: Whether to return attention weights

        Returns:
            Tuple of (output, attention_weights)
        """
        batch_size, seq_len, embed_dim = x.shape

        # Compute Q, K, V
        qkv = self.qkv(x)
        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch, num_heads, seq_len, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # Use shared attention forward (handles FlashAttention and fallback)
        attn_output, attn_weights = flash_attention_forward(
            Q=q,
            K=k,
            V=v,
            scale=None,  # Use default 1/sqrt(head_dim)
            causal=self.causal,
            dropout=self.dropout,
            training=self.training,
            attention_mask=attention_mask,
            return_weights=return_attention_weights,
        )

        # Reshape and project output
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(batch_size, seq_len, embed_dim)
        output = self.out_proj(attn_output)

        if self.resid_dropout:
            output = self.resid_dropout(output)

        return output, attn_weights


def create_flash_attention_3(
    embed_dim: int,
    num_heads: int,
    dropout: float = 0.0,
    bias: bool = True,
    config: TorchBridgeConfig | None = None
) -> FlashAttention3:
    """
    Factory function to create FlashAttention-3 module.

    Args:
        embed_dim: Embedding dimension
        num_heads: Number of attention heads
        dropout: Dropout probability
        bias: Whether to use bias
        config: TorchBridge configuration

    Returns:
        FlashAttention3 module
    """
    return FlashAttention3(
        embed_dim=embed_dim,
        num_heads=num_heads,
        dropout=dropout,
        bias=bias,
        config=config
    )


class FlashAttention3Optimizer:
    """
    Optimizer for FlashAttention-3 integration.

    Helps convert existing attention modules to FlashAttention-3.
    """

    def __init__(self, config: TorchBridgeConfig | None = None):
        """
        Initialize FlashAttention-3 optimizer.

        Args:
            config: TorchBridge configuration
        """
        self.config = config or TorchBridgeConfig()

    def convert_attention_layer(
        self,
        module: nn.Module
    ) -> nn.Module:
        """
        Convert attention module to FlashAttention-3.

        Args:
            module: Attention module to convert

        Returns:
            FlashAttention-3 module or original if conversion not possible
        """
        # Check if module is MultiheadAttention
        if isinstance(module, nn.MultiheadAttention):
            return FlashAttention3(
                embed_dim=module.embed_dim,
                num_heads=module.num_heads,
                dropout=module.dropout,
                bias=module.in_proj_bias is not None,
                config=self.config
            )

        # Return original module if not convertible
        return module

    def optimize_model(self, model: nn.Module) -> tuple[nn.Module, int]:
        """
        Convert all attention layers in model to FlashAttention-3.

        Args:
            model: PyTorch model

        Returns:
            Tuple of (optimized_model, num_conversions)
        """
        num_conversions = 0

        for name, module in model.named_children():
            if isinstance(module, nn.MultiheadAttention):
                # Replace with FlashAttention-3
                flash_attn = self.convert_attention_layer(module)
                setattr(model, name, flash_attn)
                num_conversions += 1
            elif len(list(module.children())) > 0:
                # Recursively optimize child modules
                _, child_conversions = self.optimize_model(module)
                num_conversions += child_conversions

        return model, num_conversions
