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

unsafe extern "C" {
    fn launch_rope_f32(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
                       positions: *const c_void, num_tokens: i32, num_heads: i32,
                       head_dim: i32, stream: *mut c_void);
    fn launch_rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
                        positions: *const c_void, num_tokens: i32, num_heads: i32,
                        head_dim: i32, stream: *mut c_void);
    fn launch_compute_rope_cache(cos_cache: *mut c_void, sin_cache: *mut c_void,
                                 max_seq_len: i32, half_dim: i32, theta: f32,
                                 stream: *mut c_void);
}

pub struct RopeCache {
    pub cos: GpuBuffer,
    pub sin: GpuBuffer,
    pub max_seq_len: usize,
    pub half_dim: usize,
}

impl RopeCache {
    pub fn new(max_seq_len: usize, head_dim: usize, theta: f32) -> Self {
        let half_dim = head_dim / 2;
        let nbytes = max_seq_len * half_dim * std::mem::size_of::<f32>();
        let mut cos = GpuBuffer::alloc(nbytes).expect("alloc cos_cache");
        let mut sin = GpuBuffer::alloc(nbytes).expect("alloc sin_cache");

        unsafe {
            launch_compute_rope_cache(
                cos.as_mut_ptr() as _, sin.as_mut_ptr() as _,
                max_seq_len as i32, half_dim as i32, theta, std::ptr::null_mut(),
            );
        }
        xserv_cuda::device::synchronize().unwrap();

        Self { cos, sin, max_seq_len, half_dim }
    }
}

/// Apply RoPE in-place to x.
/// x: [num_tokens, num_heads, head_dim] on GPU
/// positions: [num_tokens] (u32 on CPU, will be uploaded)
pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
    assert_eq!(x.ndim(), 3);
    assert!(x.is_contiguous());
    assert!(matches!(x.device(), Device::Cuda(_)));
    let num_tokens = x.shape()[0];
    let num_heads = x.shape()[1];
    let head_dim = x.shape()[2];
    assert_eq!(head_dim / 2, cache.half_dim);
    assert_eq!(positions.len(), num_tokens);

    let pos_bytes = unsafe {
        std::slice::from_raw_parts(
            positions.as_ptr() as *const u8,
            num_tokens * std::mem::size_of::<u32>(),
        )
    };
    let mut pos_gpu = GpuBuffer::alloc(pos_bytes.len()).expect("alloc positions");
    pos_gpu.copy_from_host(pos_bytes).unwrap();

    unsafe {
        match x.dtype() {
            DType::F32 => launch_rope_f32(
                x.data_ptr() as *mut c_void,
                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
                pos_gpu.as_ptr() as _,
                num_tokens as i32, num_heads as i32, head_dim as i32,
                std::ptr::null_mut(),
            ),
            DType::BF16 => launch_rope_bf16(
                x.data_ptr() as *mut c_void,
                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
                pos_gpu.as_ptr() as _,
                num_tokens as i32, num_heads as i32, head_dim as i32,
                std::ptr::null_mut(),
            ),
            _ => panic!("unsupported dtype for rope"),
        }
    }
    xserv_cuda::device::synchronize().unwrap();
}