use std::ffi::c_void;
use xserv_tensor::{DType, Tensor};

use crate::gemm::{CublasHandle, cublas_handle};

unsafe extern "C" {
    fn launch_moe_topk_softmax_bf16(
        router_logits: *const c_void,
        topk_ids: *mut c_void,
        topk_weights: *mut c_void,
        num_tokens: i32,
        num_experts: i32,
        top_k: i32,
        stream: *mut c_void,
    );
    fn launch_moe_replicate_bf16(
        x: *const c_void,
        x_rep: *mut c_void,
        num_tokens: i32,
        hidden: i32,
        local_experts: i32,
        stream: *mut c_void,
    );
    fn launch_moe_bias_add_3d_bf16(
        x: *mut c_void,
        bias: *const c_void,
        batch: i32,
        num_tokens: i32,
        dim: i32,
        stream: *mut c_void,
    );
    fn launch_moe_weighted_sum_bf16(
        expert_out: *const c_void,
        topk_ids: *const c_void,
        topk_weights: *const c_void,
        out: *mut c_void,
        num_tokens: i32,
        hidden: i32,
        top_k: i32,
        expert_start: i32,
        local_experts: i32,
        stream: *mut c_void,
    );

    fn launch_moe_sparse_gemv_bf16_bf16(
        x: *const c_void,
        w: *const c_void,
        topk_ids: *const c_void,
        y: *mut c_void,
        num_tokens: i32,
        n: i32,
        k: i32,
        top_k: i32,
        expert_start: i32,
        local_experts: i32,
        x_per_slot: i32,
        stream: *mut c_void,
    );

    fn launch_moe_sparse_gemv_fp8_bf16(
        x: *const c_void,
        w: *const c_void,
        w_scales: *const c_void,
        bias: *const c_void,
        topk_ids: *const c_void,
        y: *mut c_void,
        num_tokens: i32,
        n: i32,
        k: i32,
        top_k: i32,
        expert_start: i32,
        local_experts: i32,
        x_per_slot: i32,
        stream: *mut c_void,
    );
    fn launch_moe_sparse_gemv_mxfp4_bf16(
        x: *const c_void,
        w_packed: *const c_void,
        w_scales: *const c_void,
        bias: *const c_void,
        topk_ids: *const c_void,
        y: *mut c_void,
        num_tokens: i32,
        n: i32,
        k: i32,
        top_k: i32,
        expert_start: i32,
        local_experts: i32,
        x_per_slot: i32,
        stream: *mut c_void,
    );
    fn launch_moe_weighted_sum_sparse_bf16(
        down: *const c_void,
        topk_ids: *const c_void,
        topk_weights: *const c_void,
        out: *mut c_void,
        num_tokens: i32,
        hidden: i32,
        top_k: i32,
        expert_start: i32,
        local_experts: i32,
        stream: *mut c_void,
    );

    fn cublasGemmStridedBatchedEx(
        handle: CublasHandle,
        transa: i32,
        transb: i32,
        m: i32,
        n: i32,
        k: i32,
        alpha: *const c_void,
        a: *const c_void,
        a_type: i32,
        lda: i32,
        stride_a: i64,
        b: *const c_void,
        b_type: i32,
        ldb: i32,
        stride_b: i64,
        beta: *const c_void,
        c: *mut c_void,
        c_type: i32,
        ldc: i32,
        stride_c: i64,
        batch_count: i32,
        compute_type: i32,
        algo: i32,
    ) -> i32;

    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
}

const CUDA_R_16BF: i32 = 14;
const CUBLAS_COMPUTE_32F: i32 = 68;
const CUBLAS_GEMM_DEFAULT: i32 = -1;

/// GPU top-k selection + softmax over router logits.
///
/// Input:  router_logits [num_tokens, num_experts] BF16 on GPU
/// Output: (topk_ids [num_tokens, top_k] i32, topk_weights [num_tokens, top_k] f32)
pub fn moe_topk_softmax(
    router_logits: &Tensor,
    num_experts: usize,
    top_k: usize,
) -> (Tensor, Tensor) {
    assert_eq!(router_logits.ndim(), 2);
    assert_eq!(router_logits.dtype(), DType::BF16);
    assert!(router_logits.is_contiguous());
    let num_tokens = router_logits.shape()[0];
    assert_eq!(router_logits.shape()[1], num_experts);

    // NOTE: topk_ids actually holds i32 expert indices; DType has no I32, so
    // this is a raw 4-byte buffer mislabeled F32. Never read it as floats —
    // all consumers (weighted-sum / sparse GEMV kernels) cast to int*.
    let topk_ids = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
    let topk_weights = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());

    unsafe {
        launch_moe_topk_softmax_bf16(
            router_logits.data_ptr() as *const c_void,
            topk_ids.data_ptr() as *mut c_void,
            topk_weights.data_ptr() as *mut c_void,
            num_tokens as i32,
            num_experts as i32,
            top_k as i32,
            xserv_cuda::current_stream_raw(),
        );
    }

    (topk_ids, topk_weights)
}

/// Replicate x [num_tokens, hidden] → [local_experts, num_tokens, hidden].
pub fn moe_replicate(x: &Tensor, local_experts: usize) -> Tensor {
    assert_eq!(x.ndim(), 2);
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous());
    let num_tokens = x.shape()[0];
    let hidden = x.shape()[1];
    let out = Tensor::empty(
        &[local_experts, num_tokens, hidden],
        DType::BF16,
        x.device(),
    );

    unsafe {
        launch_moe_replicate_bf16(
            x.data_ptr() as *const c_void,
            out.data_ptr() as *mut c_void,
            num_tokens as i32,
            hidden as i32,
            local_experts as i32,
            xserv_cuda::current_stream_raw(),
        );
    }

    out
}

/// In-place 3D bias add: x [batch, num_tokens, dim] += bias [batch, dim].
pub fn moe_bias_add_3d(x: &Tensor, bias: &Tensor) {
    assert_eq!(x.ndim(), 3);
    assert_eq!(bias.ndim(), 2);
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous());
    let batch = x.shape()[0];
    let num_tokens = x.shape()[1];
    let dim = x.shape()[2];
    assert_eq!(bias.shape(), &[batch, dim]);

    unsafe {
        launch_moe_bias_add_3d_bf16(
            x.data_ptr() as *mut c_void,
            bias.data_ptr() as *const c_void,
            batch as i32,
            num_tokens as i32,
            dim as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
}

/// Weighted sum of expert outputs → [num_tokens, hidden].
///
/// expert_out:   [local_experts, num_tokens, hidden] BF16
/// topk_ids:     [num_tokens, top_k] i32 (global expert indices)
/// topk_weights: [num_tokens, top_k] f32
pub fn moe_weighted_sum(
    expert_out: &Tensor,
    topk_ids: &Tensor,
    topk_weights: &Tensor,
    expert_start: usize,
    local_experts: usize,
    top_k: usize,
) -> Tensor {
    assert_eq!(expert_out.ndim(), 3);
    assert_eq!(expert_out.dtype(), DType::BF16);
    let num_tokens = expert_out.shape()[1];
    let hidden = expert_out.shape()[2];

    let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, expert_out.device());

    unsafe {
        launch_moe_weighted_sum_bf16(
            expert_out.data_ptr() as *const c_void,
            topk_ids.data_ptr() as *const c_void,
            topk_weights.data_ptr() as *const c_void,
            out.data_ptr() as *mut c_void,
            num_tokens as i32,
            hidden as i32,
            top_k as i32,
            expert_start as i32,
            local_experts as i32,
            xserv_cuda::current_stream_raw(),
        );
    }

    out
}

/// Sparse MoE GEMV (BF16 weights): compute only routed experts.
/// x: [num_tokens, K] or [num_tokens * top_k, K], w_t: [local_experts, N, K].
pub fn moe_sparse_gemv_bf16(
    x: &Tensor,
    w_t: &Tensor,
    topk_ids: &Tensor,
    top_k: usize,
    expert_start: usize,
    local_experts: usize,
    x_per_slot: bool,
) -> Tensor {
    assert_eq!(x.ndim(), 2);
    assert_eq!(w_t.ndim(), 3);
    assert_eq!(x.dtype(), DType::BF16);
    assert_eq!(w_t.dtype(), DType::BF16);
    assert!(x.is_contiguous() && w_t.is_contiguous());
    let local = w_t.shape()[0];
    let n = w_t.shape()[1];
    let k = w_t.shape()[2];
    assert_eq!(local, local_experts);
    assert_eq!(x.shape()[1], k);
    let num_tokens = if x_per_slot {
        x.shape()[0] / top_k
    } else {
        x.shape()[0]
    };
    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
    unsafe {
        launch_moe_sparse_gemv_bf16_bf16(
            x.data_ptr() as *const c_void,
            w_t.data_ptr() as *const c_void,
            topk_ids.data_ptr() as *const c_void,
            y.data_ptr() as *mut c_void,
            num_tokens as i32,
            n as i32,
            k as i32,
            top_k as i32,
            expert_start as i32,
            local_experts as i32,
            if x_per_slot { 1 } else { 0 },
            xserv_cuda::current_stream_raw(),
        );
    }
    y
}

/// Sparse MoE GEMV (FP8 W8A16): compute only the routed experts.
///
/// x:        [num_tokens, K] BF16 (x_per_slot=false, gate_up) or
///           [num_tokens * top_k, K] BF16 (x_per_slot=true, down)
/// w_fp8_t:  [local_experts, N, K] FP8E4M3 (transposed weight layout)
/// w_scales: [local_experts] F32 per-expert scalar scales
/// bias:     [local_experts, N] BF16 (fused into the epilogue)
/// topk_ids: [num_tokens, top_k] i32 global expert ids (GPU)
///
/// Returns y [num_tokens, top_k, N] BF16. Slots routed to experts NOT
/// owned by this rank are left UNWRITTEN (uninitialized memory) — the
/// consumer must skip them (see moe_weighted_sum_sparse).
#[allow(clippy::too_many_arguments)]
pub fn moe_sparse_gemv_fp8(
    x: &Tensor,
    w_fp8_t: &Tensor,
    w_scales: &Tensor,
    bias: &Tensor,
    topk_ids: &Tensor,
    num_tokens: usize,
    top_k: usize,
    expert_start: usize,
    local_experts: usize,
    x_per_slot: bool,
) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous());
    assert_eq!(w_fp8_t.dtype(), DType::FP8E4M3);
    let n = w_fp8_t.shape()[1];
    let k = w_fp8_t.shape()[2];
    // The kernel reads weights as uint4 (16 FP8 values per lane) and would
    // silently skip a K%16 tail.
    assert_eq!(k % 16, 0, "sparse FP8 GEMV requires K % 16 == 0, got {k}");
    assert_eq!(x.shape()[x.ndim() - 1], k);
    assert_eq!(
        x.shape()[0],
        if x_per_slot {
            num_tokens * top_k
        } else {
            num_tokens
        }
    );

    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
    unsafe {
        launch_moe_sparse_gemv_fp8_bf16(
            x.data_ptr() as *const c_void,
            w_fp8_t.data_ptr() as *const c_void,
            w_scales.data_ptr() as *const c_void,
            bias.data_ptr() as *const c_void,
            topk_ids.data_ptr() as *const c_void,
            y.data_ptr() as *mut c_void,
            num_tokens as i32,
            n as i32,
            k as i32,
            top_k as i32,
            expert_start as i32,
            local_experts as i32,
            x_per_slot as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    y
}

/// Sparse MoE GEMV (MXFP4 W4A16): same contract as moe_sparse_gemv_fp8,
/// with packed 4-bit weights [E, N, K/2] + UE8M0 block scales [E, N, K/32].
#[allow(clippy::too_many_arguments)]
pub fn moe_sparse_gemv_mxfp4(
    x: &Tensor,
    w_packed: &Tensor,
    w_scales: &Tensor,
    bias: &Tensor,
    topk_ids: &Tensor,
    num_tokens: usize,
    top_k: usize,
    n: usize,
    k: usize,
    expert_start: usize,
    local_experts: usize,
    x_per_slot: bool,
) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous());
    // 32-element MXFP4 blocks, read as uint4 (32 nibbles) per lane.
    assert_eq!(k % 32, 0, "sparse MXFP4 GEMV requires K % 32 == 0, got {k}");
    assert_eq!(x.shape()[x.ndim() - 1], k);
    assert_eq!(
        x.shape()[0],
        if x_per_slot {
            num_tokens * top_k
        } else {
            num_tokens
        }
    );

    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
    unsafe {
        launch_moe_sparse_gemv_mxfp4_bf16(
            x.data_ptr() as *const c_void,
            w_packed.data_ptr() as *const c_void,
            w_scales.data_ptr() as *const c_void,
            bias.data_ptr() as *const c_void,
            topk_ids.data_ptr() as *const c_void,
            y.data_ptr() as *mut c_void,
            num_tokens as i32,
            n as i32,
            k as i32,
            top_k as i32,
            expert_start as i32,
            local_experts as i32,
            x_per_slot as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    y
}

/// Weighted sum over the slot axis of the sparse GEMV output.
///
/// down: [num_tokens, top_k, hidden] BF16 (non-local slots uninitialized
/// and skipped, never multiplied by zero — NaN * 0 = NaN).
pub fn moe_weighted_sum_sparse(
    down: &Tensor,
    topk_ids: &Tensor,
    topk_weights: &Tensor,
    expert_start: usize,
    local_experts: usize,
) -> Tensor {
    assert_eq!(down.ndim(), 3);
    assert_eq!(down.dtype(), DType::BF16);
    let num_tokens = down.shape()[0];
    let top_k = down.shape()[1];
    let hidden = down.shape()[2];

    let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, down.device());
    unsafe {
        launch_moe_weighted_sum_sparse_bf16(
            down.data_ptr() as *const c_void,
            topk_ids.data_ptr() as *const c_void,
            topk_weights.data_ptr() as *const c_void,
            out.data_ptr() as *mut c_void,
            num_tokens as i32,
            hidden as i32,
            top_k as i32,
            expert_start as i32,
            local_experts as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    out
}

/// Strided batched GEMM for MoE expert forward.
/// C[b] = A[b] @ B[b]  for b in 0..batch
///
/// A: [batch, M, K] BF16 contiguous
/// B: [batch, K, N] BF16 contiguous
/// Returns C: [batch, M, N] BF16
#[allow(clippy::too_many_arguments)]
pub fn batched_gemm_strided(a: &Tensor, b: &Tensor) -> Tensor {
    assert_eq!(a.ndim(), 3);
    assert_eq!(b.ndim(), 3);
    assert_eq!(a.dtype(), DType::BF16);
    assert_eq!(b.dtype(), DType::BF16);
    assert!(a.is_contiguous() && b.is_contiguous());
    assert_eq!(a.shape()[0], b.shape()[0]);
    assert_eq!(a.shape()[2], b.shape()[1]);

    let batch = a.shape()[0];
    let m = a.shape()[1];
    let k = a.shape()[2];
    let n = b.shape()[2];

    let c = Tensor::empty(&[batch, m, n], DType::BF16, a.device());

    let alpha: f32 = 1.0;
    let beta: f32 = 0.0;

    // cuBLAS column-major: we compute C^T = B^T @ A^T
    // A is [batch, M, K] row-major → A^T is [K, M] col-major, lda=K
    // B is [batch, K, N] row-major → B^T is [N, K] col-major, ldb=N? No...
    //
    // Actually for row-major: A[M,K] in memory = col-major A^T[K,M] with lda=K.
    // So we call cublasGemmStridedBatchedEx with:
    //   transa=N, transb=N
    //   m=N, n=M, k=K  (because cuBLAS sees col-major)
    //   A_cublas = B_row (pointer), lda=N
    //   B_cublas = A_row (pointer), ldb=K
    //   C_cublas = C_row (pointer), ldc=N

    let stride_a = (m * k) as i64;
    let stride_b = (k * n) as i64;
    let stride_c = (m * n) as i64;

    let handle = cublas_handle();
    unsafe {
        cublasSetStream_v2(handle, xserv_cuda::current_stream_raw());
        let status = cublasGemmStridedBatchedEx(
            handle,
            0,
            0, // CUBLAS_OP_N, CUBLAS_OP_N
            n as i32,
            m as i32,
            k as i32,
            &alpha as *const f32 as *const c_void,
            b.data_ptr() as *const c_void,
            CUDA_R_16BF,
            n as i32,
            stride_b,
            a.data_ptr() as *const c_void,
            CUDA_R_16BF,
            k as i32,
            stride_a,
            &beta as *const f32 as *const c_void,
            c.data_ptr() as *mut c_void,
            CUDA_R_16BF,
            n as i32,
            stride_c,
            batch as i32,
            CUBLAS_COMPUTE_32F,
            CUBLAS_GEMM_DEFAULT,
        );
        assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
    }

    c
}