"""Direct H3 self-attention using packed NVFP4 linears and SageAttention3."""

import torch
import torch.nn.functional as functional
from torch import nn

from .checkpoint import H3Checkpoint
from .nvfp4 import Nvfp4Linear


AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3")
PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "sol_attn", "kj_sage", "kj_chunked_ffn", "kj_head_sliced")


def attention_backend_status() -> dict[str, str]:
    """Report direct-runtime attention choices without importing ComfyUI nodes."""
    status = {name: "available" for name in AVAILABLE_BACKENDS}
    status.update({"flash4": "planned: exact Blackwell kernel adapter"})
    status.update({"easycache": "planned: approximate denoiser cache"})
    status.update({"h3_cache": "planned: approximate H3-specific cache"})
    status.update({"sol_attn": "experimental: prior H3-tested sparse Triton attention; standalone adapter pending"})
    status.update({"kj_sage": "experimental: prior H3-tested Sage patch; standalone adapter pending"})
    status.update({"kj_chunked_ffn": "planned: exact memory-lifetime adapter"})
    status.update({"kj_head_sliced": "planned: exact memory-lifetime adapter"})
    return status


def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight.to(x.dtype), eps)


def apply_split_half_rope(x: torch.Tensor, rotation: torch.Tensor) -> torch.Tensor:
    """Apply H3's split-half rotary table to `[batch, sequence, heads, dim]`."""
    rotated_width = rotation.shape[-3] * 2
    half = rotated_width // 2
    if rotated_width > x.shape[-1]:
        raise ValueError("RoPE rotation width exceeds the attention head dimension.")

    pair = torch.stack((x[..., :half], x[..., half:rotated_width]), dim=-1)
    pair = torch.matmul(rotation.to(x.dtype), pair.unsqueeze(-1)).squeeze(-1)
    return torch.cat((pair[..., 0], pair[..., 1], x[..., rotated_width:]), dim=-1)


class H3SageAttention(nn.Module):
    """One MiniMax H3 attention module, independent of ComfyUI and Raylight."""

    def __init__(
        self,
        qkv_proj: Nvfp4Linear,
        out_proj: Nvfp4Linear,
        q_norm_weight: torch.Tensor,
        k_norm_weight: torch.Tensor,
        *,
        heads: int = 56,
        head_dim: int = 128,
        eps: float = 1e-5,
        backend: str = "sage2",
    ):
        super().__init__()
        self.qkv_proj = qkv_proj
        self.out_proj = out_proj
        self.heads = heads
        self.head_dim = head_dim
        self.eps = eps
        if backend in PLANNED_BACKENDS:
            raise ValueError(f"H3 attention backend '{backend}' needs a standalone adapter and is not installed.")
        if backend not in AVAILABLE_BACKENDS:
            raise ValueError(f"Unsupported H3 attention backend: {backend}")
        self.backend = backend
        self.register_buffer("q_norm_weight", q_norm_weight, persistent=False)
        self.register_buffer("k_norm_weight", k_norm_weight, persistent=False)

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16, backend: str = "sage2"):
        return cls(
            checkpoint.nvfp4_linear(f"{prefix}.qkv_proj", output_dtype=output_dtype),
            checkpoint.nvfp4_linear(f"{prefix}.out_proj", output_dtype=output_dtype),
            checkpoint.tensor(f"{prefix}.q_norm.weight", dtype=output_dtype),
            checkpoint.tensor(f"{prefix}.k_norm.weight", dtype=output_dtype),
            backend=backend,
        )

    def forward(self, x: torch.Tensor, rope_rotation: torch.Tensor) -> torch.Tensor:
        if x.ndim != 2:
            raise ValueError("H3 attention expects `[sequence, hidden]` input.")
        sequence = x.shape[0]
        inner = self.heads * self.head_dim
        q, k, v = self.qkv_proj(x).split(inner, dim=-1)
        q = rms_norm(q.view(1, sequence, self.heads, self.head_dim), self.q_norm_weight, self.eps)
        k = rms_norm(k.view(1, sequence, self.heads, self.head_dim), self.k_norm_weight, self.eps)
        v = v.view(1, sequence, self.heads, self.head_dim)

        q = apply_split_half_rope(q, rope_rotation).transpose(1, 2).contiguous()
        k = apply_split_half_rope(k, rope_rotation).transpose(1, 2).contiguous()
        v = v.transpose(1, 2).contiguous()

        if self.backend == "sage2":
            from sageattention import sageattn

            out = sageattn(q, k, v, is_causal=False, tensor_layout="HND", smooth_k=False)
        elif self.backend == "sage3":
            from sageattn3 import sageattn3_blackwell

            out = sageattn3_blackwell(q, k, v, is_causal=False)
        else:
            out = functional.scaled_dot_product_attention(q, k, v, is_causal=False)
        return self.out_proj(out.transpose(1, 2).reshape(sequence, inner).contiguous())