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,
                xserv_cuda::current_stream_raw(),
            );
        }

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

    /// YaRN (Yet another RoPE extensioN) RoPE cache. Applies frequency-dependent
    /// interpolation so the model can extrapolate beyond its training context.
    pub fn new_yarn(
        max_seq_len: usize,
        head_dim: usize,
        theta: f64,
        factor: f64,
        original_max_pos: usize,
        beta_fast: f64,
        beta_slow: f64,
    ) -> Self {
        let half_dim = head_dim / 2;
        let dim = head_dim as f64;

        // find_correction_dim: inverse formula to find dimension from number of rotations
        let find_correction_dim = |num_rotations: f64| -> f64 {
            dim * (original_max_pos as f64 / (num_rotations * 2.0 * std::f64::consts::PI)).ln()
                / (2.0 * theta.ln())
        };

        let low_raw = find_correction_dim(beta_fast);
        let high_raw = find_correction_dim(beta_slow);
        // config has truncate=false, so use raw values (no floor/ceil)
        let low = low_raw.max(0.0);
        let high = high_raw.min((half_dim - 1) as f64);

        // Compute inv_freq with YaRN interpolation
        let mut inv_freq = vec![0.0f64; half_dim];
        for i in 0..half_dim {
            let pos_freq = theta.powf((2 * i) as f64 / dim);
            let inv_freq_extrapolation = 1.0 / pos_freq; // original
            let inv_freq_interpolation = 1.0 / (factor * pos_freq); // scaled

            // Linear ramp: 0 where we keep original, 1 where we interpolate
            let ramp = if (high - low).abs() < 0.001 {
                0.5
            } else {
                ((i as f64 - low) / (high - low)).clamp(0.0, 1.0)
            };
            let extrapolation_factor = 1.0 - ramp;

            inv_freq[i] = inv_freq_interpolation * (1.0 - extrapolation_factor)
                + inv_freq_extrapolation * extrapolation_factor;
        }

        // Attention scaling factor for YaRN: 0.1 * ln(factor) + 1.0
        let attn_factor = 0.1 * factor.ln() + 1.0;

        // Build cos/sin cache on CPU then upload
        let total = max_seq_len * half_dim;
        let mut cos_host = vec![0.0f32; total];
        let mut sin_host = vec![0.0f32; total];
        for pos in 0..max_seq_len {
            for i in 0..half_dim {
                let angle = pos as f64 * inv_freq[i];
                cos_host[pos * half_dim + i] = (angle.cos() * attn_factor) as f32;
                sin_host[pos * half_dim + i] = (angle.sin() * attn_factor) as f32;
            }
        }

        let nbytes = total * std::mem::size_of::<f32>();
        let mut cos = GpuBuffer::alloc(nbytes).expect("alloc yarn cos_cache");
        let mut sin = GpuBuffer::alloc(nbytes).expect("alloc yarn sin_cache");
        let cos_bytes =
            unsafe { std::slice::from_raw_parts(cos_host.as_ptr() as *const u8, nbytes) };
        let sin_bytes =
            unsafe { std::slice::from_raw_parts(sin_host.as_ptr() as *const u8, nbytes) };
        cos.copy_from_host(cos_bytes).unwrap();
        sin.copy_from_host(sin_bytes).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 =
        xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
    pos_gpu.copy_from_host(pos_bytes).unwrap();

    rope_inplace_device_pos(x, cache, pos_gpu.as_ptr() as *const c_void);
}

/// RoPE in-place with positions already on the GPU (u32, [num_tokens]).
/// Used by the CUDA-graph decode path, where the position lives in a
/// persistent device buffer updated outside the captured region.
pub fn rope_inplace_device_pos(x: &Tensor, cache: &RopeCache, pos_gpu: *const c_void) {
    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);

    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,
                num_tokens as i32,
                num_heads as i32,
                head_dim as i32,
                xserv_cuda::current_stream_raw(),
            ),
            DType::BF16 => launch_rope_bf16(
                x.data_ptr() as *mut c_void,
                cache.cos.as_ptr() as _,
                cache.sin.as_ptr() as _,
                pos_gpu,
                num_tokens as i32,
                num_heads as i32,
                head_dim as i32,
                xserv_cuda::current_stream_raw(),
            ),
            _ => panic!("unsupported dtype for rope"),
        }
    }
}