"""Rotary embedding dispatch (family "rotary").

Registered rows: the fused CUDA kernel (bf16 CUDA, inference-only) and the
torch complex-multiply fallback (autograd-safe).  Selection runs through
the generic dispatcher, so ``op_backend(rotary=...)`` and
``ASTR_OPS=rotary=torch`` work exactly like for attention.

Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
"""

import torch
from torch import Tensor

from astrai.extension.dispatch import (
    CallContext,
    ImplRecord,
    Spec,
    register_family,
    resolve,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.rotary import rotary_emb as _cuda_rotary

_SPEC_CUDA = Spec.cuda_device() & Spec.dtype_in(torch.bfloat16) & Spec.no_grad()


def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
    cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
    dtype = x.dtype
    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
    x_complex = torch.view_as_complex(x_)
    freqs_cis_complex = torch.complex(cos, sin).unsqueeze(-2)
    x_rotated = x_complex * freqs_cis_complex
    x_out = torch.view_as_real(x_rotated).flatten(-2)
    return x_out.to(dtype)


def _rotary_records() -> list:
    return [
        ImplRecord(
            family="rotary",
            name="cuda",
            obj=_cuda_rotary,
            spec=_SPEC_CUDA,
            available=lambda: is_available("rotary_emb"),
            priority=0,
        ),
        ImplRecord(
            family="rotary",
            name="torch",
            obj=_torch_apply,
            spec=Spec.always(),
            priority=99,
        ),
    ]


register_family("rotary", _rotary_records, lambda: _rotary_records()[-1])


def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
    """Apply rotary embedding to x.

    Args:
        x: [batch, seq_len, n_heads, head_dim] (bf16)
        freqs_cis: [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs

    Returns:
        [batch, seq_len, n_heads, head_dim] (bf16)
    """
    ctx = CallContext(
        family="rotary",
        dtype=x.dtype,
        device_cuda=x.is_cuda,
        grad_enabled=torch.is_grad_enabled(),
        raw=(x, freqs_cis),
    )
    return resolve("rotary", ctx).record.obj(x, freqs_cis)