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

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

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 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);

    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,
            std::ptr::null_mut(),
        );
    }

    (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,
            std::ptr::null_mut(),
        );
    }

    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,
            std::ptr::null_mut(),
        );
    }
}

/// 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,
            std::ptr::null_mut(),
        );
    }

    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, std::ptr::null_mut());
        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
}