//! GPU forward of the Gemma-4 MTP assistant (the speculative draft head), batched over B <= 16 sequences. Its q-only decoder layers read the TARGET's paged
//! K/V of the last sliding (28) / full (29) layer, read-only. Semantics: transformers `Gemma4AssistantForCausalLM` + vLLM `gemma4_mtp.py` / `Gemma4Proposer`
//! (query position constant = position of the last committed row).
//! Weights are quantised to FP8 per output row at load (from the public bf16 checkpoint); the lm_head is the draft's own tied table PRUNED to the most
//! frequent thinking-trace tokens (`DraftVocab`). The draft only PROPOSES tokens: the verify forward of the target decides every emitted token, so the
//! draft's precision cannot change the output (it only changes how many tokens are accepted per step).
use crate::gpu_model::GpuModel;
use anyhow::{ensure, Result};
use dllm_core::config::{AssistantConfig, LayerKind};
use dllm_core::loader::WeightStore;
use dllm_core::quant::quantize_fp8_rows;
use dllm_kernels::attn as fattn;
use dllm_kernels::linear::{gemv_fused, Prologue, Seg};
use dllm_kernels::spec::KMAX;
use dllm_kernels::{DBuf, Gpu};
use half::bf16;

pub const MAX_DRAFT_ROWS: usize = 16;

/// draft-vocabulary table: kept global token ids, ascending (so argmax ties resolve to the lowest global id, like the full-vocab argmax)
#[derive(Debug, Clone)]
pub struct DraftVocab {
    pub ids: Vec<i32>,
}
impl DraftVocab {
    pub fn from_json_file(p: impl AsRef<std::path::Path>) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(p)?)?;
        let mut ids: Vec<i32> = v["ids"].as_array().ok_or_else(|| anyhow::anyhow!("draft vocab json: no ids"))?.iter().map(|x| x.as_i64().unwrap() as i32).collect();
        ids.sort_unstable();
        ids.dedup();
        Ok(DraftVocab { ids })
    }
    pub fn len(&self) -> usize {
        self.ids.len()
    }
    pub fn is_empty(&self) -> bool {
        self.ids.is_empty()
    }
}

/// FP8 linear quantised at load: e4m3 [n, k] + f32 per-row scale
struct Lin {
    n: usize,
    k: usize,
    w: DBuf,
    s: DBuf,
}
impl Lin {
    fn seg(&self) -> Seg {
        Seg { w: self.w.ptr, s: self.s.ptr, n: self.n }
    }
}

fn bf16_bytes(t: &[f32]) -> Vec<u8> {
    t.iter().flat_map(|x| bf16::from_f32(*x).to_le_bytes()).collect()
}
fn f32_bytes(v: &[f32]) -> Vec<u8> {
    v.iter().flat_map(|x| x.to_le_bytes()).collect()
}

fn fp8_lin(gpu: &Gpu, data: &[f32], n: usize, k: usize) -> Result<Lin> {
    let q = quantize_fp8_rows(data, n, k);
    Ok(Lin { n, k, w: DBuf::from_bytes(&gpu.be, &q.w)?, s: DBuf::from_bytes(&gpu.be, &f32_bytes(&q.scales))? })
}

struct DLayer {
    hd: usize,
    nkv: usize,
    window: usize,
    /// target layer whose K/V this layer reads
    tgt_layer: usize,
    rope: usize,
    in_ln: DBuf,
    q: Lin,
    o: Lin,
    q_norm: DBuf,
    post_attn_ln: DBuf,
    pre_ff_ln: DBuf,
    gate: Lin,
    up: Lin,
    down: Lin,
    post_ff_ln: DBuf,
    layer_scalar: DBuf,
}

struct Head {
    /// rows of the pruned table
    n: usize,
    w: DBuf,
    s: DBuf,
    /// row -> global token id
    map: DBuf,
}

struct Ws {
    x: DBuf,
    h: DBuf,
    qraw: DBuf,
    q: DBuf,
    ao: DBuf,
    o: DBuf,
    mid: DBuf,
    f: DBuf,
    act: DBuf,
    m: DBuf,
    hn: DBuf,
    backbone: DBuf,
    idx: DBuf,
    val: DBuf,
    dummy: DBuf,
    /// split-KV decode partials / merge counters, lm_head scan scratch
    part: DBuf,
    cnt: DBuf,
    head: dllm_kernels::head::HeadScratch,
}

pub struct Draft {
    pub gpu: Gpu,
    pub cfg: AssistantConfig,
    layers: Vec<DLayer>,
    pre_proj: Lin,
    post_proj: Lin,
    norm: DBuf,
    head: Head,
    rope: Vec<DBuf>,
    ws: Ws,
    pub hb: usize,
    eps: f32,
}

/// what a draft step needs from the serving state (all device addresses)
#[derive(Clone, Copy)]
pub struct DraftCtx {
    /// block tables of the sequences [B][maxp] (the TARGET's page tables) and per-row offsets
    pub bt: u64,
    pub bt_off: u64,
    /// per-sequence draft query position == position of the last committed row (constant over a chain); also the last visible kv position
    pub dpos: u64,
    pub splits: usize,
}

impl Draft {
    pub fn load(gpu: Gpu, st: &WeightStore, cfg: &AssistantConfig, target_cfg: &dllm_core::config::TextConfig, vocab: &DraftVocab, max_pos: usize) -> Result<Self> {
        let be = gpu.be.clone();
        let c = &cfg.text;
        let (h, hb) = (c.hidden_size, cfg.backbone_hidden_size);
        let v = |n: &str| -> Result<DBuf> { DBuf::from_bytes(&be, &bf16_bytes(&st.load_f32(n)?.data)) };
        let lin = |n: &str| -> Result<Lin> {
            let t = st.load_f32(n)?;
            fp8_lin(&gpu, &t.data, t.rows(), t.cols())
        };
        let mut layers = vec![];
        let mut rope: Vec<DBuf> = vec![];
        let mut rope_keys: Vec<(usize, u64, u64, bool)> = vec![];
        for (i, sp) in c.layers.iter().enumerate().take(c.num_layers) {
            let p = format!("model.layers.{i}.");
            let key = (sp.head_dim, sp.rope_theta.to_bits(), sp.partial_rotary.to_bits(), sp.proportional_rope);
            let ri = match rope_keys.iter().position(|k| *k == key) {
                Some(r) => r,
                None => {
                    rope.push(DBuf::from_slice(&be, &crate::rope::cos_sin_table(sp, max_pos))?);
                    rope_keys.push(key);
                    rope.len() - 1
                }
            };
            let tgt_layer = target_cfg.last_layer_of(sp.kind).ok_or_else(|| anyhow::anyhow!("target has no {:?} layer", sp.kind))?;
            let tsp = &target_cfg.layers[tgt_layer];
            ensure!(tsp.head_dim == sp.head_dim && tsp.n_kv_heads == sp.n_kv_heads, "draft layer {i} geometry != target layer {tgt_layer}");
            layers.push(DLayer {
                hd: sp.head_dim,
                nkv: sp.n_kv_heads,
                window: if sp.kind == LayerKind::Sliding { sp.window.unwrap_or(0) } else { 0 },
                tgt_layer,
                rope: ri,
                in_ln: v(&format!("{p}input_layernorm.weight"))?,
                q: lin(&format!("{p}self_attn.q_proj.weight"))?,
                o: lin(&format!("{p}self_attn.o_proj.weight"))?,
                q_norm: v(&format!("{p}self_attn.q_norm.weight"))?,
                post_attn_ln: v(&format!("{p}post_attention_layernorm.weight"))?,
                pre_ff_ln: v(&format!("{p}pre_feedforward_layernorm.weight"))?,
                gate: lin(&format!("{p}mlp.gate_proj.weight"))?,
                up: lin(&format!("{p}mlp.up_proj.weight"))?,
                down: lin(&format!("{p}mlp.down_proj.weight"))?,
                post_ff_ln: v(&format!("{p}post_feedforward_layernorm.weight"))?,
                layer_scalar: v(&format!("{p}layer_scalar"))?,
            });
        }
        let pre_proj = lin("pre_projection.weight")?;
        let post_proj = lin("post_projection.weight")?;
        ensure!(pre_proj.k == 2 * hb && pre_proj.n == h, "pre_projection shape");
        // lm_head = the draft's tied embedding table gathered to the pruned vocabulary
        let emb = st.load_f32("model.embed_tokens.weight")?;
        let mut rows = Vec::with_capacity(vocab.len() * h);
        for &id in &vocab.ids {
            ensure!((id as usize) < emb.rows(), "draft vocab id {id} out of range");
            rows.extend_from_slice(emb.row(id as usize));
        }
        let n = rows.len() / h;
        let q = quantize_fp8_rows(&rows, n, h);
        let head = Head { n, w: DBuf::from_bytes(&be, &q.w)?, s: DBuf::from_bytes(&be, &f32_bytes(&q.scales))?, map: DBuf::from_slice(&be, &vocab.ids)? };
        let r = MAX_DRAFT_ROWS;
        let z = |bytes: usize| DBuf::zeros(&be, bytes);
        let max_q = c.layers.iter().map(|s| c.num_heads * s.head_dim).max().unwrap();
        let dff = c.intermediate_size;
        let max_part = [256usize, 512].iter().map(|&hd| fattn::decode_part_bytes(MAX_DRAFT_ROWS, c.num_heads, 32, hd)).max().unwrap();
        let ws = Ws {
            x: z(r * 2 * hb * 2)?,
            h: z(r * h * 2)?,
            qraw: z(r * max_q * 2)?,
            q: z(r * max_q * 2)?,
            ao: z(r * max_q * 2)?,
            o: z(r * h * 2)?,
            mid: z(r * h * 2)?,
            f: z(r * h * 2)?,
            act: z(r * dff * 2)?,
            m: z(r * h * 2)?,
            hn: z(r * h * 2)?,
            backbone: z(r * hb * 2)?,
            idx: z(r * 4)?,
            val: z(r * 4)?,
            dummy: z(256)?,
            part: z(max_part)?,
            cnt: z(4096 * 4)?,
            head: dllm_kernels::head::HeadScratch::new(&be)?,
        };
        Ok(Draft { gpu, cfg: cfg.clone(), layers, pre_proj, post_proj, norm: v("model.norm.weight")?, head, rope, ws, hb, eps: c.rms_eps })
    }

    /// One draft step for `b` sequences: input token of sequence s = `tok[s*tok_stride]` (embedded with the TARGET table), hidden input `hs` [b, Hb] bf16.
    /// Leaves the pruned-row argmax in `ws.idx` and, if `want_backbone`, the projected state in `ws.backbone` (the next step's `hs`).
    #[allow(clippy::too_many_arguments)]
    fn step(&self, target: &GpuModel, c: &DraftCtx, tok: u64, tok_stride: usize, hs: u64, b: usize, want_backbone: bool) -> Result<()> {
        ensure!(b >= 1 && b <= MAX_DRAFT_ROWS, "draft rows {b}");
        let gpu = &self.gpu;
        let be = &gpu.be;
        let (hb, h, eps) = (self.hb, self.cfg.text.hidden_size, self.eps);
        let ws = &self.ws;
        let (tbl, rs, scale) = target.embed_desc();
        gpu.spec_draft_input(ws.x.ptr, tbl, rs, tok, tok_stride, hs, b, hb, scale)?;
        gemv_fused(be, ws.h.ptr, h, ws.x.ptr, Prologue::Bf16, 0, eps, b, 2 * hb, &[self.pre_proj.seg()], false, None)?;
        let nh = self.cfg.text.num_heads;
        let dff = self.cfg.text.intermediate_size;
        for l in &self.layers {
            gemv_fused(be, ws.qraw.ptr, l.q.n, ws.h.ptr, Prologue::Rms, l.in_ln.ptr, eps, b, h, &[l.q.seg()], false, None)?;
            // q-only (n_kv = 0): norm(q; q_norm) + RoPE at the constant draft position
            gpu.qkv_prep(ws.q.ptr, ws.dummy.ptr, ws.dummy.ptr, ws.qraw.ptr, ws.dummy.ptr, ws.dummy.ptr, l.q_norm.ptr, ws.dummy.ptr, self.rope[l.rope].ptr, c.dpos, ws.dummy.ptr, b, nh, 0, l.hd, eps, 0)?;
            let (kp, vp) = target.kv_pool_ptrs(l.tgt_layer);
            let sp = fattn::choose_splits(c.splits.max(1) * 1024, b, l.nkv, 1);
            let (part, cnt) = if sp > 1 { (ws.part.ptr, ws.cnt.ptr) } else { (0, 0) };
            fattn::attn_decode(gpu, ws.ao.ptr, part, ws.q.ptr, kp, vp, c.bt, c.bt_off, c.dpos, c.dpos, b, nh, l.nkv, l.hd, l.window, sp, cnt, target.kv_pmap(l.tgt_layer))?;
            gemv_fused(be, ws.o.ptr, h, ws.ao.ptr, Prologue::Bf16, 0, eps, b, l.o.k, &[l.o.seg()], false, None)?;
            gpu.sandwich(ws.mid.ptr, ws.f.ptr, ws.h.ptr, ws.o.ptr, l.post_attn_ln.ptr, l.pre_ff_ln.ptr, b, h, eps)?;
            gemv_fused(be, ws.act.ptr, dff, ws.f.ptr, Prologue::Bf16, 0, eps, b, h, &[l.gate.seg(), l.up.seg()], true, None)?;
            gemv_fused(be, ws.m.ptr, h, ws.act.ptr, Prologue::Bf16, 0, eps, b, dff, &[l.down.seg()], false, None)?;
            gpu.layer_out(ws.h.ptr, ws.mid.ptr, ws.m.ptr, l.post_ff_ln.ptr, l.layer_scalar.ptr, b, h, eps)?;
        }
        gpu.rmsnorm(ws.hn.ptr, ws.h.ptr, self.norm.ptr, b, h, eps)?;
        let mut c0 = 0;
        while c0 < b {
            let m = (b - c0).min(8);
            dllm_kernels::head::lm_head_argmax(be, &ws.head, ws.idx.ptr + (c0 * 4) as u64, ws.val.ptr + (c0 * 4) as u64, ws.hn.ptr + (c0 * h * 2) as u64, self.head.w.ptr, self.head.s.ptr, m, self.head.n, h, -1, -1, 0.0, false)?;
            c0 += m;
        }
        if want_backbone {
            gemv_fused(be, ws.backbone.ptr, hb, ws.hn.ptr, Prologue::Bf16, 0, eps, b, h, &[self.post_proj.seg()], false, None)?;
        }
        Ok(())
    }

    /// greedy chain of `k` draft tokens for `b` sequences into `drafts[s*KMAX + j]`;  step 0 input = (x0, h0), step j input = (draft j, backbone j-1)
    #[allow(clippy::too_many_arguments)]
    pub fn chain(&self, target: &GpuModel, c: &DraftCtx, x0: u64, h0: u64, drafts: u64, b: usize, k: usize) -> Result<()> {
        ensure!(k >= 1 && k < KMAX);
        for j in 0..k {
            let (tok, stride, hs) = if j == 0 { (x0, 1, h0) } else { (drafts + ((j - 1) * 4) as u64, KMAX, self.ws.backbone.ptr) };
            self.step(target, c, tok, stride, hs, b, j + 1 < k)?;
            self.gpu.spec_draft_pick(drafts, self.ws.idx.ptr, self.head.map.ptr, j, b)?;
        }
        Ok(())
    }
}