//! CUDA Graph integration for batch=1 single-sequence decode.
//!
//! Uses a per-layer split graph approach:
//! - Pre-attention graph: RMSNorm + QKV projections + reshape + QK-norm + RoPE
//! - Ungraphed: KV cache append + decode attention (variable kv_len)
//! - Post-attention graph: merge_heads + O-proj + add_rmsnorm + FFN + residual
//! - Final graph: last RMSNorm + lm_head GEMV

use std::ffi::c_void;
use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
use xserv_kernels::dispatch;
use xserv_kernels::gemm::{cublas_handle, gemv_scratch_elems};

use crate::config::ModelConfig;
use crate::kv_cache::GpuKVCache;

/// Pre-allocated intermediate buffers for decode (batch=1).
/// All buffers have stable GPU addresses for CUDA Graph replay.
struct DecodeBuffers {
    // Hidden-size buffers: [1, hidden]
    x: GpuBuffer,           // running hidden state
    normed: GpuBuffer,      // rmsnorm output
    attn_out: GpuBuffer,    // attention output [1, num_heads, 1, head_dim]
    attn_merged: GpuBuffer, // merge_heads output [1, hidden]
    o_proj: GpuBuffer,      // O projection output [1, hidden]
    normed2: GpuBuffer,     // post-attn norm output [1, hidden]
    sum_out: GpuBuffer,     // add_rmsnorm sum output [1, hidden]
    down: GpuBuffer,        // down projection output [1, hidden]

    // QKV projection outputs
    q_proj: GpuBuffer, // [1, num_heads * head_dim]
    k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
    v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]

    // Reshaped: [1, H, 1, D]
    q_reshaped: GpuBuffer,
    k_reshaped: GpuBuffer,
    v_reshaped: GpuBuffer,

    // After QK-norm (same shape as reshaped)
    q_normed: GpuBuffer,
    k_normed: GpuBuffer,

    // RoPE transposed: [1, H, D]
    q_rope: GpuBuffer,
    k_rope: GpuBuffer,

    // After RoPE transpose back: [1, H, 1, D]
    q_final: GpuBuffer,
    k_final: GpuBuffer,

    // FFN intermediates
    gate: GpuBuffer,     // [1, intermediate]
    up: GpuBuffer,       // [1, intermediate]
    silu_out: GpuBuffer, // [1, intermediate]

    // GEMV fp32 scratch for deterministic K-block partials.
    fp32_hidden: GpuBuffer,       // for hidden-sized GEMV outputs
    fp32_q: GpuBuffer,            // for Q projection
    fp32_kv: GpuBuffer,           // for K/V projection
    fp32_intermediate: GpuBuffer, // for gate/up projections
    fp32_vocab: GpuBuffer,        // for lm_head

    // Token ID and position (GPU-resident, updated before replay)
    token_id_gpu: GpuBuffer, // 4 bytes (u32)
    position_gpu: GpuBuffer, // 4 bytes (u32)

    // Final output
    logits: GpuBuffer, // [1, vocab_size]
}

pub struct DecodeGraphState {
    stream: CudaStream,
    buffers: DecodeBuffers,

    // Per-layer graph pairs
    pre_attn_graphs: Vec<CudaGraph>,
    post_attn_graphs: Vec<CudaGraph>,
    final_graph: CudaGraph,

    captured: bool,

    // Model dimensions
    hidden: usize,
    num_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
    intermediate: usize,
    vocab_size: usize,
    num_layers: usize,
    eps: f32,
}

impl DecodeGraphState {
    pub fn new(config: &ModelConfig) -> Self {
        let hidden = config.hidden();
        let num_heads = config.num_heads();
        let num_kv_heads = config.num_kv_heads();
        let head_dim = config.head_dim();
        let intermediate = config.ffn_hidden();
        let vocab_size = config.vocab_size;
        let num_layers = config.num_layers();
        let eps = config.rms_norm_eps.unwrap_or(1e-6) as f32;
        let es = 2usize; // BF16 = 2 bytes

        let stream = CudaStream::new().expect("create CUDA stream for graph");

        let alloc = |size: usize| -> GpuBuffer {
            GpuBuffer::alloc(size).expect("alloc decode graph buffer")
        };

        let buffers = DecodeBuffers {
            x: alloc(hidden * es),
            normed: alloc(hidden * es),
            attn_out: alloc(num_heads * head_dim * es),
            attn_merged: alloc(hidden * es),
            o_proj: alloc(hidden * es),
            normed2: alloc(hidden * es),
            sum_out: alloc(hidden * es),
            down: alloc(hidden * es),

            q_proj: alloc(num_heads * head_dim * es),
            k_proj: alloc(num_kv_heads * head_dim * es),
            v_proj: alloc(num_kv_heads * head_dim * es),

            q_reshaped: alloc(num_heads * head_dim * es),
            k_reshaped: alloc(num_kv_heads * head_dim * es),
            v_reshaped: alloc(num_kv_heads * head_dim * es),

            q_normed: alloc(num_heads * head_dim * es),
            k_normed: alloc(num_kv_heads * head_dim * es),

            q_rope: alloc(num_heads * head_dim * es),
            k_rope: alloc(num_kv_heads * head_dim * es),

            q_final: alloc(num_heads * head_dim * es),
            k_final: alloc(num_kv_heads * head_dim * es),

            gate: alloc(intermediate * es),
            up: alloc(intermediate * es),
            silu_out: alloc(intermediate * es),

            fp32_hidden: alloc(
                gemv_scratch_elems(hidden, hidden).max(gemv_scratch_elems(intermediate, hidden))
                    * 4,
            ),
            fp32_q: alloc(gemv_scratch_elems(hidden, num_heads * head_dim) * 4),
            fp32_kv: alloc(gemv_scratch_elems(hidden, num_kv_heads * head_dim) * 4),
            fp32_intermediate: alloc(gemv_scratch_elems(hidden, intermediate) * 4),
            fp32_vocab: alloc(gemv_scratch_elems(hidden, vocab_size) * 4),

            token_id_gpu: alloc(4),
            position_gpu: alloc(4),

            logits: alloc(vocab_size * es),
        };

        let pre_attn_graphs = (0..num_layers).map(|_| CudaGraph::new()).collect();
        let post_attn_graphs = (0..num_layers).map(|_| CudaGraph::new()).collect();

        Self {
            stream,
            buffers,
            pre_attn_graphs,
            post_attn_graphs,
            final_graph: CudaGraph::new(),
            captured: false,
            hidden,
            num_heads,
            num_kv_heads,
            head_dim,
            intermediate,
            vocab_size,
            num_layers,
            eps,
        }
    }

    pub fn is_captured(&self) -> bool {
        self.captured
    }

    /// Capture all per-layer graphs. Called once after the first decode step.
    pub fn capture(
        &mut self,
        layers: &[LayerWeightPtrs],
        norm_weight: *const c_void,
        lm_head_wt: *const c_void,
        _embed_table: *const c_void,
        rope_cos: *const c_void,
        rope_sin: *const c_void,
    ) {
        let s = self.stream.as_raw();
        let h = self.hidden as i32;
        let nh = self.num_heads as i32;
        let nkv = self.num_kv_heads as i32;
        let hd = self.head_dim as i32;
        let inter = self.intermediate as i32;
        let vocab = self.vocab_size as i32;
        let eps = self.eps;

        let cublas = cublas_handle();

        // Set cuBLAS to use our stream
        unsafe {
            dispatch::set_cublas_stream(cublas, s);
        }

        for (l, lw) in layers.iter().enumerate() {
            // === Pre-attention graph ===
            self.pre_attn_graphs[l]
                .begin_capture(&self.stream)
                .expect("begin pre-attn capture");
            unsafe {
                // RMSNorm
                dispatch::rmsnorm_bf16(
                    self.buffers.x.as_ptr() as _,
                    lw.input_norm,
                    self.buffers.normed.as_mut_ptr() as _,
                    1,
                    h,
                    eps,
                    s,
                );

                // Q projection (GEMV)
                dispatch::gemv_bf16(
                    self.buffers.normed.as_ptr() as _,
                    lw.q_proj_wt,
                    self.buffers.q_proj.as_mut_ptr() as _,
                    self.buffers.fp32_q.as_mut_ptr() as _,
                    h,
                    nh * hd,
                    s,
                );

                // K projection (GEMV)
                dispatch::gemv_bf16(
                    self.buffers.normed.as_ptr() as _,
                    lw.k_proj_wt,
                    self.buffers.k_proj.as_mut_ptr() as _,
                    self.buffers.fp32_kv.as_mut_ptr() as _,
                    h,
                    nkv * hd,
                    s,
                );

                // V projection (GEMV)
                dispatch::gemv_bf16(
                    self.buffers.normed.as_ptr() as _,
                    lw.v_proj_wt,
                    self.buffers.v_proj.as_mut_ptr() as _,
                    self.buffers.fp32_kv.as_mut_ptr() as _,
                    h,
                    nkv * hd,
                    s,
                );

                // Reshape heads: [1, H*D] -> [1, H, 1, D]
                dispatch::reshape_heads_bf16(
                    self.buffers.q_proj.as_ptr() as _,
                    self.buffers.q_reshaped.as_mut_ptr() as _,
                    1,
                    nh,
                    hd,
                    s,
                );
                dispatch::reshape_heads_bf16(
                    self.buffers.k_proj.as_ptr() as _,
                    self.buffers.k_reshaped.as_mut_ptr() as _,
                    1,
                    nkv,
                    hd,
                    s,
                );
                dispatch::reshape_heads_bf16(
                    self.buffers.v_proj.as_ptr() as _,
                    self.buffers.v_reshaped.as_mut_ptr() as _,
                    1,
                    nkv,
                    hd,
                    s,
                );

                // QK norm (head-level rmsnorm: treat [1,H,1,D] as [H, D])
                dispatch::rmsnorm_bf16(
                    self.buffers.q_reshaped.as_ptr() as _,
                    lw.q_norm,
                    self.buffers.q_normed.as_mut_ptr() as _,
                    nh,
                    hd,
                    eps,
                    s,
                );
                dispatch::rmsnorm_bf16(
                    self.buffers.k_reshaped.as_ptr() as _,
                    lw.k_norm,
                    self.buffers.k_normed.as_mut_ptr() as _,
                    nkv,
                    hd,
                    eps,
                    s,
                );

                // Transpose for RoPE: [1,H,1,D] -> [1,H,D]
                dispatch::transpose_hsd_to_shd_bf16(
                    self.buffers.q_normed.as_ptr() as _,
                    self.buffers.q_rope.as_mut_ptr() as _,
                    1,
                    nh,
                    hd,
                    s,
                );
                dispatch::transpose_hsd_to_shd_bf16(
                    self.buffers.k_normed.as_ptr() as _,
                    self.buffers.k_rope.as_mut_ptr() as _,
                    1,
                    nkv,
                    hd,
                    s,
                );

                // RoPE (in-place, reads position_gpu)
                dispatch::rope_bf16(
                    self.buffers.q_rope.as_mut_ptr() as _,
                    rope_cos,
                    rope_sin,
                    self.buffers.position_gpu.as_ptr() as _,
                    1,
                    nh,
                    hd,
                    s,
                );
                dispatch::rope_bf16(
                    self.buffers.k_rope.as_mut_ptr() as _,
                    rope_cos,
                    rope_sin,
                    self.buffers.position_gpu.as_ptr() as _,
                    1,
                    nkv,
                    hd,
                    s,
                );

                // Transpose back: [1,H,D] -> [1,H,1,D]
                dispatch::transpose_shd_to_hsd_bf16(
                    self.buffers.q_rope.as_ptr() as _,
                    self.buffers.q_final.as_mut_ptr() as _,
                    1,
                    nh,
                    hd,
                    s,
                );
                dispatch::transpose_shd_to_hsd_bf16(
                    self.buffers.k_rope.as_ptr() as _,
                    self.buffers.k_final.as_mut_ptr() as _,
                    1,
                    nkv,
                    hd,
                    s,
                );
            }
            self.pre_attn_graphs[l]
                .end_capture(&self.stream)
                .expect("end pre-attn capture");

            // === Post-attention graph ===
            self.post_attn_graphs[l]
                .begin_capture(&self.stream)
                .expect("begin post-attn capture");
            unsafe {
                // Merge heads: [1,H,1,D] -> [1, hidden]
                // attn_out is written by ungraphed attention
                dispatch::merge_heads_bf16(
                    self.buffers.attn_out.as_ptr() as _,
                    self.buffers.attn_merged.as_mut_ptr() as _,
                    1,
                    nh,
                    hd,
                    s,
                );

                // O projection
                dispatch::gemv_bf16(
                    self.buffers.attn_merged.as_ptr() as _,
                    lw.o_proj_wt,
                    self.buffers.o_proj.as_mut_ptr() as _,
                    self.buffers.fp32_hidden.as_mut_ptr() as _,
                    nh * hd,
                    h,
                    s,
                );

                // Fused Add+RMSNorm: normed2 = rmsnorm(o_proj + x), sum_out = o_proj + x
                dispatch::add_rmsnorm_bf16(
                    self.buffers.o_proj.as_ptr() as _,
                    self.buffers.x.as_ptr() as _,
                    lw.post_norm,
                    self.buffers.normed2.as_mut_ptr() as _,
                    self.buffers.sum_out.as_mut_ptr() as _,
                    1,
                    h,
                    eps,
                    s,
                );

                // Gate projection
                dispatch::gemv_bf16(
                    self.buffers.normed2.as_ptr() as _,
                    lw.gate_proj_wt,
                    self.buffers.gate.as_mut_ptr() as _,
                    self.buffers.fp32_intermediate.as_mut_ptr() as _,
                    h,
                    inter,
                    s,
                );

                // Up projection
                dispatch::gemv_bf16(
                    self.buffers.normed2.as_ptr() as _,
                    lw.up_proj_wt,
                    self.buffers.up.as_mut_ptr() as _,
                    self.buffers.fp32_intermediate.as_mut_ptr() as _,
                    h,
                    inter,
                    s,
                );

                // Fused SiLU x Mul
                dispatch::silu_mul_bf16(
                    self.buffers.gate.as_ptr() as _,
                    self.buffers.up.as_ptr() as _,
                    self.buffers.silu_out.as_mut_ptr() as _,
                    inter,
                    s,
                );

                // Down projection
                dispatch::gemv_bf16(
                    self.buffers.silu_out.as_ptr() as _,
                    lw.down_proj_wt,
                    self.buffers.down.as_mut_ptr() as _,
                    self.buffers.fp32_hidden.as_mut_ptr() as _,
                    inter,
                    h,
                    s,
                );

                // x = sum_out + down (residual connection for next layer)
                dispatch::add_bf16(
                    self.buffers.sum_out.as_ptr() as _,
                    self.buffers.down.as_ptr() as _,
                    self.buffers.x.as_mut_ptr() as _,
                    h,
                    s,
                );
            }
            self.post_attn_graphs[l]
                .end_capture(&self.stream)
                .expect("end post-attn capture");
        }

        // === Final graph: norm + lm_head ===
        self.final_graph
            .begin_capture(&self.stream)
            .expect("begin final capture");
        unsafe {
            dispatch::rmsnorm_bf16(
                self.buffers.x.as_ptr() as _,
                norm_weight,
                self.buffers.normed.as_mut_ptr() as _,
                1,
                h,
                eps,
                s,
            );
            dispatch::gemv_bf16(
                self.buffers.normed.as_ptr() as _,
                lm_head_wt,
                self.buffers.logits.as_mut_ptr() as _,
                self.buffers.fp32_vocab.as_mut_ptr() as _,
                h,
                vocab,
                s,
            );
        }
        self.final_graph
            .end_capture(&self.stream)
            .expect("end final capture");

        // Reset cuBLAS back to null stream
        unsafe {
            dispatch::set_cublas_stream(cublas, std::ptr::null_mut());
        }

        self.captured = true;
    }

    /// Execute a single decode step using captured graphs.
    pub fn execute(
        &mut self,
        token_id: u32,
        position: u32,
        cache: &mut GpuKVCache,
        _layers: &[LayerWeightPtrs],
        embed_table: *const c_void,
        vocab_size: i32,
        hidden_size: i32,
    ) {
        assert!(self.captured, "must call capture() before execute()");
        let s = self.stream.as_raw();
        let nkv = self.num_kv_heads;
        let nh = self.num_heads;
        let hd = self.head_dim;
        let es = 2usize; // BF16

        // Upload token ID and position to fixed GPU buffers
        self.buffers
            .token_id_gpu
            .copy_from_host(&token_id.to_le_bytes())
            .unwrap();
        self.buffers
            .position_gpu
            .copy_from_host(&position.to_le_bytes())
            .unwrap();

        // Embedding (outside graph since token_id changes each step)
        unsafe {
            dispatch::embedding_bf16(
                embed_table,
                self.buffers.token_id_gpu.as_ptr() as _,
                self.buffers.x.as_mut_ptr() as _,
                1,
                hidden_size,
                vocab_size,
                s,
            );
        }

        for l in 0..self.num_layers {
            // Pre-attention graph (norm + QKV + reshape + QK-norm + RoPE)
            self.pre_attn_graphs[l]
                .launch(&self.stream)
                .expect("launch pre-attn graph");

            // Ungraphed: KV cache append
            // k_final shape: [1, num_kv_heads, 1, head_dim] (after RoPE pipeline)
            // v_reshaped shape: [1, num_kv_heads, 1, head_dim] (V skips RoPE)
            let pos = position as usize;

            let k_buf_size = nkv * hd * es;
            let v_buf_size = nkv * hd * es;
            let shape = [1usize, nkv, 1, hd];

            // Synchronize before accessing buffers for KV cache append
            self.stream.synchronize().expect("sync before kv cache");

            let k_view = unsafe {
                crate::kv_cache::tensor_from_gpu_buffer_pub(
                    GpuBuffer::borrow_raw(self.buffers.k_final.as_mut_ptr(), k_buf_size),
                    &shape,
                    xserv_tensor::DType::BF16,
                    0,
                )
            };
            let v_view = unsafe {
                crate::kv_cache::tensor_from_gpu_buffer_pub(
                    GpuBuffer::borrow_raw(self.buffers.v_reshaped.as_mut_ptr(), v_buf_size),
                    &shape,
                    xserv_tensor::DType::BF16,
                    0,
                )
            };
            cache.append(l, &k_view, &v_view, 1, pos);

            // Ungraphed: get full KV cache and run decode attention
            let (k_full, v_full) = cache.get_kv_len(l, pos + 1);
            let kv_len = (pos + 1) as i32;
            let scale = 1.0 / (hd as f32).sqrt();

            // Attention output written to attn_out (separate from q_final)
            unsafe {
                dispatch::decode_attention_bf16(
                    self.buffers.q_final.as_ptr() as _,
                    k_full.data_ptr() as _,
                    v_full.data_ptr() as _,
                    self.buffers.attn_out.as_mut_ptr() as _,
                    1,
                    nh as i32,
                    nkv as i32,
                    kv_len,
                    hd as i32,
                    scale,
                    s,
                );
            }

            // Synchronize before post-attention graph reads attn_out
            self.stream.synchronize().expect("sync before post-attn");

            // Post-attention graph (merge + O-proj + add_rmsnorm + FFN + residual)
            self.post_attn_graphs[l]
                .launch(&self.stream)
                .expect("launch post-attn graph");
        }

        // Final graph (norm + lm_head)
        self.final_graph
            .launch(&self.stream)
            .expect("launch final graph");

        // Sync to ensure logits are ready
        self.stream.synchronize().expect("sync after decode");
    }

    /// Get the logits buffer (for reading results after execute).
    pub fn logits_buffer(&self) -> &GpuBuffer {
        &self.buffers.logits
    }

    /// Invalidate captured graphs (e.g. when switching sequences).
    pub fn invalidate(&mut self) {
        self.captured = false;
        self.pre_attn_graphs = (0..self.num_layers).map(|_| CudaGraph::new()).collect();
        self.post_attn_graphs = (0..self.num_layers).map(|_| CudaGraph::new()).collect();
        self.final_graph = CudaGraph::new();
    }
}

unsafe impl Send for DecodeGraphState {}

/// Lightweight struct holding raw pointers to a layer's weight tensors.
/// Used to avoid passing the full model struct into the graph capture code.
pub struct LayerWeightPtrs {
    pub input_norm: *const c_void,
    pub q_proj_wt: *const c_void,
    pub k_proj_wt: *const c_void,
    pub v_proj_wt: *const c_void,
    pub o_proj_wt: *const c_void,
    pub q_norm: *const c_void,
    pub k_norm: *const c_void,
    pub post_norm: *const c_void,
    pub gate_proj_wt: *const c_void,
    pub up_proj_wt: *const c_void,
    pub down_proj_wt: *const c_void,
}

unsafe impl Send for LayerWeightPtrs {}
unsafe impl Sync for LayerWeightPtrs {}