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

import os
from typing import TYPE_CHECKING

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

from .checkpoint import H3Checkpoint
from .nvfp4 import Nvfp4Linear

if TYPE_CHECKING:
    from .distributed import SequenceParallelContext


AVAILABLE_BACKENDS = ("sage2", "cudnn_sdpa", "ck_int8", "sdpa", "flash4", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn")
PLANNED_BACKENDS = ("easycache", "h3_cache", "kj_chunked_ffn")
DEFAULT_ATTENTION_BACKEND = os.getenv("H3_DEFAULT_ATTENTION", "sage2")


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({"cudnn_sdpa": "available: forced cuDNN SDPA with no backend fallback"})
    status.update({"ck_int8": "available: approximate Comfy Kitchen INT8 Q/K/V attention"})
    status.update({"sol_attn": "experimental: sparse Triton attention for eligible non-causal H3 attention calls; falls back below H3_SOL_MIN_TOKENS unless H3_SOL_STRICT=1"})
    status.update({"flash4": "available: official FlashAttention-4 CuTeDSL Blackwell kernel (strict, no fallback)"})
    status.update({"easycache": "planned: approximate denoiser cache"})
    status.update({"h3_cache": "planned: approximate H3-specific cache"})
    status.update({"kj_chunked_ffn": "available: exact H3 MLP row chunking via H3_MLP_CHUNKS or runtime args"})
    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 run_sol_attention_bshd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor:
    """Run Sol-Attn on `[batch, sequence, heads, dim]` tensors and return the same layout."""
    try:
        from sol_kernel import sol_attn

        tau = float(os.getenv("H3_SOL_TAU", "1.3"))
        min_tokens = int(os.getenv("H3_SOL_MIN_TOKENS", "4096"))
        thresh_type = os.getenv("H3_SOL_THRESH_TYPE", "diag")
        int8_qk = os.getenv("H3_SOL_INT8_QK", "").lower() in {"1", "true", "yes", "on"}
        int8_pv = os.getenv("H3_SOL_INT8_PV", "").lower() in {"1", "true", "yes", "on"}
        if is_causal:
            raise ValueError("Sol-Attn backend only supports non-causal H3 attention")
        if q.shape[-1] != 128:
            raise ValueError(f"Sol-Attn requires head dim 128, got {q.shape[-1]}")
        if q.shape[1] < min_tokens:
            raise ValueError(f"{q.shape[1]} tokens < H3_SOL_MIN_TOKENS={min_tokens}")
        return sol_attn(
            q.contiguous(),
            k.contiguous(),
            v.contiguous(),
            tau=tau,
            thresh_type=thresh_type,
            int8_qk=int8_qk,
            int8_pv=int8_pv,
        )
    except Exception:
        if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}:
            raise
        fallback = os.getenv("H3_SOL_FALLBACK", "sage2")
        hnd = run_attention(q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), backend=fallback, is_causal=is_causal)
        return hnd.transpose(1, 2)


def qkv_to_bshd(qkv: torch.Tensor, heads: int, head_dim: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Split `[S, 3*H*D]` QKV into contiguous Sol-native `[1,S,H,D]` tensors."""
    if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() != "fused":
        raise RuntimeError("fused QKV layout disabled")
    try:
        from .nvfp4_quant import _vortex_scale_extension

        return tuple(_vortex_scale_extension().qkv_to_bshd(qkv, heads, head_dim))
    except Exception:
        if os.getenv("H3_SOL_QKV_LAYOUT_STRICT", "").lower() in {"1", "true", "yes", "on"}:
            raise
        inner = heads * head_dim
        sequence = qkv.shape[0]
        q, k, v = qkv.split(inner, dim=-1)
        return (
            q.view(1, sequence, heads, head_dim).contiguous(),
            k.view(1, sequence, heads, head_dim).contiguous(),
            v.view(1, sequence, heads, head_dim).contiguous(),
        )


def run_flash4_attention_bshd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor:
    """Run official FlashAttention-4 on `[batch, sequence, heads, dim]` tensors."""
    if not q.is_cuda or not k.is_cuda or not v.is_cuda:
        raise ValueError("FlashAttention-4 requires CUDA tensors")
    if q.dtype not in {torch.float16, torch.bfloat16} or k.dtype != q.dtype or v.dtype != q.dtype:
        raise ValueError("FlashAttention-4 requires matching FP16 or BF16 Q/K/V tensors")
    if q.shape != k.shape or q.shape != v.shape:
        raise ValueError("FlashAttention-4 requires matching Q/K/V shapes")
    if q.shape[-1] != 128:
        raise ValueError(f"H3 FlashAttention-4 requires head dim 128, got {q.shape[-1]}")

    from flash_attn.cute import flash_attn_func

    result = flash_attn_func(q.contiguous(), k.contiguous(), v.contiguous(), causal=is_causal)
    return result[0] if isinstance(result, tuple) else result


def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend: str, is_causal: bool) -> torch.Tensor:
    """Run one `[batch, heads, sequence, dim]` attention operation."""
    if backend == "sol_attn":
        return run_sol_attention_bshd(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=is_causal).transpose(1, 2)
    if backend == "kj_head_sliced":
        head_slice_size = int(os.getenv("H3_HEAD_SLICE_SIZE", "8"))
        base_backend = os.getenv("H3_HEAD_SLICE_BACKEND", "sage2")
        if head_slice_size <= 0:
            raise ValueError("H3_HEAD_SLICE_SIZE must be positive")
        outputs = [
            run_attention(q[:, start:start + head_slice_size], k[:, start:start + head_slice_size], v[:, start:start + head_slice_size], backend=base_backend, is_causal=is_causal)
            for start in range(0, q.shape[1], head_slice_size)
        ]
        return torch.cat(outputs, dim=1)
    if backend == "sage2":
        from sageattention import sageattn

        return sageattn(q, k, v, is_causal=is_causal, tensor_layout="HND", smooth_k=False)
    if backend == "flash4":
        return run_flash4_attention_bshd(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=is_causal,
        ).transpose(1, 2)
    if backend == "cudnn_sdpa":
        from torch.nn.attention import SDPBackend, sdpa_kernel

        with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
            return functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
    if backend == "ck_int8":
        if is_causal:
            raise ValueError("Comfy Kitchen INT8 attention does not support causal H3 attention")
        import comfy_kitchen

        return comfy_kitchen.int8_attention(q, k, v)
    if backend == "sage3":
        from sageattn3 import sageattn3_blackwell

        return sageattn3_blackwell(q, k, v, is_causal=is_causal)
    if backend == "sage3_mean":
        from sageattn3 import sageattn3_blackwell

        return sageattn3_blackwell(q, k, v, is_causal=is_causal, per_block_mean=True)
    if backend == "kj_sage_cuda":
        from sageattention import sageattn_qk_int8_pv_fp16_cuda

        return sageattn_qk_int8_pv_fp16_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32", tensor_layout="HND")
    if backend == "kj_sage_triton":
        from sageattention import sageattn_qk_int8_pv_fp16_triton

        return sageattn_qk_int8_pv_fp16_triton(q, k, v, is_causal=is_causal, tensor_layout="HND")
    if backend == "kj_sage_fp8":
        from sageattention import sageattn_qk_int8_pv_fp8_cuda

        return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp32", tensor_layout="HND")
    if backend == "kj_sage_fp8pp":
        from sageattention import sageattn_qk_int8_pv_fp8_cuda

        return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp16", tensor_layout="HND")
    if backend == "sdpa":
        return functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
    raise ValueError(f"Unsupported H3 attention backend: {backend}")


def run_sage_attention_nhd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    """Run Sage2 directly on projection-strided NHD Q/K/V views."""
    from sageattention import sageattn

    return sageattn(q, k, v, is_causal=False, tensor_layout="NHD", smooth_k=False)


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)


def rms_rope_split_half_(
    q: torch.Tensor,
    k: torch.Tensor,
    rotation: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run Comfy Kitchen's standalone fused H3 Q/K normalization and RoPE."""
    import comfy_kitchen  # Registers the independent CUDA extension operators.

    del comfy_kitchen
    torch.ops.comfy_kitchen.rms_rope_split_half_(
        q, k, rotation, q_weight, k_weight, eps, rotation.shape[-3] * 2
    )
    return q, k


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 = DEFAULT_ATTENTION_BACKEND,
    ):
        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 = DEFAULT_ATTENTION_BACKEND):
        return cls(
            checkpoint.nvfp4_linear(f"{prefix}.qkv_proj", output_dtype=output_dtype, role="h3_attn_qkv"),
            checkpoint.nvfp4_linear(f"{prefix}.out_proj", output_dtype=output_dtype, role="h3_attn_out"),
            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,
        sequence_parallel: "SequenceParallelContext | None" = None,
        tensor_parallel: "SequenceParallelContext | None" = None,
        modulation: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        if x.ndim != 2:
            raise ValueError("H3 attention expects `[sequence, hidden]` input.")
        if sequence_parallel is not None and tensor_parallel is not None:
            raise ValueError("choose Ulysses sequence parallelism or tensor parallelism, not both")
        if tensor_parallel is not None:
            return self._forward_tensor_parallel(x, rope_rotation, tensor_parallel)
        sequence = x.shape[0]
        inner = self.heads * self.head_dim
        qkv = self.qkv_proj.forward_modulated(x, *modulation) if modulation is not None else self.qkv_proj(x)
        q, k, v = qkv.split(inner, dim=-1)
        q = q.view(1, sequence, self.heads, self.head_dim)
        k = k.view(1, sequence, self.heads, self.head_dim)
        v = v.view(1, sequence, self.heads, self.head_dim)

        q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps)
        if (
            sequence_parallel is None
            and self.backend == "sage2"
            and os.getenv("H3_SAGE_QKV_LAYOUT", "hnd").lower() == "strided_nhd"
        ):
            out = run_sage_attention_nhd(q, k, v)
            if not out.is_contiguous():
                raise RuntimeError("Sage2 NHD output must be contiguous for zero-copy output projection")
            return self.out_proj(out.reshape(sequence, inner))
        if sequence_parallel is not None:
            q, k, v = sequence_parallel.seq_to_heads(q, k, v)
            if self.backend == "sol_attn":
                out = run_sol_attention_bshd(q, k, v, is_causal=False)
            elif self.backend == "flash4":
                out = run_flash4_attention_bshd(q, k, v, is_causal=False)
            else:
                out = run_attention(
                    q.transpose(1, 2).contiguous(),
                    k.transpose(1, 2).contiguous(),
                    v.transpose(1, 2).contiguous(),
                    backend=self.backend,
                    is_causal=False,
                ).transpose(1, 2)
            local_out = sequence_parallel.heads_to_seq(out)
            return self.out_proj(local_out.reshape(sequence, inner))
        if self.backend == "sol_attn":
            if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() == "fused":
                q, k, v = qkv_to_bshd(qkv, self.heads, self.head_dim)
            else:
                q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
            out = run_sol_attention_bshd(q, k, v, is_causal=False)
            return self.out_proj(out.reshape(sequence, inner).contiguous())
        if self.backend == "flash4":
            out = run_flash4_attention_bshd(q, k, v, is_causal=False)
            return self.out_proj(out.reshape(sequence, inner).contiguous())

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

        out = run_attention(q, k, v, backend=self.backend, is_causal=False)
        return self.out_proj(out.transpose(1, 2).reshape(sequence, inner).contiguous())

    def _forward_tensor_parallel(
        self,
        local_x: torch.Tensor,
        local_rotation: torch.Tensor,
        context: "SequenceParallelContext",
    ) -> torch.Tensor:
        """Run local-head attention with column/row-parallel NVFP4 projections."""
        local_sequence = local_x.shape[0]
        full_x = context.all_gather_rows(local_x)
        full_rotation = context.all_gather_rows(local_rotation[0]).unsqueeze(0)
        inner = self.heads * self.head_dim
        q, k, v = self.qkv_proj(full_x).split(inner, dim=-1)
        q = q.view(1, context.sequence_length, self.heads, self.head_dim)
        k = k.view(1, context.sequence_length, self.heads, self.head_dim)
        v = v.view(1, context.sequence_length, self.heads, self.head_dim)
        q, k = rms_rope_split_half_(
            q, k, full_rotation, self.q_norm_weight, self.k_norm_weight, self.eps,
        )
        if self.backend == "sol_attn":
            out = run_sol_attention_bshd(q, k, v, is_causal=False)
        elif self.backend == "flash4":
            out = run_flash4_attention_bshd(q, k, v, is_causal=False)
        else:
            out = run_attention(
                q.transpose(1, 2).contiguous(),
                k.transpose(1, 2).contiguous(),
                v.transpose(1, 2).contiguous(),
                backend=self.backend,
                is_causal=False,
            ).transpose(1, 2)
        partial = self.out_proj(out.reshape(context.sequence_length, inner).contiguous())
        local_output = context.reduce_scatter_rows(partial)
        bias = getattr(self, "tensor_parallel_output_bias", None)
        if bias is not None:
            local_output = local_output + bias.to(local_output)
        if local_output.shape[0] != local_sequence:
            raise RuntimeError("tensor-parallel attention returned the wrong local token count")
        return local_output