//! EAGLE3 speculative draft head for Qwen3-8B (Phase 25).
//!
//! Loads the AngelSlim/Qwen3-8B_eagle3 pytorch_model.bin and provides a
//! single-step forward pass that takes 3 target hidden states + the previous
//! token and returns a draft token in the target vocabulary.
//!
//! Architecture (from weights):
//! - fc:           [hidden, 3*hidden] → fuse 3 target hidden states
//! - midlayer:     1 decoder layer (attn input dim = 2*hidden)
//! - norm + lm_head: → [draft_vocab_size=32000]
//! - d2t:          draft_id → target_id offset mapping

use std::collections::HashMap;
use std::path::Path;
use xserv_kernels::*;
use xserv_tensor::{DType, Device, Tensor};

/// Target layers to hook for EAGLE3 auxiliary hidden states, for Qwen3-8B
/// (36 layers). Value comes from AngelSlim/vLLM speculators training config
/// `dflash_qwen3_8b_sharegpt_online_5k.sh` which specifies target_layer_ids
/// = "2 18 33". Must match training-time selection or EAGLE outputs are wrong.
pub const EAGLE_HOOK_LAYERS: [usize; 3] = [2, 18, 33];
const DRAFT_VOCAB_SIZE: usize = 32000;

fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
    assert_eq!(a.ndim(), 2);
    assert_eq!(b.ndim(), 2);
    matmul(a, b, GemmBackend::CuBlas)
}

pub struct Eagle3Head {
    fc_wt: Tensor,                    // [hidden, 3*hidden] transposed for matmul
    hidden_norm: Tensor,              // [hidden]
    input_layernorm: Tensor,          // [hidden]
    q_proj_wt: Tensor,                // [num_heads*head_dim, 2*hidden]
    k_proj_wt: Tensor,                // [num_kv_heads*head_dim, 2*hidden]
    v_proj_wt: Tensor,                // [num_kv_heads*head_dim, 2*hidden]
    o_proj_wt: Tensor,                // [hidden, num_heads*head_dim]
    gate_proj_wt: Tensor,             // [intermediate, hidden]
    up_proj_wt: Tensor,               // [intermediate, hidden]
    down_proj_wt: Tensor,             // [hidden, intermediate]
    post_attention_layernorm: Tensor, // [hidden]
    norm: Tensor,                     // [hidden] final
    lm_head_wt: Tensor,               // [draft_vocab, hidden]
    d2t: Vec<i64>,                    // [draft_vocab] offset mapping
    /// t2d[target_id] = true iff target_id has a corresponding draft-vocab id
    /// (i.e. can potentially be produced by EAGLE). Used to measure the
    /// coverage cap on acceptance.
    t2d: Vec<bool>,
    hidden_size: usize,
    num_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
    max_seq_len: usize,
    rope_cache: RopeCache,
    // Stateful 1-layer KV cache: [1, num_kv_heads, max_seq_len, head_dim] BF16.
    // We slice `..current_len` for attention. The head is tiny (~64 KB per
    // 1000 tokens) so pre-allocating max_seq_len wastes negligible memory.
    k_cache: Tensor,
    v_cache: Tensor,
    current_len: usize,
}

impl Eagle3Head {
    pub fn load(dir: &Path, device: u32) -> Self {
        let (weights, d2t, t2d) = load_eagle3_weights(dir, device);
        let hidden_size = 4096;
        let num_heads = 32;
        let num_kv_heads = 8;
        let head_dim = 128;
        let intermediate_size = 12288;
        let max_seq_len = 2048;
        let rope_theta = 1_000_000.0f32;

        let get = |name: &str| -> Tensor {
            weights
                .get(name)
                .unwrap_or_else(|| panic!("missing eagle3 weight: {name}"))
                .clone()
        };

        let fc_wt = get("fc.weight").transpose(0, 1).contiguous();
        let q_proj_wt = get("midlayer.self_attn.q_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let k_proj_wt = get("midlayer.self_attn.k_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let v_proj_wt = get("midlayer.self_attn.v_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let o_proj_wt = get("midlayer.self_attn.o_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let gate_proj_wt = get("midlayer.mlp.gate_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let up_proj_wt = get("midlayer.mlp.up_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let down_proj_wt = get("midlayer.mlp.down_proj.weight")
            .transpose(0, 1)
            .contiguous();
        let hidden_norm = get("midlayer.hidden_norm.weight");
        let input_layernorm = get("midlayer.input_layernorm.weight");
        let post_attention_layernorm = get("midlayer.post_attention_layernorm.weight");
        let norm = get("norm.weight");
        let lm_head_wt = get("lm_head.weight").transpose(0, 1).contiguous();

        assert_eq!(d2t.len(), DRAFT_VOCAB_SIZE);

        let rope_cache = RopeCache::new(max_seq_len, head_dim, rope_theta);

        let k_cache = Tensor::zeros(
            &[1, num_kv_heads, max_seq_len, head_dim],
            DType::BF16,
            Device::Cuda(device),
        );
        let v_cache = Tensor::zeros(
            &[1, num_kv_heads, max_seq_len, head_dim],
            DType::BF16,
            Device::Cuda(device),
        );

        Self {
            fc_wt,
            hidden_norm,
            input_layernorm,
            q_proj_wt,
            k_proj_wt,
            v_proj_wt,
            o_proj_wt,
            gate_proj_wt,
            up_proj_wt,
            down_proj_wt,
            post_attention_layernorm,
            norm,
            lm_head_wt,
            d2t,
            t2d,
            hidden_size,
            num_heads,
            num_kv_heads,
            head_dim,
            max_seq_len,
            rope_cache,
            k_cache,
            v_cache,
            current_len: 0,
        }
    }

    /// Reset the internal KV cache for a fresh sequence.
    pub fn reset(&mut self) {
        self.current_len = 0;
    }

    /// Truncate the internal KV cache to `new_len` entries. Used to discard
    /// K/V of rejected drafts after a speculative round.
    pub fn truncate_to(&mut self, new_len: usize) {
        assert!(new_len <= self.current_len);
        self.current_len = new_len;
    }

    /// Current number of committed K/V entries in the internal EAGLE cache.
    pub fn current_len(&self) -> usize {
        self.current_len
    }

    /// One draft step: produce a token in target vocabulary space.
    ///
    /// - `target_hidden`: 3 tensors [1, hidden_size] from target hook layers
    /// - `embed_table`: the target model's embed_tokens (shared, not copied)
    /// - `prev_token`: the previous committed token
    /// - `position`: the decode position for RoPE
    ///
    /// Returns (draft_token_in_target_vocab, draft_logits_tensor).
    pub fn step(
        &mut self,
        target_hidden: &[Tensor; 3],
        embed_table: &Tensor,
        prev_token: u32,
        position: usize,
    ) -> (u32, Tensor) {
        let (id, logits, _) = self.step_with_aux(target_hidden, embed_table, prev_token, position);
        (id, logits)
    }

    /// Like `step`, but also returns the final hidden state (aux) usable as
    /// the fused_h for a subsequent recursive draft step via `step_recursive`.
    pub fn step_with_aux(
        &mut self,
        target_hidden: &[Tensor; 3],
        embed_table: &Tensor,
        prev_token: u32,
        position: usize,
    ) -> (u32, Tensor, Tensor) {
        // Fuse 3 target hidden states into fused_h via fc.
        let h_cat = concat_hidden(target_hidden);
        let fused_h = matmul_2d(&h_cat, &self.fc_wt);
        self.forward_from_fused(fused_h, embed_table, prev_token, position)
    }

    /// Recursive draft step: reuses the previous EAGLE step's aux as fused_h,
    /// bypassing the fc+3-hidden fusion. Used for γ≥2 chained drafts.
    pub fn step_recursive(
        &mut self,
        fused_h: Tensor,
        embed_table: &Tensor,
        prev_token: u32,
        position: usize,
    ) -> (u32, Tensor, Tensor) {
        self.forward_from_fused(fused_h, embed_table, prev_token, position)
    }

    fn forward_from_fused(
        &mut self,
        fused_h: Tensor,
        embed_table: &Tensor,
        prev_token: u32,
        position: usize,
    ) -> (u32, Tensor, Tensor) {
        let eps = 1e-6f32;
        assert!(
            self.current_len < self.max_seq_len,
            "EAGLE KV cache overflow: {} >= {}",
            self.current_len,
            self.max_seq_len
        );

        let emb = embedding(embed_table, &[prev_token]);
        let residual = fused_h.clone();
        let emb_normed = rmsnorm(&emb, &self.input_layernorm, eps);
        let h_normed = rmsnorm(&fused_h, &self.hidden_norm, eps);
        let attn_in = concat_last_dim(&emb_normed, &h_normed);

        let q = matmul_2d(&attn_in, &self.q_proj_wt);
        let k = matmul_2d(&attn_in, &self.k_proj_wt);
        let v = matmul_2d(&attn_in, &self.v_proj_wt);

        let q_3d = q.reshape(&[1, self.num_heads, self.head_dim]);
        let k_3d = k.reshape(&[1, self.num_kv_heads, self.head_dim]);
        let positions = [position as u32];
        rope_inplace(&q_3d, &self.rope_cache, &positions);
        rope_inplace(&k_3d, &self.rope_cache, &positions);

        let v_3d = v.reshape(&[1, self.num_kv_heads, self.head_dim]);
        self.append_to_kv_cache(&k_3d, &v_3d);
        self.current_len += 1;
        let kv_len = self.current_len;
        let k_view = self.k_cache.narrow(2, 0, kv_len).contiguous();
        let v_view = self.v_cache.narrow(2, 0, kv_len).contiguous();

        let q_4d = q_3d.reshape(&[1, self.num_heads, 1, self.head_dim]);
        let attn_out = decode_attention(&q_4d, &k_view, &v_view);

        let attn_merged = attn_out.reshape(&[1, self.num_heads * self.head_dim]);
        let attn_proj = matmul_2d(&attn_merged, &self.o_proj_wt);

        let (mlp_in, residual) =
            add_rmsnorm(&attn_proj, &residual, &self.post_attention_layernorm, eps);

        let gate = matmul_2d(&mlp_in, &self.gate_proj_wt);
        let up = matmul_2d(&mlp_in, &self.up_proj_wt);
        let hidden = silu_mul(&gate, &up);
        let down = matmul_2d(&hidden, &self.down_proj_wt);

        let (x, prenorm) = add_rmsnorm(&down, &residual, &self.norm, eps);
        let logits = matmul_2d(&x, &self.lm_head_wt);

        let draft_id = argmax_bf16_single(&logits);
        let target_id = (draft_id as i64 + self.d2t[draft_id as usize]) as u32;
        // aux for recursive drafting = PRE-norm hidden (default norm_output=False
        // in vllm/llama_eagle3.py). Feeding the pre-norm state matches training.
        (target_id, logits, prenorm)
    }

    /// Write new K/V rows (shape [1, num_kv_heads, head_dim]) at position
    /// `current_len` inside the [1, num_kv_heads, max_seq_len, head_dim] cache.
    fn append_to_kv_cache(&mut self, new_k: &Tensor, new_v: &Tensor) {
        let head_bytes = self.head_dim * self.k_cache.dtype().size_bytes();
        for h in 0..self.num_kv_heads {
            for (cache, src) in [(&self.k_cache, new_k), (&self.v_cache, new_v)] {
                let dst = unsafe {
                    (cache.data_ptr() as *mut u8)
                        .add(((h * self.max_seq_len) + self.current_len) * head_bytes)
                };
                let s = unsafe { (src.data_ptr() as *const u8).add(h * head_bytes) };
                d2d(dst, s, head_bytes);
            }
        }
    }

    /// Map a draft-vocab token id to the full target-vocab id via d2t.
    pub fn map_draft_to_target(&self, draft_id: u32) -> u32 {
        (draft_id as i64 + self.d2t[draft_id as usize]) as u32
    }

    /// Returns true iff `target_id` is representable in the draft vocabulary
    /// (i.e., EAGLE could in principle produce it).
    pub fn target_id_in_draft_vocab(&self, target_id: u32) -> bool {
        self.t2d.get(target_id as usize).copied().unwrap_or(false)
    }
}

fn d2d(dst: *mut u8, src: *const u8, bytes: usize) {
    unsafe {
        xserv_cuda::ffi::cudaMemcpy(dst, src, bytes, xserv_cuda::ffi::CUDA_MEMCPY_D2D);
    }
}

fn concat_hidden(hidden: &[Tensor; 3]) -> Tensor {
    let h = hidden[0].shape()[1];
    let dtype = hidden[0].dtype();
    let device = hidden[0].device();
    let elem_bytes = dtype.size_bytes();
    let out = Tensor::empty(&[1, 3 * h], dtype, device);
    for (i, t) in hidden.iter().enumerate() {
        assert!(t.is_contiguous());
        let dst = unsafe { (out.data_ptr() as *mut u8).add(i * h * elem_bytes) };
        d2d(dst, t.data_ptr() as *const u8, h * elem_bytes);
    }
    out
}

fn concat_last_dim(a: &Tensor, b: &Tensor) -> Tensor {
    let da = a.shape()[1];
    let db = b.shape()[1];
    let dtype = a.dtype();
    let device = a.device();
    let elem_bytes = dtype.size_bytes();
    let out = Tensor::empty(&[1, da + db], dtype, device);
    d2d(
        out.data_ptr() as *mut u8,
        a.data_ptr() as *const u8,
        da * elem_bytes,
    );
    let dst = unsafe { (out.data_ptr() as *mut u8).add(da * elem_bytes) };
    d2d(dst, b.data_ptr() as *const u8, db * elem_bytes);
    out
}

fn repeat_kv_for_single_token(kv: &Tensor, repeats: usize) -> Tensor {
    if repeats == 1 {
        return kv.clone();
    }
    let nkv = kv.shape()[1];
    let d = kv.shape()[2];
    let dtype = kv.dtype();
    let device = kv.device();
    let head_bytes = d * dtype.size_bytes();
    let out = Tensor::empty(&[1, nkv * repeats, d], dtype, device);
    for h in 0..nkv {
        let src = unsafe { (kv.data_ptr() as *const u8).add(h * head_bytes) };
        for r in 0..repeats {
            let dst = unsafe { (out.data_ptr() as *mut u8).add((h * repeats + r) * head_bytes) };
            d2d(dst, src, head_bytes);
        }
    }
    out
}

/// Load EAGLE3 weights from safetensors, handling int64 d2t + bool t2d specially.
fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec<i64>, Vec<bool>) {
    let st_path = dir.join("model.safetensors");
    assert!(
        st_path.exists(),
        "Eagle3 model.safetensors not found in {}. Convert with:\n\
         python3 -c \"import torch; from safetensors.torch import save_file; \
         sd=torch.load('pytorch_model.bin', map_location='cpu', weights_only=False); \
         save_file(sd, 'model.safetensors')\"",
        dir.display()
    );

    let data = std::fs::read(&st_path)
        .unwrap_or_else(|e| panic!("failed to read {}: {e}", st_path.display()));
    let st = safetensors::SafeTensors::deserialize(&data)
        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", st_path.display()));

    let mut tensors = HashMap::new();
    let mut d2t_vec: Vec<i64> = Vec::new();
    let mut t2d_vec: Vec<bool> = Vec::new();

    for (name, view) in st.tensors() {
        if name == "t2d" {
            let raw = view.data();
            assert_eq!(view.dtype(), safetensors::Dtype::BOOL);
            t2d_vec = raw.iter().map(|&b| b != 0).collect();
            continue;
        }
        if name == "d2t" {
            let raw = view.data();
            assert_eq!(view.dtype(), safetensors::Dtype::I64);
            let n = raw.len() / 8;
            d2t_vec = (0..n)
                .map(|i| i64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap()))
                .collect();
            continue;
        }
        let dtype = match view.dtype() {
            safetensors::Dtype::BF16 => DType::BF16,
            safetensors::Dtype::F32 => DType::F32,
            safetensors::Dtype::F16 => DType::F16,
            other => {
                eprintln!("eagle3: skipping {name} with unsupported dtype {other:?}");
                continue;
            }
        };
        let shape: Vec<usize> = view.shape().to_vec();
        let raw = view.data();
        let t = crate::loader::make_tensor(raw, &shape, dtype);
        let t = t.to_device(Device::Cuda(device));
        tensors.insert(name.to_string(), t);
    }

    assert!(
        !d2t_vec.is_empty(),
        "d2t tensor not found in eagle3 weights"
    );
    assert!(
        !t2d_vec.is_empty(),
        "t2d tensor not found in eagle3 weights"
    );
    (tensors, d2t_vec, t2d_vec)
}