"""Fused Operations Kernels.

Provides optimized fused kernel implementations that combine multiple
operations to reduce memory bandwidth and kernel launch overhead.

Fused operations are particularly beneficial for transformer models where:
- Memory bandwidth is often the bottleneck
- Multiple element-wise operations can be combined
- Intermediate results don't need to be materialized

References:
    - xformers: https://github.com/facebookresearch/xformers
    - Triton: https://github.com/openai/triton
"""

from __future__ import annotations

import functools
import logging
from typing import TYPE_CHECKING

import torch
import torch.nn.functional as F

from sagellm_backend.kernels.base import Kernel

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


# Check if Triton is available for custom kernels
_TRITON_AVAILABLE = False
try:
    import triton  # type: ignore[import-untyped]
    import triton.language as tl  # type: ignore[import-untyped]

    _TRITON_AVAILABLE = True
    logger.info("✅ Triton available for fused kernels (optimized fusion)")
except ImportError:
    logger.warning(
        "⚠️  Triton not available, using PyTorch fallback for fused ops (performance impact)"
    )
    triton = None
    tl = None


def is_triton_available() -> bool:
    """Check if Triton is available for custom kernels."""
    return _TRITON_AVAILABLE


# =============================================================================
# Fused SiLU * Mul Kernel
# =============================================================================


class FusedSiLUMulKernel(Kernel):
    """Fused SiLU activation with element-wise multiply.

    Implements: output = silu(x) * y

    This is commonly used in LLaMA/Qwen style MLP blocks:
        gate = W_gate @ x
        up = W_up @ x
        output = silu(gate) * up  # This fused op

    Fusing these operations:
    - Reduces memory bandwidth (no intermediate silu(gate) tensor)
    - Reduces kernel launch overhead
    - Can be 15-30% faster than separate ops

    Example:
        kernel = FusedSiLUMulKernel()
        output = kernel(gate_proj, up_proj)
    """

    @property
    def name(self) -> str:
        """Kernel name."""
        return "fused_silu_mul"

    def forward(
        self,
        gate: torch.Tensor,
        up: torch.Tensor,
    ) -> torch.Tensor:
        """Execute fused SiLU * mul.

        Args:
            gate: Gate tensor to apply SiLU to.
            up: Up tensor to multiply with.

        Returns:
            Result of silu(gate) * up.
        """
        if _TRITON_AVAILABLE and gate.is_cuda:
            return self._triton_fused_silu_mul(gate, up)
        else:
            return self._pytorch_fused_silu_mul(gate, up)

    def _pytorch_fused_silu_mul(
        self,
        gate: torch.Tensor,
        up: torch.Tensor,
    ) -> torch.Tensor:
        """PyTorch fallback implementation.

        Args:
            gate: Gate tensor.
            up: Up tensor.

        Returns:
            Fused result.
        """
        return F.silu(gate) * up

    def _triton_fused_silu_mul(
        self,
        gate: torch.Tensor,
        up: torch.Tensor,
    ) -> torch.Tensor:
        """Triton optimized implementation.

        Falls back to PyTorch if the Triton kernel cannot be compiled
        (e.g. pyc-only install where inspect.getsourcelines has no .py to read).

        Args:
            gate: Gate tensor.
            up: Up tensor.

        Returns:
            Fused result.
        """
        kernel = _get_fused_silu_mul_kernel()
        if kernel is None:
            return self._pytorch_fused_silu_mul(gate, up)

        assert gate.shape == up.shape
        assert gate.is_contiguous() and up.is_contiguous()

        output = torch.empty_like(gate)
        n_elements = gate.numel()

        grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)  # noqa: E731
        kernel[grid](gate, up, output, n_elements, BLOCK_SIZE=1024)
        return output


@functools.lru_cache(maxsize=None)
def _get_fused_silu_mul_kernel():
    """Lazily compile the Triton fused-SiLU-mul kernel on first call.

    Defers @triton.jit decoration to call time so that module import
    succeeds even in pyc-only installs (where inspect.getsourcelines
    would raise OSError at decoration time).

    Returns:
        Compiled JITFunction, or None if Triton is unavailable or the
        source cannot be read (pyc-only install).
    """
    if not _TRITON_AVAILABLE:
        return None
    try:
        import triton as _triton
        import triton.language as _tl

        @_triton.jit
        def kernel(
            gate_ptr,
            up_ptr,
            output_ptr,
            n_elements,
            BLOCK_SIZE: _tl.constexpr,
        ):
            """Triton kernel for fused silu(gate) * up."""
            pid = _tl.program_id(0)
            block_start = pid * BLOCK_SIZE
            offsets = block_start + _tl.arange(0, BLOCK_SIZE)
            mask = offsets < n_elements
            gate = _tl.load(gate_ptr + offsets, mask=mask)
            up = _tl.load(up_ptr + offsets, mask=mask)
            silu_gate = gate * _tl.sigmoid(gate)
            result = silu_gate * up
            _tl.store(output_ptr + offsets, result, mask=mask)

        return kernel
    except OSError:
        logger.warning(
            "Triton fused_silu_mul kernel compilation skipped: "
            "no .py source available (pyc-only install). "
            "Using PyTorch fallback."
        )
        return None


# =============================================================================
# Fused RMSNorm Kernel
# =============================================================================


class FusedRMSNormKernel(Kernel):
    """Fused RMS Normalization kernel.

    Implements RMSNorm in a single fused operation:
        output = (x / sqrt(mean(x^2) + eps)) * weight

    Fusing the normalization:
    - Reduces memory reads/writes
    - Single pass over the data
    - Can be 2-3x faster than naive implementation

    Example:
        kernel = FusedRMSNormKernel()
        output = kernel(x, weight, eps=1e-6)
    """

    @property
    def name(self) -> str:
        """Kernel name."""
        return "fused_rmsnorm"

    def forward(
        self,
        x: torch.Tensor,
        weight: torch.Tensor,
        eps: float = 1e-6,
        *,
        residual: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """Execute fused RMSNorm.

        Args:
            x: Input tensor of shape (..., hidden_size).
            weight: Weight tensor of shape (hidden_size,).
            eps: Epsilon for numerical stability.
            residual: Optional residual to add before normalization.
                     If provided, returns (output, residual_out).

        Returns:
            Normalized tensor, or (normalized, residual_out) if residual provided.
        """
        if _TRITON_AVAILABLE and x.is_cuda and residual is None:
            return self._triton_fused_rmsnorm(x, weight, eps)
        else:
            return self._pytorch_fused_rmsnorm(x, weight, eps, residual)

    def _pytorch_fused_rmsnorm(
        self,
        x: torch.Tensor,
        weight: torch.Tensor,
        eps: float,
        residual: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """PyTorch fallback implementation.

        Args:
            x: Input tensor.
            weight: Weight tensor.
            eps: Epsilon.
            residual: Optional residual.

        Returns:
            Normalized tensor(s).
        """
        if residual is not None:
            x = x + residual
            residual_out = x

        # Compute RMSNorm
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        x_normalized = x * torch.rsqrt(variance + eps)
        output = x_normalized * weight

        if residual is not None:
            return output, residual_out
        return output

    def _triton_fused_rmsnorm(
        self,
        x: torch.Tensor,
        weight: torch.Tensor,
        eps: float,
    ) -> torch.Tensor:
        """Triton optimized implementation.

        Falls back to PyTorch if the Triton kernel cannot be compiled
        (e.g. pyc-only install where inspect.getsourcelines has no .py to read).

        Args:
            x: Input tensor.
            weight: Weight tensor.
            eps: Epsilon.

        Returns:
            Normalized tensor.
        """
        kernel = _get_fused_rmsnorm_kernel()
        if kernel is None:
            return self._pytorch_fused_rmsnorm(x, weight, eps, None)  # type: ignore[return-value]

        # Reshape to 2D for kernel
        orig_shape = x.shape
        x_2d = x.view(-1, orig_shape[-1])
        output = torch.empty_like(x_2d)
        n_rows, n_cols = x_2d.shape

        kernel[(n_rows,)](
            x_2d,
            weight,
            output,
            n_cols,
            eps,
            BLOCK_SIZE=triton.next_power_of_2(n_cols),
        )
        return output.view(orig_shape)


@functools.lru_cache(maxsize=None)
def _get_fused_rmsnorm_kernel():
    """Lazily compile the Triton fused-RMSNorm kernel on first call.

    Defers @triton.jit decoration to call time so that module import
    succeeds even in pyc-only installs (where inspect.getsourcelines
    would raise OSError at decoration time).

    Returns:
        Compiled JITFunction, or None if Triton is unavailable or the
        source cannot be read (pyc-only install).
    """
    if not _TRITON_AVAILABLE:
        return None
    try:
        import triton as _triton
        import triton.language as _tl

        @_triton.jit
        def kernel(
            x_ptr,
            weight_ptr,
            output_ptr,
            n_cols,
            eps,
            BLOCK_SIZE: _tl.constexpr,
        ):
            """Triton kernel for RMSNorm."""
            row_idx = _tl.program_id(0)
            col_offsets = _tl.arange(0, BLOCK_SIZE)
            mask = col_offsets < n_cols
            row_start = row_idx * n_cols
            x = _tl.load(x_ptr + row_start + col_offsets, mask=mask, other=0.0)
            weight = _tl.load(weight_ptr + col_offsets, mask=mask, other=0.0)
            x_sq = x * x
            variance = _tl.sum(x_sq, axis=0) / n_cols
            rstd = 1.0 / _tl.sqrt(variance + eps)
            x_norm = x * rstd
            output = x_norm * weight
            _tl.store(output_ptr + row_start + col_offsets, output, mask=mask)

        return kernel
    except OSError:
        logger.warning(
            "Triton fused_rmsnorm kernel compilation skipped: "
            "no .py source available (pyc-only install). "
            "Using PyTorch fallback."
        )
        return None


# =============================================================================
# Fused Add + RMSNorm Kernel
# =============================================================================


class FusedAddRMSNormKernel(Kernel):
    """Fused residual add + RMS normalization.

    Implements: output = rmsnorm(x + residual)

    This is commonly used after attention and MLP blocks:
        residual = x
        x = attention(x)
        x = rmsnorm(x + residual)  # This fused op

    Fusing saves one full tensor read/write of the residual.

    Example:
        kernel = FusedAddRMSNormKernel()
        output, new_residual = kernel(x, residual, weight, eps=1e-6)
    """

    @property
    def name(self) -> str:
        """Kernel name."""
        return "fused_add_rmsnorm"

    def forward(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
        weight: torch.Tensor,
        eps: float = 1e-6,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Execute fused add + RMSNorm.

        Args:
            x: Input tensor.
            residual: Residual tensor to add.
            weight: Normalization weight.
            eps: Epsilon for numerical stability.

        Returns:
            Tuple of (normalized_output, updated_residual).
        """
        if _TRITON_AVAILABLE and x.is_cuda:
            return self._triton_fused_add_rmsnorm(x, residual, weight, eps)
        else:
            return self._pytorch_fused_add_rmsnorm(x, residual, weight, eps)

    def _pytorch_fused_add_rmsnorm(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
        weight: torch.Tensor,
        eps: float,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """PyTorch fallback implementation.

        Args:
            x: Input tensor.
            residual: Residual tensor.
            weight: Weight tensor.
            eps: Epsilon.

        Returns:
            (normalized_output, updated_residual).
        """
        # Add residual
        residual = x + residual

        # RMSNorm
        variance = residual.pow(2).mean(dim=-1, keepdim=True)
        residual_normalized = residual * torch.rsqrt(variance + eps)
        output = residual_normalized * weight

        return output, residual

    def _triton_fused_add_rmsnorm(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
        weight: torch.Tensor,
        eps: float,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Triton optimized implementation.

        Note: Currently falls back to PyTorch as the Triton kernel
        would need careful memory management for the residual output.

        Args:
            x: Input tensor.
            residual: Residual tensor.
            weight: Weight tensor.
            eps: Epsilon.

        Returns:
            (normalized_output, updated_residual).
        """
        # TODO: Implement Triton kernel with proper residual handling
        return self._pytorch_fused_add_rmsnorm(x, residual, weight, eps)


# =============================================================================
# Fused QKV Projection Kernel
# =============================================================================


class FusedQKVProjectionKernel(Kernel):
    """Fused Query-Key-Value projection.

    Combines three separate linear projections into one:
        q = x @ W_q
        k = x @ W_k
        v = x @ W_v

    Into a single matmul:
        qkv = x @ W_qkv
        q, k, v = split(qkv)

    Benefits:
    - Single matmul instead of three
    - Better GPU utilization
    - Reduced kernel launch overhead

    Example:
        kernel = FusedQKVProjectionKernel()
        q, k, v = kernel(hidden_states, qkv_weight, num_heads, head_dim)
    """

    @property
    def name(self) -> str:
        """Kernel name."""
        return "fused_qkv_projection"

    def forward(
        self,
        hidden_states: torch.Tensor,
        qkv_weight: torch.Tensor,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        *,
        qkv_bias: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Execute fused QKV projection.

        Args:
            hidden_states: Input tensor (batch, seq_len, hidden_size).
            qkv_weight: Fused QKV weight (q_dim + k_dim + v_dim, hidden_size).
            num_heads: Number of query heads.
            num_kv_heads: Number of key/value heads.
            head_dim: Dimension per head.
            qkv_bias: Optional fused QKV bias.

        Returns:
            Tuple of (query, key, value) tensors.
        """
        batch_size, seq_len, _ = hidden_states.shape

        # Single fused matmul
        qkv = F.linear(hidden_states, qkv_weight, qkv_bias)

        # Split into Q, K, V
        q_dim = num_heads * head_dim
        kv_dim = num_kv_heads * head_dim

        q, k, v = qkv.split([q_dim, kv_dim, kv_dim], dim=-1)

        # Reshape for attention
        q = q.view(batch_size, seq_len, num_heads, head_dim)
        k = k.view(batch_size, seq_len, num_kv_heads, head_dim)
        v = v.view(batch_size, seq_len, num_kv_heads, head_dim)

        return q, k, v


# =============================================================================
# Fused Rotary Position Embedding Kernel
# =============================================================================


class FusedRotaryEmbeddingKernel(Kernel):
    """Fused Rotary Position Embedding (RoPE) kernel.

    Applies rotary position embeddings to query and key tensors
    in a single fused operation.

    RoPE: x' = x * cos(θ) + rotate_half(x) * sin(θ)

    Example:
        kernel = FusedRotaryEmbeddingKernel()
        q_rot, k_rot = kernel(query, key, cos, sin, position_ids)
    """

    @property
    def name(self) -> str:
        """Kernel name."""
        return "fused_rope"

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        position_ids: torch.Tensor | None = None,
        *,
        unsqueeze_dim: int = 1,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Apply rotary embeddings to Q and K.

        Args:
            query: Query tensor (..., seq_len, num_heads, head_dim).
            key: Key tensor (..., seq_len, num_kv_heads, head_dim).
            cos: Cosine cache tensor.
            sin: Sine cache tensor.
            position_ids: Optional position indices.
            unsqueeze_dim: Dimension to unsqueeze cos/sin for broadcasting.

        Returns:
            Tuple of (rotated_query, rotated_key).
        """
        if position_ids is not None:
            # Gather cos/sin at specific positions
            cos = cos[position_ids].unsqueeze(unsqueeze_dim)
            sin = sin[position_ids].unsqueeze(unsqueeze_dim)
        else:
            cos = cos.unsqueeze(unsqueeze_dim)
            sin = sin.unsqueeze(unsqueeze_dim)

        # Apply rotary embedding to both Q and K
        q_embed = self._apply_rotary(query, cos, sin)
        k_embed = self._apply_rotary(key, cos, sin)

        return q_embed, k_embed

    def _apply_rotary(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
    ) -> torch.Tensor:
        """Apply rotary embedding to a single tensor.

        Args:
            x: Input tensor.
            cos: Cosine values.
            sin: Sine values.

        Returns:
            Rotated tensor.
        """
        # Rotate half
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        x_rotated = torch.cat((-x2, x1), dim=-1)

        # Apply rotation
        return (x * cos) + (x_rotated * sin)


# =============================================================================
# Registry functions
# =============================================================================


def get_fused_kernel(name: str) -> Kernel:
    """Get a fused kernel by name.

    Args:
        name: Kernel name.

    Returns:
        Kernel instance.

    Raises:
        KeyError: If kernel not found.
    """
    kernels = {
        "fused_silu_mul": FusedSiLUMulKernel,
        "fused_rmsnorm": FusedRMSNormKernel,
        "fused_add_rmsnorm": FusedAddRMSNormKernel,
        "fused_qkv_projection": FusedQKVProjectionKernel,
        "fused_rope": FusedRotaryEmbeddingKernel,
    }

    if name not in kernels:
        raise KeyError(f"Fused kernel '{name}' not found. Available: {list(kernels.keys())}")

    return kernels[name]()


def list_fused_kernels() -> list[str]:
    """List all available fused kernels.

    Returns:
        List of kernel names.
    """
    return [
        "fused_silu_mul",
        "fused_rmsnorm",
        "fused_add_rmsnorm",
        "fused_qkv_projection",
        "fused_rope",
    ]
