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

unsafe extern "C" {
    fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
    fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
    fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
    fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
    fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
}

/// [S, H*D] → [1, H, S, D] on GPU (BF16)
pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
    unsafe {
        launch_reshape_heads_bf16(
            x.data_ptr() as _, out.data_ptr() as *mut c_void,
            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
        );
    }
    xserv_cuda::device::synchronize().unwrap();
    out
}

/// [1, H, S, D] → [S, H*D] on GPU (BF16)
pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let hidden = num_heads * head_dim;
    let out = Tensor::zeros(&[seq_len, hidden], DType::BF16, x.device());
    unsafe {
        launch_merge_heads_bf16(
            x.data_ptr() as _, out.data_ptr() as *mut c_void,
            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
        );
    }
    xserv_cuda::device::synchronize().unwrap();
    out
}

/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let out = Tensor::zeros(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
    unsafe {
        launch_transpose_hsd_to_shd_bf16(
            x.data_ptr() as _, out.data_ptr() as *mut c_void,
            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
        );
    }
    xserv_cuda::device::synchronize().unwrap();
    out
}

/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
    unsafe {
        launch_transpose_shd_to_hsd_bf16(
            x.data_ptr() as _, out.data_ptr() as *mut c_void,
            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
        );
    }
    xserv_cuda::device::synchronize().unwrap();
    out
}

/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
    if n_rep == 1 { return x.clone(); }
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let kv_heads = x.shape()[1];
    let seq_len = x.shape()[2];
    let head_dim = x.shape()[3];
    let new_heads = kv_heads * n_rep;
    let out = Tensor::zeros(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
    unsafe {
        launch_repeat_kv_bf16(
            x.data_ptr() as _, out.data_ptr() as *mut c_void,
            kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, std::ptr::null_mut(),
        );
    }
    xserv_cuda::device::synchronize().unwrap();
    out
}