use super::checkpoint::{read_buffer, read_u32, read_u64, write_buffer, write_u32, write_u64};
use super::*;
use crate::settings::EngineSsdSettings;

const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4GLM01";
const CHECKPOINT_VERSION: u32 = 1;
const CACHE_F16: bool = true;
const DECODE_FLUSH_LAYERS: usize = 4;

#[derive(Clone, Copy)]
struct SparseWeights {
    router: Weight,
    bias: Weight,
    gate: Weight,
    up: Weight,
    down: Weight,
    shared_gate: Weight,
    shared_up: Weight,
    shared_down: Weight,
}

#[derive(Clone, Copy)]
struct DenseWeights {
    gate: Weight,
    up: Weight,
    down: Weight,
}

struct GlmLayer {
    attn_norm: Weight,
    q_a: Weight,
    q_a_norm: Weight,
    q_b: Weight,
    kv_a: Weight,
    kv_norm: Weight,
    k_b: Weight,
    v_b: Weight,
    output: Weight,
    indexer_k: Weight,
    indexer_q: Weight,
    indexer_k_norm: Weight,
    indexer_k_bias: Weight,
    indexer_proj: Weight,
    ffn_norm: Weight,
    dense: Option<DenseWeights>,
    sparse: Option<SparseWeights>,
}

struct GlmWeights {
    embedding: Weight,
    output_norm: Weight,
    output: Weight,
    layers: Vec<GlmLayer>,
}

impl GlmWeights {
    fn bind(model: &Model) -> Result<Self, String> {
        let main = &model.main;
        let normal_layers = model.shape.layers - model.shape.nextn;
        let layers = (0..normal_layers)
            .map(|index| {
                let required = |suffix: &str| Weight::bind(main, &format!("blk.{index}.{suffix}"));
                let dense = (index < model.shape.leading_dense)
                    .then(|| {
                        Ok::<_, String>(DenseWeights {
                            gate: required("ffn_gate.weight")?,
                            up: required("ffn_up.weight")?,
                            down: required("ffn_down.weight")?,
                        })
                    })
                    .transpose()?;
                let sparse = (index >= model.shape.leading_dense)
                    .then(|| {
                        Ok::<_, String>(SparseWeights {
                            router: required("ffn_gate_inp.weight")?,
                            bias: required("exp_probs_b.bias")?,
                            gate: required("ffn_gate_exps.weight")?,
                            up: required("ffn_up_exps.weight")?,
                            down: required("ffn_down_exps.weight")?,
                            shared_gate: required("ffn_gate_shexp.weight")?,
                            shared_up: required("ffn_up_shexp.weight")?,
                            shared_down: required("ffn_down_shexp.weight")?,
                        })
                    })
                    .transpose()?;
                Ok(GlmLayer {
                    attn_norm: required("attn_norm.weight")?,
                    q_a: required("attn_q_a.weight")?,
                    q_a_norm: required("attn_q_a_norm.weight")?,
                    q_b: required("attn_q_b.weight")?,
                    kv_a: required("attn_kv_a_mqa.weight")?,
                    kv_norm: required("attn_kv_a_norm.weight")?,
                    k_b: required("attn_k_b.weight")?,
                    v_b: required("attn_v_b.weight")?,
                    output: required("attn_output.weight")?,
                    indexer_k: required("indexer.attn_k.weight")?,
                    indexer_q: required("indexer.attn_q_b.weight")?,
                    indexer_k_norm: required("indexer.k_norm.weight")?,
                    indexer_k_bias: required("indexer.k_norm.bias")?,
                    indexer_proj: required("indexer.proj.weight")?,
                    ffn_norm: required("ffn_norm.weight")?,
                    dense,
                    sparse,
                })
            })
            .collect::<Result<_, String>>()?;
        Ok(Self {
            embedding: Weight::bind(main, "token_embd.weight")?,
            output_norm: Weight::bind(main, "output_norm.weight")?,
            output: Weight::bind(main, "output.weight")?,
            layers,
        })
    }
}

struct LayerCache {
    kv: Buffer,
    rope: Buffer,
    indexer: Option<Buffer>,
}

impl LayerCache {
    fn allocate(shape: super::super::Shape, layer: usize, context: u32) -> Result<Self, String> {
        Ok(Self {
            kv: Buffer::bytes(u64::from(context) * shape.kv_lora * 2)?,
            rope: Buffer::bytes(u64::from(context) * shape.rot * 2)?,
            indexer: full_indexer_layer(shape, layer)
                .then(|| Buffer::bytes(u64::from(context) * shape.indexer_head_dim * 2))
                .transpose()?,
        })
    }
}

struct GlmScratch {
    current: Buffer,
    next: Buffer,
    attn_norm: Buffer,
    q_rank: Buffer,
    q_rank_norm: Buffer,
    q: Buffer,
    kv_raw: Buffer,
    indexer_k: Buffer,
    indexer_q: Buffer,
    indexer_weights: Buffer,
    indexer_scores: Buffer,
    indexer_selected: Buffer,
    qk_low: Buffer,
    heads: Buffer,
    attn_out: Buffer,
    after_attn: Buffer,
    ffn_norm: Buffer,
    ffn_gate: Buffer,
    ffn_up: Buffer,
    ffn_mid: Buffer,
    ffn_out: Buffer,
    ffn_sum: Buffer,
    router_logits: Buffer,
    router_probs: Buffer,
    router_selected: Buffer,
    router_weights: Buffer,
    output_norm: Buffer,
    logits: Buffer,
}

impl GlmScratch {
    fn allocate(model: &Model, context: u32) -> Result<Self, String> {
        let shape = model.shape;
        let q = shape.heads * shape.key_mla;
        let heads = shape.heads * shape.value_mla;
        Ok(Self {
            current: Buffer::floats(shape.embd)?,
            next: Buffer::floats(shape.embd)?,
            attn_norm: Buffer::floats(shape.embd)?,
            q_rank: Buffer::floats(shape.lora_q)?,
            q_rank_norm: Buffer::floats(shape.lora_q)?,
            q: Buffer::floats(q)?,
            kv_raw: Buffer::floats(shape.head_dim)?,
            indexer_k: Buffer::floats(shape.indexer_head_dim)?,
            indexer_q: Buffer::floats(shape.indexer_heads * shape.indexer_head_dim)?,
            indexer_weights: Buffer::floats(shape.indexer_heads)?,
            indexer_scores: Buffer::floats(u64::from(context))?,
            indexer_selected: Buffer::bytes(shape.indexer_top_k * 4)?,
            qk_low: Buffer::floats(shape.heads * shape.kv_lora)?,
            heads: Buffer::floats(heads)?,
            attn_out: Buffer::floats(shape.embd)?,
            after_attn: Buffer::floats(shape.embd)?,
            ffn_norm: Buffer::floats(shape.embd)?,
            ffn_gate: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?,
            ffn_up: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?,
            ffn_mid: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?,
            ffn_out: Buffer::floats(shape.embd)?,
            ffn_sum: Buffer::floats(shape.embd)?,
            router_logits: Buffer::floats(shape.experts)?,
            router_probs: Buffer::floats(shape.experts)?,
            router_selected: Buffer::bytes(shape.experts_used * 4)?,
            router_weights: Buffer::floats(shape.experts_used)?,
            output_norm: Buffer::floats(shape.embd)?,
            logits: Buffer::floats(shape.vocab)?,
        })
    }
}

pub(in crate::engine) struct GlmExecutor {
    weights: GlmWeights,
    scratch: GlmScratch,
    caches: Vec<LayerCache>,
    logits: Vec<f32>,
    tokens: Vec<i32>,
    context: u32,
    quality: bool,
    ssd: EngineSsdSettings,
    checkpoint_tag: [u8; 32],
    model_modified: (u64, u32),
    model_identity: [u8; 32],
    _context: Context,
    model: Model,
}

impl GlmExecutor {
    pub(super) fn open(
        model: Model,
        context: u32,
        quality: bool,
        ssd: EngineSsdSettings,
    ) -> Result<Self, String> {
        if context == 0 || u64::from(context) > model.shape.original_context {
            return Err(format!(
                "GLM context must be between 1 and {} tokens",
                model.shape.original_context
            ));
        }
        let weights = GlmWeights::bind(&model)?;
        let admission = admission_bytes(&model, &weights, context, ssd)?;
        let context_handle = Context::open(&model, quality, ssd.enabled, admission)?;
        configure_streaming(&model, &weights, ssd)?;
        let scratch = GlmScratch::allocate(&model, context)?;
        let caches = (0..weights.layers.len())
            .map(|layer| LayerCache::allocate(model.shape, layer, context))
            .collect::<Result<_, _>>()?;
        let model_modified = fs::metadata(model.main.path())
            .and_then(|metadata| metadata.modified())
            .ok()
            .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok())
            .map(|duration| (duration.as_secs(), duration.subsec_nanos()))
            .unwrap_or_default();
        let model_identity = model.checkpoint_identity();
        Ok(Self {
            weights,
            scratch,
            caches,
            logits: vec![0.0; model.shape.vocab as usize],
            tokens: Vec::new(),
            context,
            quality,
            ssd,
            checkpoint_tag: [0; 32],
            model_modified,
            model_identity,
            _context: context_handle,
            model,
        })
    }

    pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
        let shape = self.model.shape;
        if token < 0 || token as u64 >= shape.vocab {
            return Err(format!("token {token} is outside the vocabulary"));
        }
        let pos = u32::try_from(self.tokens.len()).map_err(|_| "GLM position overflow")?;
        if pos >= self.context {
            return Err(format!(
                "the GLM Metal executor supports {} tokens",
                self.context
            ));
        }
        let map = self.model.main.map_ptr().cast();
        let size = self.model.main.len();
        let commands = Commands::begin()?;
        call(
            unsafe {
                ds4_gpu_embed_token_quant_tensor(
                    self.scratch.current.raw(),
                    map,
                    size,
                    self.weights.embedding.offset,
                    self.weights.embedding.kind,
                    shape.vocab as u32,
                    token as u32,
                    shape.embd as u32,
                )
            },
            "GLM token embedding",
        )?;
        let mut selected_count = 0;
        for (layer_index, ((layer, cache), ordinal)) in self
            .weights
            .layers
            .iter()
            .zip(&self.caches)
            .zip(0_u32..)
            .enumerate()
        {
            self.encode_layer(layer, cache, layer_index, ordinal, pos, &mut selected_count)?;
            std::mem::swap(&mut self.scratch.current, &mut self.scratch.next);
            if (layer_index + 1).is_multiple_of(DECODE_FLUSH_LAYERS)
                && layer_index + 1 < self.weights.layers.len()
            {
                call(
                    unsafe { ds4_gpu_flush_commands() },
                    "flushing the GLM decode graph",
                )?;
            }
        }
        norm(
            &self.scratch.output_norm,
            &self.scratch.current,
            self.weights.output_norm,
            shape.embd as u32,
            shape.rms_epsilon,
            map,
            size,
        )?;
        project(
            &self.scratch.logits,
            self.weights.output,
            shape.embd,
            shape.vocab,
            &self.scratch.output_norm,
            map,
            size,
            self.ssd.enabled,
        )?;
        commands.finish()?;
        self.scratch.logits.read_f32(&mut self.logits)?;
        self.tokens.push(token);
        Ok(())
    }

    fn encode_layer(
        &self,
        layer: &GlmLayer,
        cache: &LayerCache,
        _layer_index: usize,
        ordinal: u32,
        pos: u32,
        selected_count: &mut u32,
    ) -> Result<(), String> {
        let shape = self.model.shape;
        let map = self.model.main.map_ptr().cast();
        let size = self.model.main.len();
        let q_dim = shape.heads * shape.key_mla;
        let q_nope = shape.key_mla - shape.rot;
        norm(
            &self.scratch.attn_norm,
            &self.scratch.current,
            layer.attn_norm,
            shape.embd as u32,
            shape.rms_epsilon,
            map,
            size,
        )?;
        project(
            &self.scratch.q_rank,
            layer.q_a,
            shape.embd,
            shape.lora_q,
            &self.scratch.attn_norm,
            map,
            size,
            self.ssd.enabled,
        )?;
        project(
            &self.scratch.kv_raw,
            layer.kv_a,
            shape.embd,
            shape.head_dim,
            &self.scratch.attn_norm,
            map,
            size,
            self.ssd.enabled,
        )?;
        call(
            unsafe {
                ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
                    self.scratch.q_rank_norm.raw(),
                    self.scratch.q_rank.raw(),
                    map,
                    size,
                    layer.q_a_norm.offset,
                    shape.lora_q as u32,
                    cache.kv.raw(),
                    cache.rope.raw(),
                    self.scratch.kv_raw.raw(),
                    layer.kv_norm.offset,
                    pos,
                    1,
                    self.context,
                    shape.head_dim as u32,
                    shape.kv_lora as u32,
                    shape.rot as u32,
                    CACHE_F16,
                    shape.rms_epsilon,
                )
            },
            "storing GLM compact KV",
        )?;
        project(
            &self.scratch.q,
            layer.q_b,
            shape.lora_q,
            q_dim,
            &self.scratch.q_rank_norm,
            map,
            size,
            self.ssd.enabled,
        )?;
        call(
            unsafe {
                ds4_gpu_glm_rope_tail_tensor(
                    self.scratch.q.raw(),
                    1,
                    shape.heads as u32,
                    shape.key_mla as u32,
                    shape.rot as u32,
                    pos,
                    0,
                    shape.rope_base,
                    1.0,
                    0.0,
                    1.0,
                    0.0,
                    0.0,
                )
            },
            "applying GLM query RoPE",
        )?;

        if let Some(indexer_cache) = &cache.indexer {
            project(
                &self.scratch.indexer_k,
                layer.indexer_k,
                shape.embd,
                shape.indexer_head_dim,
                &self.scratch.current,
                map,
                size,
                self.ssd.enabled,
            )?;
            call(
                unsafe {
                    ds4_gpu_glm_store_indexer_k_tensor(
                        indexer_cache.raw(),
                        self.scratch.indexer_k.raw(),
                        map,
                        size,
                        layer.indexer_k_norm.offset,
                        layer.indexer_k_bias.offset,
                        pos,
                        1,
                        self.context,
                        shape.indexer_head_dim as u32,
                        shape.rot as u32,
                        0,
                        1.0e-6,
                        shape.rope_base,
                        1.0,
                        0.0,
                        1.0,
                        0.0,
                        0.0,
                        CACHE_F16,
                    )
                },
                "storing the GLM indexer key",
            )?;
            let visible = pos + 1;
            *selected_count = visible.min(shape.indexer_top_k as u32);
            if visible <= shape.indexer_top_k as u32 {
                call(
                    unsafe {
                        ds4_gpu_glm_fill_selected_range_tensor(
                            self.scratch.indexer_selected.raw(),
                            *selected_count,
                        )
                    },
                    "selecting the visible GLM context",
                )?;
            } else {
                project(
                    &self.scratch.indexer_q,
                    layer.indexer_q,
                    shape.lora_q,
                    shape.indexer_heads * shape.indexer_head_dim,
                    &self.scratch.q_rank_norm,
                    map,
                    size,
                    self.ssd.enabled,
                )?;
                call(
                    unsafe {
                        ds4_gpu_glm_indexer_rope_tail_tensor(
                            self.scratch.indexer_q.raw(),
                            1,
                            shape.indexer_heads as u32,
                            shape.indexer_head_dim as u32,
                            shape.rot as u32,
                            pos,
                            0,
                            shape.rope_base,
                            1.0,
                            0.0,
                            1.0,
                            0.0,
                            0.0,
                        )
                    },
                    "applying GLM indexer RoPE",
                )?;
                f32_project(
                    &self.scratch.indexer_weights,
                    layer.indexer_proj,
                    shape.embd,
                    shape.indexer_heads,
                    &self.scratch.current,
                    map,
                    size,
                )?;
                let scale = 1.0 / ((shape.indexer_heads * shape.indexer_head_dim) as f32).sqrt();
                call(
                    unsafe {
                        ds4_gpu_glm_indexer_score_one_tensor(
                            self.scratch.indexer_scores.raw(),
                            self.scratch.indexer_q.raw(),
                            self.scratch.indexer_weights.raw(),
                            indexer_cache.raw(),
                            visible,
                            shape.indexer_heads as u32,
                            shape.indexer_head_dim as u32,
                            scale,
                            CACHE_F16,
                        )
                    },
                    "scoring the GLM indexer",
                )?;
                call(
                    unsafe {
                        ds4_gpu_indexer_topk_tensor(
                            self.scratch.indexer_selected.raw(),
                            self.scratch.indexer_scores.raw(),
                            visible,
                            1,
                            *selected_count,
                        )
                    },
                    "selecting GLM indexed attention rows",
                )?;
            }
        }
        if *selected_count == 0 {
            return Err("GLM indexer did not select an attention context".into());
        }
        call(
            unsafe {
                ds4_gpu_glm_qk_lowrank_typed_tensor(
                    self.scratch.qk_low.raw(),
                    self.scratch.q.raw(),
                    map,
                    size,
                    layer.k_b.offset,
                    layer.k_b.kind,
                    shape.heads as u32,
                    shape.kv_lora as u32,
                    q_nope as u32,
                    shape.key_mla as u32,
                )
            },
            "projecting the GLM low-rank query",
        )?;
        call(
            unsafe {
                ds4_gpu_glm_attention_indexed_decode_typed_tensor(
                    self.scratch.heads.raw(),
                    self.scratch.q.raw(),
                    self.scratch.qk_low.raw(),
                    cache.kv.raw(),
                    cache.rope.raw(),
                    map,
                    size,
                    layer.v_b.offset,
                    layer.v_b.kind,
                    self.scratch.indexer_selected.raw(),
                    *selected_count,
                    self.context,
                    CACHE_F16,
                    shape.heads as u32,
                    shape.kv_lora as u32,
                    q_nope as u32,
                    shape.rot as u32,
                    shape.value_mla as u32,
                    0,
                    shape.rope_base,
                    1.0,
                    0.0,
                    1.0,
                    0.0,
                    0.0,
                )
            },
            "running GLM indexed attention",
        )?;
        project(
            &self.scratch.attn_out,
            layer.output,
            shape.heads * shape.value_mla,
            shape.embd,
            &self.scratch.heads,
            map,
            size,
            self.ssd.enabled,
        )?;
        call(
            unsafe {
                ds4_gpu_add_rms_norm_weight_tensor(
                    self.scratch.ffn_norm.raw(),
                    self.scratch.after_attn.raw(),
                    self.scratch.current.raw(),
                    self.scratch.attn_out.raw(),
                    map,
                    size,
                    layer.ffn_norm.offset,
                    shape.embd as u32,
                    shape.rms_epsilon,
                )
            },
            "normalizing the GLM FFN input",
        )?;

        if let Some(dense) = layer.dense {
            project(
                &self.scratch.ffn_gate,
                dense.gate,
                shape.embd,
                shape.ff_dense,
                &self.scratch.ffn_norm,
                map,
                size,
                self.ssd.enabled,
            )?;
            project(
                &self.scratch.ffn_up,
                dense.up,
                shape.embd,
                shape.ff_dense,
                &self.scratch.ffn_norm,
                map,
                size,
                self.ssd.enabled,
            )?;
            call(
                unsafe {
                    ds4_gpu_swiglu_tensor(
                        self.scratch.ffn_mid.raw(),
                        self.scratch.ffn_gate.raw(),
                        self.scratch.ffn_up.raw(),
                        shape.ff_dense as u32,
                        0.0,
                        1.0,
                    )
                },
                "activating the dense GLM FFN",
            )?;
            project(
                &self.scratch.ffn_out,
                dense.down,
                shape.ff_dense,
                shape.embd,
                &self.scratch.ffn_mid,
                map,
                size,
                self.ssd.enabled,
            )?;
            call(
                unsafe {
                    ds4_gpu_add_tensor(
                        self.scratch.next.raw(),
                        self.scratch.after_attn.raw(),
                        self.scratch.ffn_out.raw(),
                        shape.embd as u32,
                    )
                },
                "adding the dense GLM residual",
            )?;
        } else if let Some(sparse) = layer.sparse {
            f32_project(
                &self.scratch.router_logits,
                sparse.router,
                shape.embd,
                shape.experts,
                &self.scratch.ffn_norm,
                map,
                size,
            )?;
            call(
                unsafe {
                    ds4_gpu_glm_router_select_tensor(
                        self.scratch.router_selected.raw(),
                        self.scratch.router_weights.raw(),
                        self.scratch.router_probs.raw(),
                        map,
                        size,
                        sparse.bias.offset,
                        self.scratch.router_logits.raw(),
                        shape.experts as u32,
                        shape.experts_used as u32,
                        shape.expert_weight_scale,
                    )
                },
                "routing the GLM experts",
            )?;
            let force_resident = self.ssd.enabled
                && ordinal.saturating_sub(shape.leading_dense) < self.ssd.full_layers;
            if self.ssd.enabled && !force_resident {
                let table = expert_table(map, size, ordinal, shape, sparse);
                call(
                    unsafe {
                        ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
                            &table,
                            self.scratch.router_selected.raw(),
                            shape.experts_used as u32,
                        )
                    },
                    "loading selected GLM experts",
                )?;
            }
            let (gate_expert, gate_row) = expert_layout(sparse.gate, shape.experts);
            let (up_expert, up_row) = expert_layout(sparse.up, shape.experts);
            let (down_expert, down_row) = expert_layout(sparse.down, shape.experts);
            call(
                unsafe {
                    ds4_gpu_glm_routed_moe_one_tensor(
                        self.scratch.ffn_out.raw(),
                        self.scratch.ffn_mid.raw(),
                        map,
                        size,
                        sparse.gate.offset,
                        sparse.up.offset,
                        sparse.down.offset,
                        sparse.gate.kind,
                        sparse.up.kind,
                        sparse.down.kind,
                        gate_expert,
                        gate_row,
                        up_expert,
                        up_row,
                        down_expert,
                        down_row,
                        shape.embd as u32,
                        shape.ff_expert as u32,
                        shape.embd as u32,
                        self.scratch.router_selected.raw(),
                        self.scratch.router_weights.raw(),
                        shape.experts as u32,
                        shape.experts_used as u32,
                        ordinal,
                        self.scratch.ffn_norm.raw(),
                        force_resident,
                    )
                },
                "running the routed GLM experts",
            )?;
            project(
                &self.scratch.ffn_gate,
                sparse.shared_gate,
                shape.embd,
                shape.ff_expert,
                &self.scratch.ffn_norm,
                map,
                size,
                self.ssd.enabled,
            )?;
            project(
                &self.scratch.ffn_up,
                sparse.shared_up,
                shape.embd,
                shape.ff_expert,
                &self.scratch.ffn_norm,
                map,
                size,
                self.ssd.enabled,
            )?;
            call(
                unsafe {
                    ds4_gpu_swiglu_tensor(
                        self.scratch.ffn_mid.raw(),
                        self.scratch.ffn_gate.raw(),
                        self.scratch.ffn_up.raw(),
                        shape.ff_expert as u32,
                        0.0,
                        1.0,
                    )
                },
                "activating the shared GLM expert",
            )?;
            project(
                &self.scratch.ffn_sum,
                sparse.shared_down,
                shape.ff_expert,
                shape.embd,
                &self.scratch.ffn_mid,
                map,
                size,
                self.ssd.enabled,
            )?;
            call(
                unsafe {
                    ds4_gpu_add3_tensor(
                        self.scratch.next.raw(),
                        self.scratch.after_attn.raw(),
                        self.scratch.ffn_out.raw(),
                        self.scratch.ffn_sum.raw(),
                        shape.embd as u32,
                    )
                },
                "adding the sparse GLM residual",
            )?;
        }
        Ok(())
    }

    pub(super) fn prefill(
        &mut self,
        tokens: &[i32],
        mut progress: impl FnMut(u32) -> bool,
    ) -> Result<usize, String> {
        let mut completed = 0;
        for &token in tokens {
            if !progress(self.position()) {
                break;
            }
            self.eval(token)?;
            completed += 1;
        }
        Ok(completed)
    }

    pub(super) fn logits(&self) -> &[f32] {
        &self.logits
    }
    pub(super) fn model(&self) -> &Model {
        &self.model
    }
    pub(super) fn context(&self) -> u32 {
        self.context
    }
    pub(super) fn position(&self) -> u32 {
        self.tokens.len() as u32
    }
    pub(super) fn tokens(&self) -> &[i32] {
        &self.tokens
    }
    pub(super) fn checkpoint_tag(&self) -> [u8; 32] {
        self.checkpoint_tag
    }
    pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) {
        self.checkpoint_tag = tag;
    }

    pub(super) fn reset(&mut self) -> Result<(), String> {
        self.scratch = GlmScratch::allocate(&self.model, self.context)?;
        self.caches = (0..self.weights.layers.len())
            .map(|layer| LayerCache::allocate(self.model.shape, layer, self.context))
            .collect::<Result<_, _>>()?;
        self.tokens.clear();
        self.checkpoint_tag = [0; 32];
        Ok(())
    }

    pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
        if !tokens.starts_with(&self.tokens) {
            self.reset()?;
        }
        Ok(self.tokens.len())
    }

    pub(super) fn save_checkpoint(
        &mut self,
        path: &Path,
        tag: [u8; 32],
        progress: &mut impl FnMut(u64),
    ) -> Result<(), String> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).map_err(|e| e.to_string())?;
        }
        let temporary = path.with_extension("tmp");
        let mut file = File::create(&temporary).map_err(|e| e.to_string())?;
        file.write_all(CHECKPOINT_MAGIC)
            .map_err(|e| e.to_string())?;
        let shape = self.model.shape;
        for value in [
            CHECKPOINT_VERSION,
            self.context,
            self.position(),
            self.quality as u32,
            self.weights.layers.len() as u32,
            shape.kv_lora as u32,
            shape.rot as u32,
            shape.indexer_head_dim as u32,
            shape.vocab as u32,
        ] {
            write_u32(&mut file, value)?;
        }
        write_u64(&mut file, self.model.main.len())?;
        write_u64(&mut file, self.model_modified.0)?;
        write_u32(&mut file, self.model_modified.1)?;
        file.write_all(&self.model_identity)
            .map_err(|e| e.to_string())?;
        file.write_all(&tag).map_err(|e| e.to_string())?;
        for &token in &self.tokens {
            write_u32(&mut file, token as u32)?;
        }
        for &logit in &self.logits {
            write_u32(&mut file, logit.to_bits())?;
        }
        let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
        let rows = u64::from(self.position());
        for cache in &self.caches {
            write_buffer(
                &mut file,
                &cache.kv,
                0,
                rows * shape.kv_lora * 2,
                &mut chunk,
                progress,
            )?;
            write_buffer(
                &mut file,
                &cache.rope,
                0,
                rows * shape.rot * 2,
                &mut chunk,
                progress,
            )?;
            if let Some(indexer) = &cache.indexer {
                write_buffer(
                    &mut file,
                    indexer,
                    0,
                    rows * shape.indexer_head_dim * 2,
                    &mut chunk,
                    progress,
                )?;
            }
        }
        file.sync_all().map_err(|e| e.to_string())?;
        fs::rename(&temporary, path).map_err(|e| e.to_string())?;
        self.checkpoint_tag = tag;
        Ok(())
    }

    pub(super) fn load_checkpoint(
        &mut self,
        path: &Path,
        progress: &mut impl FnMut(u64),
    ) -> Result<bool, String> {
        let mut file = match File::open(path) {
            Ok(file) => file,
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
            Err(error) => return Err(error.to_string()),
        };
        let mut magic = [0; 8];
        file.read_exact(&mut magic).map_err(|e| e.to_string())?;
        if &magic != CHECKPOINT_MAGIC
            || read_u32(&mut file)? != CHECKPOINT_VERSION
            || read_u32(&mut file)? != self.context
        {
            return Err("GLM checkpoint does not match the current executor".into());
        }
        let token_count = read_u32(&mut file)?;
        let shape = self.model.shape;
        if token_count > self.context
            || read_u32(&mut file)? != self.quality as u32
            || read_u32(&mut file)? != self.weights.layers.len() as u32
            || read_u32(&mut file)? != shape.kv_lora as u32
            || read_u32(&mut file)? != shape.rot as u32
            || read_u32(&mut file)? != shape.indexer_head_dim as u32
            || read_u32(&mut file)? != shape.vocab as u32
            || read_u64(&mut file)? != self.model.main.len()
            || read_u64(&mut file)? != self.model_modified.0
            || read_u32(&mut file)? != self.model_modified.1
        {
            return Err("GLM checkpoint was written for a different model or configuration".into());
        }
        let mut identity = [0; 32];
        file.read_exact(&mut identity).map_err(|e| e.to_string())?;
        if identity != self.model_identity {
            return Err("GLM checkpoint model identity changed".into());
        }
        let mut tag = [0; 32];
        file.read_exact(&mut tag).map_err(|e| e.to_string())?;
        let mut tokens = Vec::with_capacity(token_count as usize);
        for _ in 0..token_count {
            let token = read_u32(&mut file)?;
            if u64::from(token) >= self.model.shape.vocab {
                return Err("GLM checkpoint token is invalid".into());
            }
            tokens.push(token as i32);
        }
        let mut logits = Vec::with_capacity(shape.vocab as usize);
        for _ in 0..shape.vocab {
            logits.push(f32::from_bits(read_u32(&mut file)?));
        }
        self.reset()?;
        let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
        let rows = u64::from(token_count);
        for cache in &self.caches {
            read_buffer(
                &mut file,
                &cache.kv,
                0,
                rows * shape.kv_lora * 2,
                &mut chunk,
                progress,
            )?;
            read_buffer(
                &mut file,
                &cache.rope,
                0,
                rows * shape.rot * 2,
                &mut chunk,
                progress,
            )?;
            if let Some(indexer) = &cache.indexer {
                read_buffer(
                    &mut file,
                    indexer,
                    0,
                    rows * shape.indexer_head_dim * 2,
                    &mut chunk,
                    progress,
                )?;
            }
        }
        let mut trailing = [0];
        if file.read(&mut trailing).map_err(|e| e.to_string())? != 0 {
            return Err("GLM checkpoint has trailing data".into());
        }
        self.tokens = tokens;
        self.logits = logits;
        self.checkpoint_tag = tag;
        Ok(true)
    }
}

fn full_indexer_layer(shape: super::super::Shape, layer: usize) -> bool {
    layer < (shape.layers - shape.nextn) as usize
        && (layer < shape.leading_dense as usize || (layer >= 6 && (layer - 6).is_multiple_of(4)))
}

fn expert_layout(weight: Weight, experts: u64) -> (u64, u64) {
    let expert = weight.bytes / experts;
    (expert, expert / weight.dims[1])
}

fn expert_table(
    map: *const c_void,
    size: u64,
    layer: u32,
    shape: super::super::Shape,
    weights: SparseWeights,
) -> StreamExpertTable {
    StreamExpertTable {
        model_map: map,
        model_size: size,
        layer,
        total_experts: shape.experts as u32,
        gate_offset: weights.gate.offset,
        up_offset: weights.up.offset,
        down_offset: weights.down.offset,
        gate_expert_bytes: weights.gate.bytes / shape.experts,
        down_expert_bytes: weights.down.bytes / shape.experts,
    }
}

fn configure_streaming(
    model: &Model,
    weights: &GlmWeights,
    ssd: EngineSsdSettings,
) -> Result<(), String> {
    if !ssd.enabled {
        return Ok(());
    }
    let sparse = weights
        .layers
        .iter()
        .find_map(|layer| layer.sparse)
        .ok_or("GLM model has no routed expert layers")?;
    let gate = sparse.gate.bytes / model.shape.experts;
    let up = sparse.up.bytes / model.shape.experts;
    let down = sparse.down.bytes / model.shape.experts;
    let per_expert = gate
        .checked_add(up)
        .and_then(|bytes| bytes.checked_add(down))
        .ok_or("GLM expert size overflow")?;
    let budget = if ssd.cache_experts != 0 {
        ssd.cache_experts
    } else if ssd.cache_bytes != 0 {
        u32::try_from((ssd.cache_bytes / per_expert).max(1)).unwrap_or(u32::MAX)
    } else {
        unsafe { ds4_gpu_stream_expert_cache_budget_for_expert_size(gate, down) }
    };
    if budget == 0 {
        return Err("GLM SSD streaming has no memory for an expert cache".into());
    }
    unsafe {
        ds4_gpu_set_streaming_expert_cache_expert_bytes(per_expert);
        ds4_gpu_set_streaming_expert_cache_budget(budget);
    }
    Ok(())
}

fn admission_bytes(
    model: &Model,
    weights: &GlmWeights,
    context: u32,
    ssd: EngineSsdSettings,
) -> Result<u64, String> {
    let shape = model.shape;
    let normal_layers = u64::from(shape.layers - shape.nextn);
    let indexer_layers = (0..normal_layers as usize)
        .filter(|layer| full_indexer_layer(shape, *layer))
        .count() as u64;
    let per_layer = u64::from(context)
        .checked_mul((shape.kv_lora + shape.rot) * 2)
        .ok_or("GLM compact-cache size overflow")?;
    let kv = normal_layers
        .checked_mul(per_layer)
        .and_then(|bytes| {
            bytes.checked_add(indexer_layers * u64::from(context) * shape.indexer_head_dim * 2)
        })
        .ok_or("GLM compact-cache size overflow")?;
    let resident = if ssd.enabled {
        model
            .main
            .tensors
            .iter()
            .filter(|(name, _)| {
                !name.ends_with("ffn_gate_exps.weight")
                    && !name.ends_with("ffn_up_exps.weight")
                    && !name.ends_with("ffn_down_exps.weight")
            })
            .try_fold(0_u64, |total, (_, tensor)| total.checked_add(tensor.bytes))
            .ok_or("GLM resident tensor size overflow")?
    } else {
        model.main.len() - model.main.data_offset()
    };
    let sparse_expert = weights
        .layers
        .iter()
        .find_map(|layer| layer.sparse)
        .map(|weights| {
            weights.gate.bytes / shape.experts
                + weights.up.bytes / shape.experts
                + weights.down.bytes / shape.experts
        })
        .unwrap_or(0);
    let cache = if ssd.enabled {
        if ssd.cache_bytes != 0 {
            ssd.cache_bytes
        } else if ssd.cache_experts != 0 {
            sparse_expert.saturating_mul(u64::from(ssd.cache_experts))
        } else {
            // Reference auto policy reserves roughly one seventh of the available
            // budget for full layers and uses the remainder for selected experts.
            16 * 1024 * 1024 * 1024_u64
        }
    } else {
        0
    };
    let full_layers = sparse_expert
        .saturating_mul(shape.experts)
        .saturating_mul(u64::from(ssd.full_layers));
    resident
        .checked_add(kv)
        .and_then(|bytes| bytes.checked_add(cache))
        .and_then(|bytes| bytes.checked_add(full_layers))
        .and_then(|bytes| bytes.checked_add(512 * 1024 * 1024))
        .ok_or_else(|| "GLM runtime memory size overflow".into())
}

fn norm(
    out: &Buffer,
    input: &Buffer,
    weight: Weight,
    width: u32,
    epsilon: f32,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_rms_norm_weight_tensor(
                out.raw(),
                input.raw(),
                map,
                size,
                weight.offset,
                width,
                epsilon,
            )
        },
        "normalizing GLM activations",
    )
}

#[allow(clippy::too_many_arguments)]
fn project(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    map: *const c_void,
    size: u64,
    streaming: bool,
) -> Result<(), String> {
    let result = unsafe {
        if streaming {
            ds4_gpu_matmul_quant_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                weight.kind,
                input,
                output,
                x.raw(),
                1,
            )
        } else {
            ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                weight.kind,
                input,
                output,
                x.raw(),
                1,
            )
        }
    };
    call(result, "projecting GLM activations")
}

#[allow(clippy::too_many_arguments)]
fn f32_project(
    out: &Buffer,
    weight: Weight,
    input: u64,
    output: u64,
    x: &Buffer,
    map: *const c_void,
    size: u64,
) -> Result<(), String> {
    call(
        unsafe {
            ds4_gpu_matmul_f32_tensor(
                out.raw(),
                map,
                size,
                weight.offset,
                input,
                output,
                x.raw(),
                1,
            )
        },
        "projecting GLM F32 activations",
    )
}

#[cfg(test)]
mod tests {
    use super::{GlmExecutor, full_indexer_layer};
    use crate::engine::{GLM, Model};
    use crate::model::ModelChoice;
    use crate::settings::EngineSsdSettings;
    use std::path::Path;

    #[test]
    fn dsa_indexer_schedule_matches_the_reference() {
        assert!(full_indexer_layer(GLM, 0));
        assert!(full_indexer_layer(GLM, 2));
        assert!(!full_indexer_layer(GLM, 3));
        assert!(full_indexer_layer(GLM, 6));
        assert!(full_indexer_layer(GLM, 10));
        assert!(!full_indexer_layer(GLM, 11));
        assert!(!full_indexer_layer(GLM, 78));
    }

    #[test]
    #[ignore = "requires the 197 GiB GLM 5.2 checkpoint and Apple Metal"]
    fn resident_and_streamed_glm_match_the_short_code_fixture() {
        let path = std::env::var("DS4_GLM_MODEL").unwrap_or_else(|_| {
            "../ds4/models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf".into()
        });
        let prompt = b"Complete the C statement with the next exact token only:\nreturn snprintf(buf, sizeof(buf), \"%d\", value";
        for streamed in [false, true] {
            let model = Model::open_main(Path::new(&path), ModelChoice::Glm52).unwrap();
            let tokens = model.tokenize(std::str::from_utf8(prompt).unwrap());
            let mut executor = GlmExecutor::open(
                model,
                4096,
                true,
                EngineSsdSettings {
                    enabled: streamed,
                    cold: false,
                    cache_experts: 0,
                    cache_bytes: 0,
                    full_layers: 0,
                    full_layers_set: false,
                    preload_experts: 0,
                },
            )
            .unwrap();
            assert_eq!(executor.prefill(&tokens, |_| true).unwrap(), tokens.len());
            let token = executor
                .logits()
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.total_cmp(b.1))
                .unwrap()
                .0 as i32;
            assert_eq!(
                executor.model().token_bytes(token).as_deref(),
                Some(b")".as_slice())
            );
        }
    }
}