# SPDX-License-Identifier: MIT

from typing import List, Optional, Sequence, Tuple

import torch
from torch import Tensor

from ..jit.core import compile_ops


@compile_ops("module_grouped_gemm")
def ck_grouped_gemm(
    a_tensors: List[Tensor],
    b_tensors: List[Tensor],
) -> List[Tensor]: ...


@compile_ops("module_grouped_gemm")
def ck_grouped_gemm_out(
    a_tensors: List[Tensor],
    b_tensors: List[Tensor],
    c_tensors: List[Tensor],
) -> List[Tensor]: ...


# CK tile alignment for the low-level kernel (see grouped_gemm_kernels.cu).
_MOE_M_ALIGN = {
    torch.float16: 64,
    torch.bfloat16: 64,
    torch.float8_e4m3fn: 128,
    torch.int8: 32,
}
_MOE_NK_ALIGN = {
    torch.float16: dict(n=128, k=128),
    torch.bfloat16: dict(n=128, k=128),
    torch.float8_e4m3fn: dict(n=128, k=128),
    torch.int8: dict(n=32, k=128),
}


def _moe_output_dtype(dtype: torch.dtype) -> torch.dtype:
    if dtype is torch.int8:
        return torch.int32
    if dtype is torch.float8_e4m3fn:
        return torch.float32
    return dtype


def _align_up(x: int, align: int) -> int:
    return ((x + align - 1) // align) * align


def _validate_moe_fixed_nk(b_tensors: Sequence[Tensor], dtype: torch.dtype) -> Tuple[int, int]:
    if not b_tensors:
        raise ValueError("ck_grouped_gemm_moe: b_tensors must not be empty")
    n0, k0 = b_tensors[0].shape
    nk = _MOE_NK_ALIGN[dtype]
    if n0 % nk["n"] != 0 or k0 % nk["k"] != 0 or k0 < nk["k"]:
        raise ValueError(
            f"ck_grouped_gemm_moe: fixed N/K must satisfy N % {nk['n']} == 0, "
            f"K % {nk['k']} == 0, K >= {nk['k']} for {dtype}, got N={n0}, K={k0}"
        )
    for i, b in enumerate(b_tensors):
        if b.shape != (n0, k0):
            raise ValueError(
                f"ck_grouped_gemm_moe: all B tensors must share the same [N, K], "
                f"group {i} has {tuple(b.shape)} vs expected ({n0}, {k0})"
            )
    return n0, k0


def _pad_a_rows(a: Tensor, m_align: int) -> Tuple[Tensor, int, int]:
    m_orig = a.size(0)
    m_pad = _align_up(m_orig, m_align)
    if m_pad == m_orig:
        return a, m_orig, m_pad
    a_pad = a.new_zeros(m_pad, a.size(1))
    a_pad[:m_orig].copy_(a)
    return a_pad, m_orig, m_pad


def ck_grouped_gemm_moe(
    a_tensors: List[Tensor],
    b_tensors: List[Tensor],
) -> List[Tensor]:
    """
    MOE-friendly grouped GEMM with per-group dynamic M and fixed N/K.

    Each group computes C_i = A_i @ B_i^T. A_i may have arbitrary M_i >= 1;
    rows are zero-padded to the CK M-tile boundary before launch, then outputs
    are sliced back to the logical M_i.
    """
    if len(a_tensors) != len(b_tensors):
        raise ValueError("ck_grouped_gemm_moe: a and b tensor lists must have the same length")

    dtype = a_tensors[0].dtype
    m_align = _MOE_M_ALIGN[dtype]
    _validate_moe_fixed_nk(b_tensors, dtype)

    a_padded: List[Tensor] = []
    m_orig_list: List[int] = []
    for a, b in zip(a_tensors, b_tensors):
        if a.dtype != dtype or b.dtype != dtype:
            raise ValueError("ck_grouped_gemm_moe: all tensors must share the same dtype")
        if a.size(1) != b.size(1):
            raise ValueError("ck_grouped_gemm_moe: K mismatch between A and B")
        if a.size(0) <= 0:
            raise ValueError("ck_grouped_gemm_moe: M must be positive")
        a_pad, m_orig, _ = _pad_a_rows(a, m_align)
        a_padded.append(a_pad)
        m_orig_list.append(m_orig)

    c_padded = ck_grouped_gemm(a_padded, b_tensors)
    n = b_tensors[0].size(0)
    out_dtype = _moe_output_dtype(dtype)
    return [
        c[:m_orig, :n].to(out_dtype) if c.size(0) != m_orig else c
        for c, m_orig in zip(c_padded, m_orig_list)
    ]


def ck_grouped_gemm_moe_out(
    a_tensors: List[Tensor],
    b_tensors: List[Tensor],
    c_tensors: List[Tensor],
) -> List[Tensor]:
    """
    MOE grouped GEMM writing into caller-provided logical C tensors [M_i, N].

    Padded A/C buffers are allocated internally; only the valid M_i rows are
    copied into c_tensors.
    """
    if not (len(a_tensors) == len(b_tensors) == len(c_tensors)):
        raise ValueError("ck_grouped_gemm_moe_out: a, b, c lists must have the same length")

    dtype = a_tensors[0].dtype
    m_align = _MOE_M_ALIGN[dtype]
    n, _ = _validate_moe_fixed_nk(b_tensors, dtype)
    out_dtype = _moe_output_dtype(dtype)

    a_padded: List[Tensor] = []
    c_padded: List[Tensor] = []
    m_orig_list: List[int] = []

    for a, b, c in zip(a_tensors, b_tensors, c_tensors):
        if a.dtype != dtype or b.dtype != dtype:
            raise ValueError("ck_grouped_gemm_moe_out: a/b dtype mismatch")
        if c.dtype != out_dtype:
            raise ValueError(f"ck_grouped_gemm_moe_out: c dtype must be {out_dtype}")
        if a.size(1) != b.size(1):
            raise ValueError("ck_grouped_gemm_moe_out: K mismatch between A and B")
        m_orig = a.size(0)
        if c.shape != (m_orig, n):
            raise ValueError(
                f"ck_grouped_gemm_moe_out: c shape {tuple(c.shape)} != ({m_orig}, {n})"
            )
        a_pad, m_orig, m_pad = _pad_a_rows(a, m_align)
        a_padded.append(a_pad)
        m_orig_list.append(m_orig)
        if m_pad == m_orig:
            c_padded.append(c)
        else:
            c_padded.append(c.new_empty(m_pad, n))

    ck_grouped_gemm_out(a_padded, b_tensors, c_padded)

    for c, c_pad, m_orig in zip(c_tensors, c_padded, m_orig_list):
        if c_pad.data_ptr() != c.data_ptr():
            c.copy_(c_pad[:m_orig])
    return c_tensors


class GroupedGemmMoeBuffers:
    """
    Reusable padded A/C buffers for MOE inference with fixed N/K per expert.

    Avoids per-forward allocation when max tokens per expert is bounded.
    """

    def __init__(
        self,
        num_groups: int,
        n: int,
        k: int,
        dtype: torch.dtype,
        max_m: int,
        device: Optional[torch.device] = None,
    ):
        if num_groups <= 0:
            raise ValueError("GroupedGemmMoeBuffers: num_groups must be positive")
        nk = _MOE_NK_ALIGN[dtype]
        if n % nk["n"] != 0 or k % nk["k"] != 0 or k < nk["k"]:
            raise ValueError(f"GroupedGemmMoeBuffers: invalid fixed N={n}, K={k} for {dtype}")

        self.num_groups = num_groups
        self.n = n
        self.k = k
        self.dtype = dtype
        self.m_align = _MOE_M_ALIGN[dtype]
        self.max_m_pad = _align_up(max_m, self.m_align)
        self.out_dtype = _moe_output_dtype(dtype)
        dev = device or torch.device("cuda")

        self.a_bufs = [
            torch.zeros(self.max_m_pad, k, device=dev, dtype=dtype)
            for _ in range(num_groups)
        ]
        self.c_bufs = [
            torch.zeros(self.max_m_pad, n, device=dev, dtype=self.out_dtype)
            for _ in range(num_groups)
        ]

    def _ensure_capacity(self, m_orig: int) -> int:
        m_pad = _align_up(m_orig, self.m_align)
        if m_pad > self.max_m_pad:
            raise ValueError(
                f"GroupedGemmMoeBuffers: M={m_orig} exceeds configured max_m "
                f"(padded max {self.max_m_pad})"
            )
        return m_pad

    def run(
        self,
        a_tensors: Sequence[Tensor],
        b_tensors: Sequence[Tensor],
        c_tensors: Optional[Sequence[Tensor]] = None,
    ) -> List[Tensor]:
        if len(a_tensors) != self.num_groups or len(b_tensors) != self.num_groups:
            raise ValueError("GroupedGemmMoeBuffers: group count mismatch")

        a_padded: List[Tensor] = []
        c_padded: List[Tensor] = []
        m_orig_list: List[int] = []
        logical_c: List[Tensor] = []

        for i, (a, b) in enumerate(zip(a_tensors, b_tensors)):
            if b.shape != (self.n, self.k):
                raise ValueError(f"GroupedGemmMoeBuffers: B[{i}] shape {tuple(b.shape)} != ({self.n}, {self.k})")
            m_orig = a.size(0)
            m_pad = self._ensure_capacity(m_orig)
            m_orig_list.append(m_orig)

            a_buf = self.a_bufs[i]
            a_buf.zero_()
            a_buf[:m_orig].copy_(a)
            a_padded.append(a_buf[:m_pad])

            if c_tensors is not None:
                c = c_tensors[i]
                if c.shape != (m_orig, self.n):
                    raise ValueError(f"GroupedGemmMoeBuffers: c[{i}] shape mismatch")
                logical_c.append(c)
                c_padded.append(self.c_bufs[i][:m_pad])
            else:
                c_padded.append(self.c_bufs[i][:m_pad])

        if c_tensors is not None:
            ck_grouped_gemm_out(a_padded, list(b_tensors), c_padded)
            for c, c_pad, m_orig in zip(logical_c, c_padded, m_orig_list):
                c.copy_(c_pad[:m_orig])
            return list(logical_c)

        c_full = ck_grouped_gemm_out(a_padded, list(b_tensors), c_padded)
        return [c[:m_orig].clone() for c, m_orig in zip(c_full, m_orig_list)]
