import torch
import pytest

from aiter import dtypes
from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits


def _cdiv(x: int, y: int) -> int:
    return (x + y - 1) // y


def _kv_cache_cast_to_fp8_fn(x: torch.Tensor) -> torch.Tensor:
    num_blocks, block_size, num_heads, head_dim = x.shape
    assert num_heads == 1
    x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4)
    sf = x_amax / 240.0
    x_scaled = (x * (1.0 / sf)).to(dtypes.fp8)
    x_fp8 = torch.empty(
        (num_blocks, block_size * (head_dim + 4)),
        device=x.device,
        dtype=torch.uint8,
    )
    x_fp8[:, : block_size * head_dim] = x_scaled.view(num_blocks, block_size * head_dim).view(
        dtype=torch.uint8
    )
    x_fp8[:, block_size * head_dim : block_size * head_dim + 4 * block_size] = sf.view(
        num_blocks, block_size
    ).view(dtype=torch.uint8)
    return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4)


def _decode_packed_kv_cache(kv_cache_fp8: torch.Tensor, hidden_dim: int, block_size: int):
    num_blocks = kv_cache_fp8.size(0)
    flat = kv_cache_fp8.view(num_blocks, block_size * (hidden_dim + 4))
    kv_fp8 = flat[..., : block_size * hidden_dim].view(num_blocks, block_size, hidden_dim)
    kv_scale = flat[..., block_size * hidden_dim :].view(torch.float32).view(
        num_blocks, block_size, 1
    )
    keys = kv_fp8.view(dtypes.fp8).float() * kv_scale
    return keys


def _ref_logits_from_quantized_inputs(
    q_fp8: torch.Tensor,
    kv_cache_fp8: torch.Tensor,
    weights: torch.Tensor,
    context_lens: torch.Tensor,
    kv_indices: torch.Tensor,
    max_model_len: int,
):
    batch_size, next_n, heads, hidden_dim = q_fp8.size()
    block_size = kv_cache_fp8.size(1)
    logits = torch.full(
        (batch_size * next_n, max_model_len),
        float("-inf"),
        device=q_fp8.device,
        dtype=torch.float32,
    )
    q = q_fp8.float()
    keys_by_block = _decode_packed_kv_cache(kv_cache_fp8, hidden_dim, block_size)

    for b in range(batch_size):
        ctx_len = int(context_lens[b].item())
        for n in range(next_n):
            row = b * next_n + n
            q_offset = ctx_len - next_n + n
            for blk in range(_cdiv(ctx_len, block_size)):
                blk_idx = int(kv_indices[b, blk].item())
                kx = keys_by_block[blk_idx]
                k_offsets = torch.arange(
                    blk * block_size,
                    (blk + 1) * block_size,
                    device=q_fp8.device,
                    dtype=torch.int32,
                )
                scores = torch.matmul(q[b, n], kx.transpose(0, 1))
                weighted = torch.relu(scores) * weights[row].unsqueeze(1)
                s = weighted.sum(dim=0)
                valid = (k_offsets < ctx_len) & (k_offsets <= q_offset)
                logits[row, blk * block_size : (blk + 1) * block_size] = torch.where(
                    valid, s, torch.full_like(s, float("-inf"))
                )
    return logits


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
@pytest.mark.parametrize(
    "batch_size,next_n,heads,hidden_dim,max_model_len,context_lens_list,seed",
    [
        # Small irregular sanity cases
        (1, 1, 64, 128, 64, [59], 0),
        (2, 2, 64, 128, 128, [73, 121], 1),
        # Bench-like shape trend: next_n=2, batch in {1,2,4}
        (1, 2, 64, 128, 1024, [987], 2),
        (2, 2, 64, 128, 2048, [1536, 1987], 3),
        (4, 2, 64, 128, 1024, [211, 777, 943, 1001], 4),
        # Bench-like shape trend: next_n=1
        (1, 1, 64, 128, 4096, [3072], 5),
        (2, 1, 64, 128, 4096, [2049, 3991], 6),
        (4, 1, 64, 128, 2048, [401, 888, 1333, 1999], 7),
    ],
)
def test_deepgemm_fp8_paged_mqa_logits_correctness_fn(
    batch_size,
    next_n,
    heads,
    hidden_dim,
    max_model_len,
    context_lens_list,
    seed,
):
    torch.manual_seed(seed)
    block_size = 1
    num_blocks = max_model_len

    q = torch.randn((batch_size, next_n, heads, hidden_dim), device="cuda", dtype=torch.bfloat16)
    q_fp8 = q.to(dtypes.fp8)
    kv = torch.randn((num_blocks, block_size, 1, hidden_dim), device="cuda", dtype=torch.bfloat16)
    kv_fp8 = _kv_cache_cast_to_fp8_fn(kv)
    weights = torch.randn((batch_size * next_n, heads), device="cuda", dtype=torch.float32)
    context_lens = torch.tensor(context_lens_list, device="cuda", dtype=torch.int32)
    assert context_lens.numel() == batch_size
    assert torch.all(context_lens >= next_n)

    kv_indices = torch.zeros((batch_size, max_model_len), device="cuda", dtype=torch.int32)
    for b in range(batch_size):
        kv_indices[b, : context_lens[b]] = torch.randperm(max_model_len, device="cuda")[
            : context_lens[b]
        ]

    out = torch.full(
        (batch_size * next_n, max_model_len),
        float("-inf"),
        device="cuda",
        dtype=torch.float32,
    )

    deepgemm_fp8_paged_mqa_logits(
        q_fp8,
        kv_fp8,
        weights,
        out,
        context_lens,
        kv_indices,
        max_model_len,
    )
    ref = _ref_logits_from_quantized_inputs(
        q_fp8, kv_fp8, weights, context_lens, kv_indices, max_model_len
    )

    positions = torch.arange(max_model_len, device="cuda").unsqueeze(0)
    row_to_batch = torch.arange(batch_size * next_n, device="cuda") // next_n
    row_to_n = torch.arange(batch_size * next_n, device="cuda") % next_n
    valid_mask = positions <= (
        context_lens[row_to_batch] - next_n + row_to_n
    ).unsqueeze(1)

    assert torch.isfinite(out[valid_mask]).all()
    torch.testing.assert_close(out[valid_mask], ref[valid_mask], rtol=5e-2, atol=5e-2)
