mod gguf;
#[cfg(target_os = "macos")]
mod kvstore;
#[cfg(target_os = "macos")]
mod metal;
mod tokenizer;
mod validation;

#[cfg(target_os = "macos")]
use crate::metrics::{KvLookup, Metrics};
use crate::model::ModelChoice;
use crate::settings::TurnSettings;
use crate::settings::{EngineSettings, ReasoningMode};
use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
#[cfg(target_os = "macos")]
use kvstore::{KvStore, StoreReason};
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
#[cfg(target_os = "macos")]
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(target_os = "macos")]
use std::time::Instant;
use tokenizer::Tokenizer;
use validation::{validate_dspark, validate_main};

pub(crate) use validation::validate_model_artifact;

#[cfg(target_os = "macos")]
pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache;
#[cfg(target_os = "macos")]
pub(crate) use metal::configure_sources as configure_metal_sources;

const DENSE: &[u32] = &[Q8_0, Q4_K, Q4_0];
const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K];
const PLAIN: &[u32] = &[F16, F32];
const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0];

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ModelFamily {
    DeepSeek,
    Glm,
}

#[derive(Clone, Copy)]
struct Shape {
    model: ModelChoice,
    family: ModelFamily,
    layers: u32,
    embd: u64,
    vocab: u64,
    heads: u64,
    head_kv: u64,
    head_dim: u64,
    value_dim: u64,
    rot: u64,
    out_groups: u64,
    lora_q: u64,
    lora_o: u64,
    experts: u64,
    experts_used: u64,
    expert_shared: u64,
    ff_expert: u64,
    ff_dense: u64,
    hash_layers: u32,
    sliding_window: u64,
    indexer_heads: u64,
    indexer_head_dim: u64,
    indexer_top_k: u64,
    hc: u64,
    hc_sinkhorn: u64,
    nextn: u32,
    leading_dense: u32,
    kv_lora: u64,
    key_mla: u64,
    value_mla: u64,
    rms_epsilon: f32,
    hc_epsilon: f32,
    expert_weight_scale: f32,
    swiglu_clamp: f32,
    rope_base: f32,
    rope_scale: f32,
    rope_beta_fast: f32,
    rope_beta_slow: f32,
    compress_rope_base: f32,
    original_context: u64,
}

const FLASH: Shape = Shape {
    model: ModelChoice::DeepSeekV4Flash,
    family: ModelFamily::DeepSeek,
    layers: 43,
    embd: 4096,
    vocab: 129_280,
    heads: 64,
    head_kv: 1,
    head_dim: 512,
    value_dim: 512,
    rot: 64,
    out_groups: 8,
    lora_q: 1024,
    lora_o: 1024,
    experts: 256,
    experts_used: 6,
    expert_shared: 1,
    ff_expert: 2048,
    ff_dense: 0,
    hash_layers: 3,
    sliding_window: 128,
    indexer_heads: 64,
    indexer_head_dim: 128,
    indexer_top_k: 512,
    hc: 4,
    hc_sinkhorn: 20,
    nextn: 0,
    leading_dense: 0,
    kv_lora: 0,
    key_mla: 0,
    value_mla: 0,
    rms_epsilon: 1.0e-6,
    hc_epsilon: 1.0e-6,
    expert_weight_scale: 1.5,
    swiglu_clamp: 10.0,
    rope_base: 10_000.0,
    rope_scale: 16.0,
    rope_beta_fast: 32.0,
    rope_beta_slow: 1.0,
    compress_rope_base: 160_000.0,
    original_context: 65_536,
};

const PRO: Shape = Shape {
    model: ModelChoice::DeepSeekV4Pro,
    layers: 61,
    embd: 7168,
    heads: 128,
    out_groups: 16,
    lora_q: 1536,
    experts: 384,
    ff_expert: 3072,
    indexer_top_k: 1024,
    expert_weight_scale: 2.5,
    ..FLASH
};

const GLM: Shape = Shape {
    model: ModelChoice::Glm52,
    family: ModelFamily::Glm,
    layers: 79,
    embd: 6144,
    vocab: 154_880,
    heads: 64,
    head_kv: 1,
    head_dim: 576,
    value_dim: 512,
    rot: 64,
    out_groups: 0,
    lora_q: 2048,
    lora_o: 0,
    experts: 256,
    experts_used: 8,
    expert_shared: 1,
    ff_expert: 2048,
    ff_dense: 12_288,
    hash_layers: 0,
    sliding_window: 0,
    indexer_heads: 32,
    indexer_head_dim: 128,
    indexer_top_k: 2048,
    hc: 0,
    hc_sinkhorn: 0,
    nextn: 1,
    leading_dense: 3,
    kv_lora: 512,
    key_mla: 256,
    value_mla: 256,
    rms_epsilon: 1.0e-5,
    hc_epsilon: 0.0,
    expert_weight_scale: 2.5,
    swiglu_clamp: 0.0,
    rope_base: 8_000_000.0,
    rope_scale: 1.0,
    rope_beta_fast: 0.0,
    rope_beta_slow: 0.0,
    compress_rope_base: 0.0,
    original_context: 1_048_576,
};

pub(crate) struct Model {
    main: Gguf,
    support: Option<Gguf>,
    shape: Shape,
    tokenizer: Tokenizer,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct ModelSummary {
    pub(crate) model: ModelChoice,
    pub(crate) mapped_bytes: u64,
    pub(crate) tensor_count: usize,
    pub(crate) vocabulary_size: usize,
    pub(crate) support_loaded: bool,
}

impl Model {
    #[allow(dead_code)]
    pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
        let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
        if let Some(path) = &settings.artifacts.mtp {
            let support = Gguf::open(path)?;
            validate_dspark(&support, &model.shape)?;
            model.support = Some(support);
        }
        Ok(model)
    }

    fn open_main(path: &Path, expected: ModelChoice) -> Result<Self, String> {
        let main = Gguf::open(path)?;
        let shape = validate_main(&main, expected)?;
        let tokenizer = Tokenizer::load(&main, shape.family)?;
        if tokenizer.vocab_size() != shape.vocab as usize {
            return Err(format!(
                "tokenizer has {} entries, expected {}",
                tokenizer.vocab_size(),
                shape.vocab
            ));
        }
        Ok(Self {
            main,
            support: None,
            shape,
            tokenizer,
        })
    }

    pub(crate) fn summary(&self) -> ModelSummary {
        ModelSummary {
            model: self.shape.model,
            mapped_bytes: self.main.len() + self.support.as_ref().map_or(0, Gguf::len),
            tensor_count: self.main.tensors.len()
                + self
                    .support
                    .as_ref()
                    .map_or(0, |support| support.tensors.len()),
            vocabulary_size: self.tokenizer.vocab_size(),
            support_loaded: self.support.is_some(),
        }
    }

    pub(crate) fn tokenize(&self, text: &str) -> Vec<i32> {
        self.tokenizer.tokenize(text)
    }

    pub(crate) fn render_prompt(
        &self,
        system: &str,
        prompt: &str,
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.tokenizer.encode_chat(system, prompt, reasoning)
    }

    fn render_conversation(
        &self,
        system: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.tokenizer
            .encode_conversation(system, messages, reasoning)
    }

    fn render_continuation(
        &self,
        prompt: &str,
        reasoning: ReasoningMode,
        skip_previous_eos: bool,
    ) -> Vec<i32> {
        self.tokenizer
            .encode_continuation(prompt, reasoning, skip_previous_eos)
    }

    pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
        self.tokenizer.token_bytes(token)
    }

    pub(crate) fn eos_token(&self) -> i32 {
        self.tokenizer.eos()
    }

    pub(crate) fn is_stop_token(&self, token: i32) -> bool {
        self.tokenizer.is_stop(token)
    }

    pub(crate) fn is_think_start_token(&self, token: i32) -> bool {
        self.tokenizer.is_think_start(token)
    }

    pub(crate) fn is_think_end_token(&self, token: i32) -> bool {
        self.tokenizer.is_think_end(token)
    }

    pub(crate) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool {
        self.is_stop_token(token)
            || (reasoning == ReasoningMode::Direct
                && (self.is_think_start_token(token) || self.is_think_end_token(token)))
    }

    pub(crate) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
        self.main.tensor_data(name)
    }
}

#[cfg(target_os = "macos")]
pub(crate) struct Generator {
    executor: metal::Executor,
    checkpoint: Option<PathBuf>,
    /// Token frontier of the last transient store, so continued checkpoints are
    /// spaced like ds4's `continued_last_store_tokens`.
    last_store_tokens: u32,
    metrics: Arc<Metrics>,
}

#[derive(Clone)]
pub(crate) struct ChatTurn {
    pub(crate) user: bool,
    pub(crate) tool: bool,
    pub(crate) skip_previous_eos: bool,
    pub(crate) reasoning: Option<String>,
    pub(crate) reasoning_complete: bool,
    pub(crate) content: String,
}

pub(crate) struct GenerationOutput {
    pub(crate) message: ChatTurn,
    pub(crate) prompt_tokens: u32,
    pub(crate) cached_tokens: u32,
    pub(crate) completion_tokens: u32,
    pub(crate) finish_reason: &'static str,
    pub(crate) previous_checkpoint_bytes: Option<u64>,
    pub(crate) checkpoint_bytes: u64,
}

#[cfg(target_os = "macos")]
impl Generator {
    pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
        if settings.model != ModelChoice::DeepSeekV4Flash {
            return Err("local generation currently supports DeepSeek V4 Flash only".into());
        }
        if settings.speculative.dspark || settings.ssd.enabled || settings.steering.file.is_some() {
            return Err(
                "DSpark, SSD streaming, and steering are not yet available in the Rust executor"
                    .into(),
            );
        }
        let model = Model::open(settings)?;
        let executor = metal::Executor::open(
            model,
            settings.context_tokens.max(1) as u32,
            settings.execution.quality,
            settings.execution.prefill_chunk,
        )?;
        Ok(Self {
            executor,
            checkpoint: None,
            last_store_tokens: 0,
            metrics,
        })
    }

    pub(crate) fn summary(&self) -> ModelSummary {
        self.executor.model().summary()
    }

    pub(crate) fn generate(
        &mut self,
        checkpoint: &Path,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        mut emit: impl FnMut(bool, String),
        mut progress: impl FnMut(u32, u32, Option<f32>),
    ) -> Result<GenerationOutput, String> {
        let history = messages
            .split_last()
            .map_or(messages, |(_, history)| history);
        self.select_checkpoint(
            checkpoint,
            conversation_tag(&settings.system_prompt, settings.reasoning_mode, history),
        )?;
        let (mut output, prompt_complete) =
            self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
        let mut completed = messages.to_vec();
        completed.push(output.message.clone());
        output.previous_checkpoint_bytes =
            std::fs::metadata(checkpoint).ok().map(|item| item.len());
        self.save_checkpoint(
            checkpoint,
            if prompt_complete {
                conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
            } else {
                [0; 32]
            },
        )?;
        output.checkpoint_bytes = std::fs::metadata(checkpoint)
            .map(|item| item.len())
            .unwrap_or(0);
        Ok(output)
    }

    pub(crate) fn generate_transient(
        &mut self,
        directory: &Path,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        mut emit: impl FnMut(bool, String),
        mut progress: impl FnMut(u32, u32, Option<f32>),
    ) -> Result<GenerationOutput, String> {
        let history = messages
            .split_last()
            .map_or(messages, |(_, history)| history);
        let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
        let history_key =
            conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
        let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
        let mut previous_checkpoint = self.checkpoint.clone();
        if self.executor.checkpoint_tag() != history_tag {
            if let Some(entry) = store.find(&history_key, self.executor.context()) {
                if self.select_checkpoint(&entry.checkpoint, entry.tag)? {
                    store.touch(&entry)?;
                    self.last_store_tokens = entry.tokens;
                    previous_checkpoint = Some(entry.checkpoint);
                } else {
                    store.discard(&entry);
                    self.last_store_tokens = 0;
                    previous_checkpoint = None;
                }
            } else {
                self.executor.reset()?;
                self.checkpoint = None;
                self.last_store_tokens = 0;
                previous_checkpoint = None;
                self.metrics.kv_lookup(KvLookup::Miss);
            }
        } else {
            self.metrics.kv_lookup(KvLookup::MemoryHit);
        }

        let (mut output, prompt_complete) =
            self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
        let mut completed = messages.to_vec();
        completed.push(output.message.clone());
        output.previous_checkpoint_bytes = previous_checkpoint
            .as_deref()
            .and_then(|path| std::fs::metadata(path).ok())
            .map(|item| item.len());
        if prompt_complete {
            let completed_key =
                conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
            let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into();
            if !settings.kv_cache.stores(
                self.executor.position(),
                history.is_empty(),
                self.last_store_tokens,
            ) {
                // A gated store still leaves the finished conversation in the
                // live KV, so mark it and drop the stale file association.
                self.executor.note_checkpoint_tag(completed_tag);
                self.checkpoint = None;
                return Ok(output);
            }
            let completed_checkpoint = store.checkpoint_path(&completed_key);
            self.save_checkpoint(&completed_checkpoint, completed_tag)?;
            let retained = store.record(
                &completed_checkpoint,
                &completed_key,
                completed_tag,
                self.executor.position(),
                self.executor.context(),
                if history.is_empty() {
                    StoreReason::Cold
                } else {
                    StoreReason::Continued
                },
            )?;
            output.checkpoint_bytes = retained
                .then(|| std::fs::metadata(&completed_checkpoint).ok())
                .flatten()
                .map_or(0, |item| item.len());
            if retained {
                self.last_store_tokens = self.executor.position();
            }
            self.checkpoint = retained.then_some(completed_checkpoint);
        }
        Ok(output)
    }

    fn select_checkpoint(
        &mut self,
        checkpoint: &Path,
        expected_tag: [u8; 32],
    ) -> Result<bool, String> {
        if self.checkpoint.as_deref() == Some(checkpoint) {
            let found = self.executor.checkpoint_tag() == expected_tag;
            self.metrics.kv_lookup(if found {
                KvLookup::MemoryHit
            } else {
                KvLookup::Miss
            });
            return Ok(found);
        }
        self.executor.reset()?;
        self.metrics.kv_read_started();
        let started = Instant::now();
        let loaded = self
            .executor
            .load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes));
        self.metrics
            .kv_read_finished(started.elapsed(), loaded.is_err());
        let found = matches!(loaded, Ok(true)) && self.executor.checkpoint_tag() == expected_tag;
        let lookup = match loaded {
            Ok(true) if found => KvLookup::DiskHit,
            Ok(true) | Ok(false) => KvLookup::Miss,
            Err(_) => {
                self.executor.reset()?;
                let _ = std::fs::remove_file(checkpoint);
                KvLookup::Invalid
            }
        };
        self.metrics.kv_lookup(lookup);
        self.checkpoint = Some(checkpoint.to_owned());
        Ok(found)
    }

    fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
        self.metrics.kv_write_started();
        let started = Instant::now();
        let result = self
            .executor
            .save_checkpoint(checkpoint, tag, &mut |bytes| {
                self.metrics.kv_write_bytes(bytes);
            });
        self.metrics
            .kv_write_finished(started.elapsed(), result.is_err());
        result
    }

    fn generate_inner(
        &mut self,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        emit: &mut impl FnMut(bool, String),
        progress: &mut impl FnMut(u32, u32, Option<f32>),
    ) -> Result<(GenerationOutput, bool), String> {
        let tokens = match messages.split_last() {
            Some((latest, history))
                if latest.user
                    && self.executor.checkpoint_tag()
                        == conversation_tag(
                            &settings.system_prompt,
                            settings.reasoning_mode,
                            history,
                        ) =>
            {
                let mut tokens = self.executor.tokens().to_vec();
                tokens.extend(self.executor.model().render_continuation(
                    &latest.content,
                    settings.reasoning_mode,
                    latest.skip_previous_eos,
                ));
                tokens
            }
            _ => self.executor.model().render_conversation(
                &settings.system_prompt,
                messages,
                settings.reasoning_mode,
            ),
        };
        if tokens.is_empty() {
            return Err("the rendered prompt is empty".into());
        }
        let max_context = self.executor.context() as usize;
        if tokens.len() >= max_context {
            return Err(format!(
                "Prompt has {} tokens, but the configured context size is {max_context} tokens",
                tokens.len()
            ));
        }
        let reused = self.executor.align_prompt(&tokens)?;
        self.metrics.kv_prefix_reused(reused);
        progress(self.executor.position(), self.executor.context(), None);
        let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
        let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
        let mut generated = ChatTurn {
            user: false,
            tool: false,
            skip_previous_eos: false,
            reasoning: reasoning.then(String::new),
            reasoning_complete: !reasoning,
            content: String::new(),
        };
        let mut emitted_reasoning = 0;
        let mut emitted_content = 0;
        let mut pending_utf8 = Vec::new();
        let prompt_tokens = tokens.len();
        let suffix = &tokens[reused..];
        let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
            let context = self.executor.context();
            self.executor.prefill(suffix, |used| {
                progress(used, context, None);
                !cancelled.load(Ordering::Relaxed)
            })?
        } else {
            let mut completed = 0;
            for &token in suffix {
                if cancelled.load(Ordering::Relaxed) {
                    break;
                }
                self.executor.eval(token)?;
                completed += 1;
                progress(self.executor.position(), self.executor.context(), None);
            }
            completed
        };
        if completed != suffix.len() {
            return Ok((
                GenerationOutput {
                    message: generated,
                    prompt_tokens: prompt_tokens as u32,
                    cached_tokens: reused as u32,
                    completion_tokens: 0,
                    finish_reason: "stop",
                    previous_checkpoint_bytes: None,
                    checkpoint_bytes: 0,
                },
                false,
            ));
        }
        progress(self.executor.position(), self.executor.context(), Some(0.0));
        let generation_started = Instant::now();
        let mut generated_tokens = 0_u32;
        for _ in 0..settings
            .max_generated_tokens
            .max(0)
            .min((max_context - self.executor.position() as usize) as i32)
        {
            if cancelled.load(Ordering::Relaxed) {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                flush_generated(
                    &mut generated,
                    &mut emitted_reasoning,
                    &mut emitted_content,
                    &settings.stops,
                    emit,
                );
                return Ok((
                    GenerationOutput {
                        message: generated,
                        prompt_tokens: prompt_tokens as u32,
                        cached_tokens: reused as u32,
                        completion_tokens: generated_tokens,
                        finish_reason: "stop",
                        previous_checkpoint_bytes: None,
                        checkpoint_bytes: 0,
                    },
                    true,
                ));
            }
            let token = sample(
                self.executor.logits(),
                settings.temperature,
                settings.top_p,
                settings.min_p,
                settings.top_k,
                &mut rng,
            );
            if self
                .executor
                .model()
                .is_stop_token_for_reasoning(token, settings.reasoning_mode)
            {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                flush_generated(
                    &mut generated,
                    &mut emitted_reasoning,
                    &mut emitted_content,
                    &settings.stops,
                    emit,
                );
                return Ok((
                    GenerationOutput {
                        message: generated,
                        prompt_tokens: prompt_tokens as u32,
                        cached_tokens: reused as u32,
                        completion_tokens: generated_tokens,
                        finish_reason: "stop",
                        previous_checkpoint_bytes: None,
                        checkpoint_bytes: 0,
                    },
                    true,
                ));
            }
            if self.executor.model().is_think_start_token(token) {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                reasoning = true;
                generated.reasoning.get_or_insert_default();
            } else if self.executor.model().is_think_end_token(token) {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                if let Some(reasoning_text) = &mut generated.reasoning
                    && emit_safe_text(
                        reasoning_text,
                        &mut emitted_reasoning,
                        &settings.stops,
                        true,
                        true,
                        emit,
                    )
                {
                    return Ok((
                        GenerationOutput {
                            message: generated,
                            prompt_tokens: prompt_tokens as u32,
                            cached_tokens: reused as u32,
                            completion_tokens: generated_tokens + 1,
                            finish_reason: "stop",
                            previous_checkpoint_bytes: None,
                            checkpoint_bytes: 0,
                        },
                        false,
                    ));
                }
                reasoning = false;
                generated.reasoning_complete = true;
                emit(false, String::new());
            } else if let Some(bytes) = self.executor.model().token_bytes(token) {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &bytes, false);
                let stopped = if reasoning {
                    let text = generated.reasoning.get_or_insert_default();
                    emit_safe_text(
                        text,
                        &mut emitted_reasoning,
                        &settings.stops,
                        false,
                        true,
                        emit,
                    )
                } else {
                    generated.reasoning_complete = true;
                    emit_safe_text(
                        &mut generated.content,
                        &mut emitted_content,
                        &settings.stops,
                        false,
                        false,
                        emit,
                    )
                };
                if stopped {
                    return Ok((
                        GenerationOutput {
                            message: generated,
                            prompt_tokens: prompt_tokens as u32,
                            cached_tokens: reused as u32,
                            completion_tokens: generated_tokens + 1,
                            finish_reason: "stop",
                            previous_checkpoint_bytes: None,
                            checkpoint_bytes: 0,
                        },
                        false,
                    ));
                }
            }
            self.executor.eval(token)?;
            generated_tokens += 1;
            progress(
                self.executor.position(),
                self.executor.context(),
                Some(
                    generated_tokens as f32
                        / generation_started.elapsed().as_secs_f32().max(1.0e-6),
                ),
            );
        }
        append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
        flush_generated(
            &mut generated,
            &mut emitted_reasoning,
            &mut emitted_content,
            &settings.stops,
            emit,
        );
        Ok((
            GenerationOutput {
                message: generated,
                prompt_tokens: prompt_tokens as u32,
                cached_tokens: reused as u32,
                completion_tokens: generated_tokens,
                finish_reason: "length",
                previous_checkpoint_bytes: None,
                checkpoint_bytes: 0,
            },
            true,
        ))
    }
}

#[cfg(target_os = "macos")]
fn append_generated_bytes(
    generated: &mut ChatTurn,
    reasoning: bool,
    pending: &mut Vec<u8>,
    bytes: &[u8],
    final_flush: bool,
) {
    let text = if reasoning {
        generated.reasoning.get_or_insert_default()
    } else {
        &mut generated.content
    };
    pending.extend_from_slice(bytes);
    loop {
        match std::str::from_utf8(pending) {
            Ok(valid) => {
                text.push_str(valid);
                pending.clear();
                return;
            }
            Err(error) => {
                let valid = error.valid_up_to();
                text.push_str(std::str::from_utf8(&pending[..valid]).unwrap());
                pending.drain(..valid);
                match error.error_len() {
                    Some(length) => {
                        text.push('\u{fffd}');
                        pending.drain(..length);
                    }
                    None if final_flush => {
                        text.push_str(&String::from_utf8_lossy(pending));
                        pending.clear();
                        return;
                    }
                    None => return,
                }
            }
        }
    }
}

#[cfg(target_os = "macos")]
fn flush_generated(
    generated: &mut ChatTurn,
    emitted_reasoning: &mut usize,
    emitted_content: &mut usize,
    stops: &[String],
    emit: &mut impl FnMut(bool, String),
) {
    if let Some(reasoning) = &mut generated.reasoning {
        let _ = emit_safe_text(reasoning, emitted_reasoning, stops, true, true, emit);
    }
    let _ = emit_safe_text(
        &mut generated.content,
        emitted_content,
        stops,
        true,
        false,
        emit,
    );
}

#[cfg(target_os = "macos")]
fn emit_safe_text(
    text: &mut String,
    emitted: &mut usize,
    stops: &[String],
    final_flush: bool,
    reasoning: bool,
    emit: &mut impl FnMut(bool, String),
) -> bool {
    let stop = stops
        .iter()
        .filter_map(|stop| {
            text[*emitted..]
                .find(stop)
                .map(|position| *emitted + position)
        })
        .min();
    if let Some(stop) = stop {
        if stop > *emitted {
            emit(reasoning, text[*emitted..stop].to_owned());
        }
        text.truncate(stop);
        *emitted = stop;
        return true;
    }
    let hold = if final_flush {
        0
    } else {
        stops
            .iter()
            .map(|stop| stop.len().saturating_sub(1))
            .max()
            .unwrap_or(0)
    };
    let mut safe = text.len().saturating_sub(hold);
    while safe > *emitted && !text.is_char_boundary(safe) {
        safe -= 1;
    }
    if safe > *emitted {
        emit(reasoning, text[*emitted..safe].to_owned());
        *emitted = safe;
    }
    false
}

#[cfg(target_os = "macos")]
fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> Vec<u8> {
    fn text(output: &mut Vec<u8>, value: &str) {
        output.extend_from_slice(&(value.len() as u64).to_le_bytes());
        output.extend_from_slice(value.as_bytes());
    }

    let mut output = b"DS4Server chat checkpoint v3".to_vec();
    text(&mut output, system);
    output.push(match reasoning {
        ReasoningMode::Direct => 0,
        ReasoningMode::High => 1,
        ReasoningMode::Max => 2,
    });
    for message in messages {
        output.push(u8::from(message.user));
        output.push(u8::from(message.tool));
        output.push(u8::from(message.skip_previous_eos));
        match &message.reasoning {
            Some(reasoning) => {
                output.push(1);
                text(&mut output, reasoning);
            }
            None => output.push(0),
        }
        output.push(u8::from(message.reasoning_complete));
        text(&mut output, &message.content);
    }
    output
}

#[cfg(target_os = "macos")]
fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
    Sha256::digest(conversation_key(system, reasoning, messages)).into()
}

#[cfg(target_os = "macos")]
fn sample(
    logits: &[f32],
    temperature: f32,
    top_p: f32,
    min_p: f32,
    top_k: i32,
    rng: &mut Rng,
) -> i32 {
    if temperature <= 0.0 {
        return logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .map_or(0, |(index, _)| index as i32);
    }
    let maximum = logits
        .iter()
        .copied()
        .filter(|value| value.is_finite())
        .fold(f32::NEG_INFINITY, f32::max);
    if !maximum.is_finite() {
        return 0;
    }
    let top_p = if top_p <= 0.0 || top_p > 1.0 {
        1.0
    } else {
        top_p
    };
    let min_p = min_p.max(0.0);
    let mut probabilities: Vec<(usize, f32)> = logits
        .iter()
        .enumerate()
        .filter(|(_, logit)| logit.is_finite())
        .map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
        .filter(|(_, probability)| *probability >= min_p)
        .collect();
    if probabilities.is_empty() {
        return logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .map_or(0, |(index, _)| index as i32);
    }
    if top_p < 1.0 || top_k > 0 {
        probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
        if top_k > 0 {
            probabilities.truncate(probabilities.len().min(top_k as usize));
        }
    }
    if top_p < 1.0 {
        let total: f32 = probabilities
            .iter()
            .map(|(_, probability)| probability)
            .sum();
        let mut kept = 0.0;
        let count = probabilities
            .iter()
            .position(|(_, probability)| {
                kept += *probability;
                kept / total >= top_p
            })
            .map_or(probabilities.len(), |index| index + 1);
        probabilities.truncate(count);
    }
    let kept_total: f32 = probabilities.iter().map(|(_, p)| p).sum();
    let mut choice = rng.unit() * kept_total;
    for (token, probability) in &probabilities {
        choice -= probability;
        if choice <= 0.0 {
            return *token as i32;
        }
    }
    probabilities.last().map_or(0, |(token, _)| *token as i32)
}

#[cfg(target_os = "macos")]
struct Rng(u64);

#[cfg(target_os = "macos")]
impl Rng {
    fn new(seed: u64) -> Self {
        Self(seed.max(1))
    }

    fn unit(&mut self) -> f32 {
        let mut value = self.0;
        if value == 0 {
            value = 0x9e37_79b9_7f4a_7c15;
        }
        value ^= value >> 12;
        value ^= value << 25;
        value ^= value >> 27;
        self.0 = value;
        let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d);
        ((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0
    }
}

#[cfg(all(test, target_os = "macos"))]
mod sampling_tests {
    use super::*;

    #[test]
    fn zero_temperature_is_greedy() {
        let mut rng = Rng::new(1);
        assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, 0, &mut rng), 1);
    }

    #[test]
    fn top_k_and_stream_stops_are_applied_before_output() {
        let mut rng = Rng::new(1);
        assert_eq!(sample(&[1.0, 4.0, 2.0], 1.0, 1.0, 0.0, 1, &mut rng), 1);

        let mut text = "hello STOP hidden".to_owned();
        let mut emitted = 0;
        let mut chunks = Vec::new();
        assert!(emit_safe_text(
            &mut text,
            &mut emitted,
            &["STOP".into()],
            false,
            false,
            &mut |_, chunk| chunks.push(chunk),
        ));
        assert_eq!(text, "hello ");
        assert_eq!(chunks, ["hello "]);
    }

    #[test]
    fn split_utf8_token_bytes_are_joined_before_decoding() {
        let mut generated = ChatTurn {
            user: false,
            tool: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: String::new(),
        };
        let mut pending = Vec::new();

        append_generated_bytes(&mut generated, false, &mut pending, &[0xf0, 0x9f], false);
        assert!(generated.content.is_empty());
        append_generated_bytes(&mut generated, false, &mut pending, &[0x98, 0x8a], false);

        assert_eq!(generated.content, "😊");
        assert!(pending.is_empty());
    }

    #[test]
    fn checkpoint_tag_covers_the_canonical_chat_state() {
        let mut messages = vec![ChatTurn {
            user: true,
            tool: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: "Hello".into(),
        }];
        let tag = conversation_tag("System", ReasoningMode::High, &messages);
        assert_eq!(
            tag,
            conversation_tag("System", ReasoningMode::High, &messages)
        );
        assert_ne!(
            tag,
            conversation_tag("Changed", ReasoningMode::High, &messages)
        );
        assert_ne!(
            tag,
            conversation_tag("System", ReasoningMode::Direct, &messages)
        );
        messages[0].skip_previous_eos = true;
        assert_ne!(
            tag,
            conversation_tag("System", ReasoningMode::High, &messages)
        );
        let prefix = conversation_key("System", ReasoningMode::High, &messages);
        messages.push(ChatTurn {
            user: false,
            tool: false,
            skip_previous_eos: false,
            reasoning: Some("because".into()),
            reasoning_complete: true,
            content: "Hi".into(),
        });
        assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix));
    }

    #[test]
    #[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"]
    fn metal_executes_real_flash_token() {
        configure_metal_sources().unwrap();
        let path = Path::new(env!("CARGO_MANIFEST_DIR")).join(
            "../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
        );
        let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap();
        let tokens = model.render_prompt(
            "You are a helpful assistant",
            "Hello",
            ReasoningMode::Direct,
        );
        assert_eq!(tokens.len(), 10);
        let mut executor = metal::Executor::open(model, 32_768, false, 0).unwrap();
        assert_eq!(executor.context(), 32_768);
        for &token in &tokens {
            executor.eval(token).unwrap();
        }
        assert!(executor.logits().iter().all(|logit| logit.is_finite()));
        let argmax = executor
            .logits()
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap();
        eprintln!(
            "Rust logits: argmax={} value={} logit0={}",
            argmax.0,
            argmax.1,
            executor.logits()[0]
        );
        assert_eq!(argmax.0, 19_923);
        assert!((executor.logits()[0] - -7.675_424).abs() < 0.1);
        let next = argmax.0 as i32;
        let checkpoint =
            std::env::temp_dir().join(format!("ds4-rust-kv-{}.bin", std::process::id()));
        executor
            .save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
            .unwrap();
        executor.eval(next).unwrap();
        let continued_logit = executor.logits()[0];
        let continued_argmax = executor
            .logits()
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0;
        drop(executor);

        let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap();
        let mut restored = metal::Executor::open(model, 32_768, false, 0).unwrap();
        assert!(restored.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
        assert_eq!(restored.position(), tokens.len() as u32);
        assert_eq!(restored.tokens(), tokens);
        assert_eq!(restored.checkpoint_tag(), [7; 32]);
        restored.eval(next).unwrap();
        assert_eq!(
            restored
                .logits()
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.total_cmp(b.1))
                .unwrap()
                .0,
            continued_argmax
        );
        assert!((restored.logits()[0] - continued_logit).abs() < 1.0e-5);
        std::fs::remove_file(checkpoint).unwrap();
    }
}