aprender-serve 0.70.2

//! High-level inference API for CLI tools
//!
//! This module provides a simple, high-level API for running inference
//! that can be used by CLI tools like `apr run` and `apr chat`.
//!
//! # Architecture (APR-CLI-DELEGATE-001)
//!
//! ```text
//! ┌─────────────┐     ┌─────────────┐     ┌─────────────┐
//! │  apr-cli    │ --> │  realizar   │ --> │   trueno    │
//! │  (100 LOC)  │     │   infer.rs  │     │   SIMD/GPU  │
//! └─────────────┘     └─────────────┘     └─────────────┘
//! ```
//!
//! The `apr run` command delegates ALL inference to this module.
//! This eliminates ~1800 lines of duplicated code in apr-cli.
//!
//! # Example
//!
//! ```rust,ignore
//! use realizar::infer::{InferenceConfig, run_inference};
//!
//! let config = InferenceConfig::new("model.gguf")
//!     .with_prompt("Hello, world!")
//!     .with_max_tokens(32);
//!
//! let result = run_inference(config)?;
//! println!("{}", result.text);
//! ```

use crate::error::{RealizarError, Result};
use crate::format::{detect_format, ModelFormat};
use std::path::PathBuf;
use std::time::Instant;

/// PMAT-173 / GH-321: Convert GGML quantization type to human-readable string.
/// Uses unified `GgmlQuantType` enum — single source of truth.
pub(crate) fn qtype_to_dtype_str(qtype: u32) -> &'static str {
    crate::gguf::admitted_from_id(qtype).map_or("Unknown", crate::gguf::GgmlQuantType::as_str)
}

/// #4006: the `quant=` label for a GGUF model: the transformer BODY, not the head.
///
/// It printed `lm_head_weight.qtype`. Unsloth "UD" files tie the head to a
/// high-precision `token_embd` while the blocks are mixed, so
/// Qwen3.5-0.8B-UD-IQ2_XXS (95 IQ2_XXS block tensors) was labelled `Q5_K` and a
/// receipt quoting it said none of the IQ kernels ran.
///
/// `body` is the qtype of every 2-D projection weight in the blocks. One type
/// prints as that type; several print as `mixed(A×n,B×m,…)`, most frequent first
/// (ties by name). `lm_head=<qtype>` is appended when the head differs from the
/// dominant body type.
pub(crate) fn body_quant_label(body: &[u32], lm_head: u32) -> String {
    // The full GGML name table, not the admitted-kernel one: a label names what
    // the file holds, whether or not a GPU kernel exists for it.
    let name = |q: u32| {
        trueno_quant::GgmlType::from_id(q)
            .map_or_else(|| format!("ggml type {q}"), |t| t.as_str().to_string())
    };
    let mut counts: std::collections::BTreeMap<u32, usize> = std::collections::BTreeMap::new();
    for &q in body {
        *counts.entry(q).or_insert(0) += 1;
    }
    let mut ranked: Vec<(String, usize, u32)> =
        counts.into_iter().map(|(q, n)| (name(q), n, q)).collect();
    ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
    let Some((_, _, dominant)) = ranked.first().cloned() else {
        return name(lm_head);
    };
    let body_label = if ranked.len() == 1 {
        ranked[0].0.clone()
    } else {
        let parts: Vec<String> = ranked.iter().map(|(n, c, _)| format!("{n}×{c}")).collect();
        format!("mixed({})", parts.join(","))
    };
    if lm_head == dominant {
        body_label
    } else {
        format!("{body_label} lm_head={}", name(lm_head))
    }
}

/// The qtype of every 2-D block tensor (`blk.*`, `n_dims >= 2`) in the GGUF header,
/// for [`body_quant_label`]. Read from the FILE, not the loaded model struct: the
/// Qwen3.5 hybrid path builds only a base model (embeddings, final norm, head) with
/// no layers, and its block tensors are the ones that decide the label.
pub(crate) fn body_qtypes(gguf: &crate::gguf::GGUFModel) -> Vec<u32> {
    gguf.tensors
        .iter()
        .filter(|t| t.name.starts_with("blk.") && t.n_dims >= 2)
        .map(|t| t.qtype)
        .collect()
}

/// #4006 (APR paths): the qtype of every 2-D projection weight in a LOADED model's
/// layers, for [`body_quant_label`]. For `.apr` files, whose loader builds every
/// layer; GGUF uses [`body_qtypes`] on the header, because the qwen35 hybrid
/// builds no layers.
pub(crate) fn model_body_qtypes(model: &crate::gguf::OwnedQuantizedModel) -> Vec<u32> {
    use crate::gguf::OwnedQKVWeights;
    let mut out = Vec::new();
    for layer in model.layers() {
        match &layer.qkv_weight {
            OwnedQKVWeights::Fused(t) => out.push(t.qtype),
            OwnedQKVWeights::Separate { q, k, v } => out.extend([q.qtype, k.qtype, v.qtype]),
        }
        out.push(layer.attn_output_weight.qtype);
        out.push(layer.ffn_up_weight.qtype);
        out.push(layer.ffn_down_weight.qtype);
        if let Some(gate) = layer.ffn_gate_weight.as_ref() {
            out.push(gate.qtype);
        }
    }
    out
}

/// #4006 (SafeTensors): the GGML id for a float SafeTensors dtype (F32 0, F16 1,
/// BF16 30), so [`body_quant_label`] names SafeTensors weights the same way.
/// Integer dtypes are not weights and are skipped.
pub(crate) fn safetensors_dtype_ggml_id(
    dtype: &crate::safetensors::SafetensorsDtype,
) -> Option<u32> {
    use crate::safetensors::SafetensorsDtype as D;
    match dtype {
        D::F32 => Some(0),
        D::F16 => Some(1),
        D::BF16 => Some(30),
        _ => None,
    }
}

/// #4006: the `quant=` label for a SafeTensors file, read from its header: every
/// 2-D float tensor under `.layers.` is the body, `lm_head.weight` (else the tied
/// `embed_tokens`) is the head. `unknown (…)` when the header cannot be read, never
/// a guessed type.
pub(crate) fn safetensors_quant_label(path: &std::path::Path) -> String {
    let model = match crate::safetensors::MappedSafeTensorsModel::load(path) {
        Ok(m) => m,
        Err(e) => return format!("unknown (header unreadable: {e})"),
    };
    let mut body = Vec::new();
    let mut head = None;
    for name in model.tensor_names() {
        let Some(info) = model.get_tensor_info(name) else {
            continue;
        };
        let Some(id) = safetensors_dtype_ggml_id(&info.dtype) else {
            continue;
        };
        if name.contains(".layers.") && info.shape.len() >= 2 {
            body.push(id);
        } else if name == "lm_head.weight"
            || (head.is_none() && name.ends_with("embed_tokens.weight"))
        {
            head = Some(id);
        }
    }
    match head {
        Some(h) => body_quant_label(&body, h),
        None if body.is_empty() => "unknown (no float weights in the header)".to_string(),
        // No head tensor: label the body alone (the head clause only appears on a mismatch).
        None => {
            let dominant = body_quant_label(&body, u32::MAX);
            dominant
                .split(" lm_head=")
                .next()
                .unwrap_or(&dominant)
                .to_string()
        },
    }
}

/// Configuration for inference
#[derive(Debug, Clone)]
pub struct InferenceConfig {
    /// Path to model file (GGUF, APR, or SafeTensors)
    pub model_path: PathBuf,
    /// Text prompt for generation
    pub prompt: Option<String>,
    /// Token IDs for generation (alternative to prompt)
    pub input_tokens: Option<Vec<u32>>,
    /// Maximum tokens to generate
    pub max_tokens: usize,
    /// Temperature for sampling (0.0 = greedy)
    pub temperature: f32,
    /// Top-k sampling (0 = disabled)
    pub top_k: usize,
    /// Top-p (nucleus) sampling threshold (None = disabled, i.e. 1.0)
    /// PMAT-823: previously dropped — `apr run --top-p` was a no-op.
    pub top_p: Option<f32>,
    /// RNG seed for stochastic sampling
    /// PMAT-823: previously dropped — `apr run --seed` was a no-op.
    pub seed: u64,
    /// Repetition penalty (1.0 = no penalty)
    /// PMAT-823: previously dropped — `apr run --repeat-penalty` was a no-op.
    pub repeat_penalty: f32,
    /// Context window for repetition penalty
    /// PMAT-823: previously dropped — `apr run --repeat-last-n` was a no-op.
    pub repeat_last_n: usize,
    /// Disable GPU acceleration
    pub no_gpu: bool,
    /// #3757: the user EXPLICITLY asked for an accelerator (`--gpu`, or
    /// `--backend cuda|wgpu|gpu`), as classified by
    /// `apr-cli::registry::Request::wanted` — the same signal `reconcile_accelerator`
    /// already consumes, threaded one step further so the ATTEMPT is gated by it
    /// and not only the post-hoc verdict.
    ///
    /// Without this the bare `apr run model.gguf` enters the GH-559 wgpu
    /// fallback, dequantizes the whole model to F32, fails wgpu's own cpu-parity
    /// gate and falls back — paying 1.7 GB and 2.5x the wall time to reach the
    /// identical CPU answer. Measured on `release/0.69.1-batch-2` @ 9f8836c71,
    /// qwen2.5-coder-1.5b-q4_k_m: default 7607 ms vs `--no-gpu` 3035 ms.
    pub accel_forced: bool,
    /// Enable inference tracing (APR-TRACE-001)
    pub trace: bool,
    /// Verbose tracing output
    pub trace_verbose: bool,
    /// Trace output file path
    pub trace_output: Option<PathBuf>,
    /// Specific trace steps to capture
    pub trace_steps: Option<Vec<String>>,
    /// Show verbose loading/progress output
    pub verbose: bool,
    /// Stop token IDs for early termination (GH-373)
    pub stop_tokens: Vec<u32>,
    /// INTERNAL: Use mock backend for testing (PMAT-COV-95)
    #[doc(hidden)]
    pub use_mock_backend: bool,
    /// #3672: apply the model's chat template even when neither its metadata nor its file
    /// name marks it as an instruct model (`apr run --chat`). The prompt stays raw text:
    /// the template is applied once, by `prepare_tokens`, never pre-wrapped by a caller.
    pub force_chat_template: bool,
    /// #3723: `--thinking on|off`. `None` renders what production always has; `Some(true)`
    /// removes the empty `<think>` prefill so the model reasons, and is refused by name on a
    /// template with no thinking mode ([`crate::chat_template::apply_thinking_mode`]).
    pub thinking: Option<bool>,
}

/// The top-k a SAMPLED generation uses when the caller names none (#3754).
///
/// `apr run` defaulted `--top-k` to 1, and every decode loop treats `top_k == 1` as greedy,
/// so `apr run --temperature 0.8` decoded greedily and said nothing: a sampling flag that did
/// nothing. `apr chat` and `apr serve` already sampled with 40, which is also the llama.cpp
/// and Ollama default. This is that number, declared once. Greedy is still
/// `temperature == 0.0` or an explicit `top_k == 1`; `0` disables the filter.
pub const DEFAULT_TOP_K: usize = 40;

/// The top-k a generation runs with: `1` (greedy) at temperature 0, else the caller's value,
/// else [`DEFAULT_TOP_K`].
#[must_use]
pub fn sampling_top_k(temperature: f32, requested: Option<usize>) -> usize {
    if temperature == 0.0 {
        1
    } else {
        requested.unwrap_or(DEFAULT_TOP_K)
    }
}

impl InferenceConfig {
    /// Create a new inference config for a model file
    #[must_use]
    pub fn new(model_path: impl Into<PathBuf>) -> Self {
        Self {
            model_path: model_path.into(),
            prompt: None,
            input_tokens: None,
            max_tokens: 32,
            temperature: 0.0, // Greedy by default
            // PMAT-823: a default config forwards the byte-identical greedy generation
            // config, top_k 1 included. A caller that samples names its top-k:
            // `apr run` sends `--top-k` (default DEFAULT_TOP_K), others use `sampling_top_k`.
            top_k: 1,
            // PMAT-823: defaults chosen so a config with no sampling flags
            // forwards to the SAME greedy QuantizedGenerateConfig as before
            // (top_p 1.0 / seed 42 / repeat_penalty 1.0 / repeat_last_n 64).
            top_p: None,
            seed: 42,
            repeat_penalty: 1.0,
            repeat_last_n: 64,
            no_gpu: false,
            accel_forced: false,
            trace: false,
            trace_verbose: false,
            trace_output: None,
            trace_steps: None,
            verbose: false,
            stop_tokens: Vec::new(),
            use_mock_backend: false,
            force_chat_template: false,
            thinking: None,
        }
    }

    /// Set the text prompt
    #[must_use]
    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.prompt = Some(prompt.into());
        self
    }

    /// Set input tokens directly
    #[must_use]
    pub fn with_input_tokens(mut self, tokens: Vec<u32>) -> Self {
        self.input_tokens = Some(tokens);
        self
    }

    /// Set maximum tokens to generate
    #[must_use]
    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
        self.max_tokens = max_tokens;
        self
    }

    /// Set temperature (0.0 = greedy)
    #[must_use]
    pub fn with_temperature(mut self, temperature: f32) -> Self {
        self.temperature = temperature;
        self
    }

    /// Set top-k sampling
    #[must_use]
    pub fn with_top_k(mut self, top_k: usize) -> Self {
        self.top_k = top_k;
        self
    }

    /// Set top-p (nucleus) sampling threshold (PMAT-823).
    #[must_use]
    pub fn with_top_p(mut self, top_p: Option<f32>) -> Self {
        self.top_p = top_p;
        self
    }

    /// Set the RNG seed for stochastic sampling (PMAT-823).
    #[must_use]
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = seed;
        self
    }

    /// Set the repetition penalty (1.0 = no penalty) (PMAT-823).
    #[must_use]
    pub fn with_repeat_penalty(mut self, repeat_penalty: f32) -> Self {
        self.repeat_penalty = repeat_penalty;
        self
    }

    /// Set the repetition-penalty context window (PMAT-823).
    #[must_use]
    pub fn with_repeat_last_n(mut self, repeat_last_n: usize) -> Self {
        self.repeat_last_n = repeat_last_n;
        self
    }

    /// Disable GPU acceleration
    #[must_use]
    pub fn without_gpu(mut self) -> Self {
        self.no_gpu = true;
        self
    }

    /// #3757: record that the user explicitly asked for an accelerator.
    #[must_use]
    pub fn with_accel_forced(mut self, accel_forced: bool) -> Self {
        self.accel_forced = accel_forced;
        self
    }

    /// Enable verbose output
    #[must_use]
    pub fn with_verbose(mut self, verbose: bool) -> Self {
        self.verbose = verbose;
        self
    }

    /// #3672: apply the chat template even when metadata and file name say base model.
    #[must_use]
    pub fn with_force_chat_template(mut self, force: bool) -> Self {
        self.force_chat_template = force;
        self
    }

    /// #3723: the thinking mode `--thinking on|off` asks for (`None`: the production default).
    #[must_use]
    pub fn with_thinking(mut self, thinking: Option<bool>) -> Self {
        self.thinking = thinking;
        self
    }

    /// Enable inference tracing
    #[must_use]
    pub fn with_trace(mut self, trace: bool) -> Self {
        self.trace = trace;
        self
    }

    /// Set trace output file path
    #[must_use]
    pub fn with_trace_output(mut self, path: impl Into<PathBuf>) -> Self {
        self.trace_output = Some(path.into());
        self
    }

    /// Set stop token IDs for early termination (GH-373)
    #[must_use]
    pub fn with_stop_tokens(mut self, stop_tokens: Vec<u32>) -> Self {
        self.stop_tokens = stop_tokens;
        self
    }

    /// PMAT-823: Copy ALL sampling parameters from this `InferenceConfig` into a
    /// `QuantizedGenerateConfig`, so CLI flags (`--temperature`, `--top-k`,
    /// `--top-p`, `--seed`, `--repeat-penalty`, `--repeat-last-n`) actually reach
    /// the decode loop instead of being silently dropped to greedy defaults.
    ///
    /// `top_p: None` maps to the disabled threshold `1.0` (matching
    /// `QuantizedGenerateConfig::default().top_p`), so a default config produces
    /// a byte-identical greedy `gen_config`.
    pub(crate) fn apply_sampling_to(&self, gen_config: &mut crate::gguf::QuantizedGenerateConfig) {
        gen_config.temperature = self.temperature;
        gen_config.top_k = self.top_k;
        gen_config.top_p = self.top_p.unwrap_or(1.0);
        gen_config.seed = self.seed;
        gen_config.repeat_penalty = self.repeat_penalty;
        gen_config.repeat_last_n = self.repeat_last_n;
    }
}

// ============================================================================
// PreparedTokens - Compile-time chat template enforcement (PMAT-236)
// ============================================================================

/// Tokenized input that has been processed through chat template formatting.
///
/// # Compile-time enforcement (Poka-Yoke)
///
/// The inner `Vec<u32>` is **private** - the only way to construct `PreparedTokens`
/// is via `prepare_tokens()`, which ALWAYS applies chat template formatting for
/// instruct models. This makes it a **compile error** to pass raw tokens to
/// inference functions, preventing the bug where SafeTensors inference skipped
/// chat template application (producing "4" then garbage).
///
/// # Theoretical basis
///
/// Shingo, S. (1986). *Zero Quality Control: Source Inspection and the Poka-Yoke System*.
/// Brady, E. (2017). *Type-Driven Development with Idris*.
///
/// # References
///
/// - PMAT-236: Chat template enforcement for multi-format inference
/// - GH-205: SafeTensors inference garbage root cause
#[derive(Debug, Clone)]
pub struct PreparedTokens {
    /// Tokenized input (PRIVATE - enforces construction via prepare_tokens only)
    tokens: Vec<u32>,
    /// Number of input tokens (for separating prefill from generated tokens)
    input_count: usize,
}

impl PreparedTokens {
    /// Access the prepared token IDs (read-only).
    #[must_use]
    pub fn tokens(&self) -> &[u32] {
        &self.tokens
    }

    /// Number of input tokens.
    #[must_use]
    pub fn input_count(&self) -> usize {
        self.input_count
    }
}

/// Prepare tokens for inference, applying chat template for instruct models.
///
/// This is the ONLY way to create `PreparedTokens`. It handles:
/// 1. Format detection (GGUF vs SafeTensors vs APR)
/// 2. Architecture detection (Qwen2, LLaMA, Phi, etc.)
/// 3. Chat template application for instruct models
/// 4. Tokenization using the appropriate tokenizer
///
/// # Chat Template Rules
///
/// - If model name/architecture contains "instruct", chat template is applied
/// - GGUF: uses embedded tokenizer + architecture from metadata
/// - SafeTensors: uses sibling tokenizer.json + config.json architecture
/// - APR: uses sibling tokenizer.json + model metadata
///
/// # Errors
///
/// Returns error if the model cannot be read or tokenization fails.
pub fn prepare_tokens(config: &InferenceConfig, format: &ModelFormat) -> Result<PreparedTokens> {
    // If raw token IDs are provided, use them directly (user knows what they're doing)
    if let Some(ref tokens) = config.input_tokens {
        return Ok(PreparedTokens {
            input_count: tokens.len(),
            tokens: tokens.clone(),
        });
    }

    let prompt = match config.prompt {
        Some(ref p) => p.clone(),
        None => {
            return Ok(PreparedTokens {
                tokens: vec![1u32],
                input_count: 1,
            })
        },
    };

    match format {
        ModelFormat::Gguf => prepare_tokens_gguf(config, &prompt),
        ModelFormat::SafeTensors => prepare_tokens_safetensors(config, &prompt),
        ModelFormat::Apr => prepare_tokens_apr(config, &prompt),
    }
}

/// Prepare tokens for GGUF format (chat template from GGUF metadata)
///
/// GH-278: Only apply chat template when the GGUF actually contains one in its
/// metadata (`tokenizer.chat_template`). Previously, ALL models with known
/// architectures (llama, qwen2, etc.) got chat-template wrapping even if they
/// were base completion models, causing complete output divergence vs llama.cpp.
///
/// BOS token: Prepend BOS when the model metadata says `add_bos_token = true`
/// or when a BOS token ID exists and `add_bos_token` is not explicitly false.
/// This matches llama.cpp behavior for LLaMA-family models.
/// #3723: apply `--thinking` to a rendered prompt. A prompt that no chat template rendered has
/// no thinking mode, so ON is refused there too; OFF and `None` leave every prompt unchanged.
fn thinking_mode(config: &InferenceConfig, formatted: String) -> Result<String> {
    if config.thinking.is_none() {
        return Ok(formatted);
    }
    crate::chat_template::apply_thinking_mode(&formatted, config.thinking)
}

/// #3990: the `tokenizer_config.json` beside a SafeTensors model, when it declares a chat template
/// (a string, or the list form with a `default` entry, as transformers reads it). Rendering, and
/// bos/eos, are left to the ONE reader, `chat_template::render_official_from_tokenizer_config`;
/// this only decides whether the model has a template of its own, so a model WITHOUT one takes
/// the built-in formatter quietly while a template that fails to render is warned about.
fn sibling_tokenizer_config(model_path: &std::path::Path) -> Option<String> {
    let text = std::fs::read_to_string(model_path.with_file_name("tokenizer_config.json")).ok()?;
    let v: serde_json::Value = serde_json::from_str(&text).ok()?;
    let declared = match v.get("chat_template")? {
        serde_json::Value::String(s) => !s.is_empty(),
        serde_json::Value::Array(list) => list
            .iter()
            .any(|t| t.get("name").and_then(serde_json::Value::as_str) == Some("default")),
        _ => false,
    };
    declared.then_some(text)
}

fn prepare_tokens_gguf(config: &InferenceConfig, prompt: &str) -> Result<PreparedTokens> {
    use crate::chat_template::{format_messages, ChatMessage};
    use crate::gguf::{GGUFValue, MappedGGUFModel};

    let mapped = MappedGGUFModel::from_path(&config.model_path)?;
    let gguf_arch = mapped.model.architecture().unwrap_or("transformer");

    // GH-278: Check if model actually has a chat template in its GGUF metadata.
    // Base models (SmolLM-135M, GPT-2) don't have one — only instruct/chat models do.
    let has_chat_template = mapped
        .model
        .metadata
        .get("tokenizer.chat_template")
        .is_some_and(|v| matches!(v, GGUFValue::String(s) if !s.is_empty()));

    let model_name = config
        .model_path
        .file_name()
        .and_then(|n| n.to_str())
        .unwrap_or("");
    let filename_instruct = model_name.to_lowercase().contains("instruct")
        || model_name.to_lowercase().contains("-chat");

    // Only apply chat template if the model actually has one, or filename says instruct
    let formatted_prompt = if config.force_chat_template || has_chat_template || filename_instruct {
        let template_hint = apr_arch_to_template_hint(gguf_arch, model_name);
        let messages = vec![ChatMessage::user(prompt)];
        // #3990: the GGUF's own tokenizer.chat_template, when it carries one.
        let own = has_chat_template.then_some(|t: Option<bool>| {
            crate::chat_template::render_official_for_model(&mapped.model, &messages, t)
        });
        crate::chat_template::official_or_legacy(
            own,
            || {
                format_messages(&messages, Some(template_hint))
                    .unwrap_or_else(|_| prompt.to_string())
            },
            config.thinking,
        )?
    } else {
        thinking_mode(config, prompt.to_string())?
    };

    if config.verbose {
        eprintln!(
            "[DEBUG] has_chat_template={}, filename_instruct={}",
            has_chat_template, filename_instruct
        );
        eprintln!(
            "[DEBUG] formatted_prompt={:?}",
            log_head(&formatted_prompt, 200)
        );
    }

    let mut tokens = mapped.model.encode(&formatted_prompt).ok_or_else(|| {
        RealizarError::InferenceError(format!(
            "Tokenizer encode failed for GGUF model (no tokenizer data in GGUF file?). \
                 Prompt length: {} chars",
            formatted_prompt.len()
        ))
    })?;

    // GH-278: Prepend BOS token to match llama.cpp behavior.
    // llama.cpp adds BOS when add_bos_token is true (default for LLaMA-family).
    // Only add if not already present AND model has a BOS token defined.
    let add_bos = match mapped
        .model
        .metadata
        .get(crate::gguf::keys::TOKENIZER_ADD_BOS)
    {
        Some(GGUFValue::Bool(b)) => *b,
        // GH-326: Derive BOS default from architecture constraints, not hardcoded string.
        // Models with absolute position embeddings (GPT-2, BERT) use BPE → no BOS.
        // Models with RoPE (LLaMA, Qwen, Mistral) use SentencePiece → add BOS.
        _ => {
            let arch = mapped
                .model
                .metadata
                .get(crate::gguf::keys::GENERAL_ARCHITECTURE)
                .and_then(|v| {
                    if let GGUFValue::String(s) = v {
                        Some(s.as_str())
                    } else {
                        None
                    }
                })
                // R-01 (Meyer DbC): "unknown" — don't pretend unidentified model is LLaMA.
                .unwrap_or("unknown");
            let constraints = crate::gguf::ArchConstraints::from_architecture(arch);
            constraints.positional_encoding != crate::gguf::PositionalEncoding::Absolute
        },
    };

    if add_bos {
        if let Some(bos_id) = mapped.model.bos_token_id() {
            if tokens.first() != Some(&bos_id) {
                tokens.insert(0, bos_id);
            }
        }
    }

    if config.verbose {
        eprintln!(
            "[DEBUG] add_bos={}, encoded {} tokens: {:?}",
            add_bos,
            tokens.len(),
            &tokens[..tokens.len().min(30)]
        );
    }

    Ok(PreparedTokens {
        input_count: tokens.len(),
        tokens,
    })
}

/// Prepare tokens for SafeTensors format (chat template from config.json)
fn prepare_tokens_safetensors(config: &InferenceConfig, prompt: &str) -> Result<PreparedTokens> {
    use crate::apr::AprV2Model;
    use crate::chat_template::{format_messages, ChatMessage};
    use crate::safetensors::SafetensorsConfig;

    // Load config.json for architecture detection
    let st_config = SafetensorsConfig::load_from_sibling(&config.model_path);
    let architecture = st_config
        .as_ref()
        .map(SafetensorsConfig::architecture)
        .unwrap_or_default();

    let model_name = config
        .model_path
        .file_name()
        .and_then(|n| n.to_str())
        .unwrap_or("");

    // Detect instruct model from architecture or filename
    let arch_lower = architecture.to_lowercase();
    let is_instruct = config.force_chat_template
        || arch_lower.contains("instruct")
        || model_name.to_lowercase().contains("instruct")
        || matches!(
            arch_lower.as_str(),
            "qwen2forcausallm" | "llamaforcausallm" | "mistralforcausallm" | "phiforcausallm"
        );

    let formatted_prompt = if is_instruct {
        let template_hint = safetensors_arch_to_template_hint(&architecture, model_name);
        let messages = vec![ChatMessage::user(prompt)];
        // #3990: the sibling tokenizer_config.json's own chat_template, when it declares one.
        let tc = sibling_tokenizer_config(&config.model_path);
        let msgs = &messages;
        let own = tc.as_deref().map(|json| {
            move |t: Option<bool>| {
                crate::chat_template::render_official_from_tokenizer_config(json, msgs, t)
            }
        });
        crate::chat_template::official_or_legacy(
            own,
            || {
                format_messages(&messages, Some(template_hint))
                    .unwrap_or_else(|_| prompt.to_string())
            },
            config.thinking,
        )?
    } else {
        thinking_mode(config, prompt.to_string())?
    };

    let tokens =
        AprV2Model::encode_text(&config.model_path, &formatted_prompt).ok_or_else(|| {
            RealizarError::InferenceError(format!(
                "Tokenizer encode failed for SafeTensors model (no tokenizer.json sibling?). \
                 Prompt length: {} chars",
                formatted_prompt.len()
            ))
        })?;

    Ok(PreparedTokens {
        input_count: tokens.len(),
        tokens,
    })
}

/// Prepare tokens for APR format (chat template from model metadata).
///
/// GH-1623: Only apply chat template when the APR actually contains one in
/// its metadata (`tokenizer.chat_template`). Mirrors GGUF path (GH-278).
///
/// Previously, ALL models with known architectures (qwen2 / llama / mistral /
/// phi) OR with ChatML special tokens in vocab got chat-template wrapping —
/// even base completion models like `qwen2.5-coder-0.5b` (base) which carry
/// the Qwen tokenizer with `<|im_start|>` in vocab but are NOT instruct. The
/// over-triggering produced garbage output for base models.
///
/// New detection rule (matches GGUF):
///   has_chat_template_in_metadata || filename hint (`instruct` / `-chat`)
///
/// PMAT-237's note about hash-named APR files remains valid — if a file is
/// hash-named AND lacks a `tokenizer.chat_template` in its metadata, we now
/// correctly treat it as a base model (the previous broad heuristic was
/// silently wrapping such files even when they were base completion models).
fn prepare_tokens_apr(config: &InferenceConfig, prompt: &str) -> Result<PreparedTokens> {
    use crate::apr::AprV2Model;
    use crate::chat_template::{format_messages, ChatMessage};

    let model_name = config
        .model_path
        .file_name()
        .and_then(|n| n.to_str())
        .unwrap_or("");

    let (apr_arch, own_template) = if config.model_path.extension().is_some_and(|e| e == "apr") {
        match AprV2Model::load(&config.model_path) {
            Ok(model) => {
                let meta = model.metadata();
                let arch = meta.architecture.clone().unwrap_or_default();
                let tmpl = meta
                    .extra
                    .get("tokenizer.chat_template")
                    .and_then(|v| v.as_str())
                    .filter(|s| !s.is_empty())
                    .map(str::to_string);
                (arch, tmpl)
            },
            Err(_) => (String::new(), None),
        }
    } else {
        (String::new(), None)
    };
    let has_chat_template = own_template.is_some();

    let filename_instruct = model_name.to_lowercase().contains("instruct")
        || model_name.to_lowercase().contains("-chat");

    let is_instruct = config.force_chat_template || has_chat_template || filename_instruct;

    let formatted_prompt = if is_instruct {
        let template_hint = apr_arch_to_template_hint(&apr_arch, model_name);
        let messages = vec![ChatMessage::user(prompt)];
        // #3990: the .apr's own tokenizer.chat_template. The .apr carries no bos/eos STRINGS,
        // so they stay undefined (a Qwen template references neither).
        let msgs = &messages;
        let own = own_template.as_deref().map(|tpl| {
            move |t: Option<bool>| {
                crate::chat_template::render_official(tpl, None, None, msgs, true, t)
            }
        });
        crate::chat_template::official_or_legacy(
            own,
            || {
                format_messages(&messages, Some(template_hint))
                    .unwrap_or_else(|_| prompt.to_string())
            },
            config.thinking,
        )?
    } else {
        thinking_mode(config, prompt.to_string())?
    };

    let tokens =
        AprV2Model::encode_text(&config.model_path, &formatted_prompt).ok_or_else(|| {
            RealizarError::InferenceError(format!(
                "Tokenizer encode failed for APR model (no tokenizer in APR metadata?). \
                 Prompt length: {} chars",
                formatted_prompt.len()
            ))
        })?;

    Ok(PreparedTokens {
        input_count: tokens.len(),
        tokens,
    })
}

/// Map SafeTensors architecture string to chat template hint.
///
/// GH-317/318: Contract-driven — uses normalize_architecture() from
/// tensor-names-v1.yaml. No contains() heuristics, no model_name fallback.
fn safetensors_arch_to_template_hint(architecture: &str, _model_name: &str) -> &'static str {
    crate::tensor_names::normalize_architecture(architecture)
}

/// #4018: at most the first `max` bytes of `s`, cut at a CHAR BOUNDARY, for a log or error line.
///
/// `&s[..s.len().min(max)]` panicked ("byte index N is not a char boundary") whenever byte `max`
/// fell inside a multi-byte UTF-8 char, so `apr run -v` crashed on a non-ASCII prompt instead of
/// answering. The cut floors to the previous boundary: at most 3 bytes short, since a char is at
/// most 4 (the CRUX judge reads a logged prompt of >= max-3 bytes as possibly cut, #3962 B2).
/// `str::floor_char_boundary` would do this, but is not stable at this crate's rust-version.
pub(crate) fn log_head(s: &str, max: usize) -> &str {
    let mut end = s.len().min(max);
    while !s.is_char_boundary(end) {
        end -= 1;
    }
    &s[..end]
}

#[cfg(test)]
mod log_head_4018 {
    use super::log_head;

    /// MUST-RED (#4018): `apr run -v` logs the formatted prompt's head, and a prompt whose byte 200
    /// falls INSIDE a multi-byte char panicked ("byte index 200 is not a char boundary"). ChatML around
    /// `x` + 80 CJK chars puts byte 200 mid-char.
    #[test]
    fn a_non_ascii_prompt_cut_mid_char_does_not_panic() {
        let p = format!(
            "<|im_start|>user\nx{}<|im_end|>\n<|im_start|>assistant\n",
            "\u{6c34}".repeat(80)
        );
        assert!(
            !p.is_char_boundary(200),
            "the fixture must put byte 200 mid-char"
        );
        let head = log_head(&p, 200);
        assert!(head.len() <= 200 && head.len() >= 197, "{}", head.len());
        assert!(p.starts_with(head));
    }

    #[test]
    fn ascii_and_short_inputs_are_unchanged() {
        assert_eq!(log_head("What is 2+2?", 200), "What is 2+2?");
        let a = "a".repeat(300);
        assert_eq!(log_head(&a, 200).len(), 200);
        assert_eq!(log_head("", 200), "");
    }

    /// #4018: the SITES use it -- the helper alone proves nothing if a caller still byte-slices.
    #[test]
    fn no_log_site_byte_slices_text_any_more() {
        for (f, old) in [
            (
                "src/infer/mod.rs",
                "&formatted_prompt[..formatted_prompt.len().min(200)]",
            ),
            (
                "src/infer/inference_result.rs",
                "&raw_text[..raw_text.len().min(200)]",
            ),
        ] {
            let src = std::fs::read_to_string(format!("{}/{f}", env!("CARGO_MANIFEST_DIR")))
                .expect("source");
            // Scan the code ABOVE this test module: the module itself names the old pattern, and a
            // source guard that reads its own assertion strings is satisfied by them (#3907's lesson).
            let code = src.split("mod log_head_4018").next().unwrap_or(&src);
            assert!(
                !code.contains(old),
                "{f} still slices text at a fixed byte length: {old}"
            );
        }
    }
}

include!("inference_result.rs");
include!("gguf_gpu_generate.rs");
include!("mod_log_transformer_eos.rs");
include!("mod_05.rs");
include!("batch.rs");

/// #3714: qwen3moe backend selection — the CUDA forward, or the CPU chain with a printed reason.
pub mod qwen3_moe_dispatch;
pub mod qwen3_moe_generate;
pub mod run_report;

// #3760: `apr run` on a SafeTensors model samples.
#[cfg(test)]
#[path = "tests_quant_label_4006.rs"]
mod tests_quant_label_4006;
#[cfg(test)]
#[path = "tests_sampling_3760.rs"]
mod tests_sampling_3760;
// #3754: a sampling flag given alone samples (DEFAULT_TOP_K).
#[cfg(test)]
#[path = "tests_sampling_default_3754.rs"]
mod tests_sampling_default_3754;
// #4268: dense `run` and `run --batch` go through the one engine.
#[cfg(test)]
#[path = "tests_dense_session_4268.rs"]
mod tests_dense_session_4268;

#[cfg(test)]
mod sibling_tokenizer_config_3990 {
    use super::sibling_tokenizer_config;

    fn with_config(json: Option<&str>) -> Option<String> {
        let dir = tempfile::tempdir().expect("tempdir");
        if let Some(j) = json {
            std::fs::write(dir.path().join("tokenizer_config.json"), j).expect("write");
        }
        sibling_tokenizer_config(&dir.path().join("model.safetensors"))
    }

    /// A declared template -- a string, or the list form with a `default` entry -- is the model's own;
    /// no file, no key, an empty string or a list without `default` is not, and takes the built-in
    /// formatter quietly (#3990: rendering and bos/eos belong to render_official_from_tokenizer_config).
    #[test]
    fn a_declared_template_is_found_in_both_forms_and_nothing_else_is() {
        assert!(with_config(Some(r#"{"chat_template": "{{ messages }}"}"#)).is_some());
        assert!(with_config(Some(
            r#"{"chat_template": [{"name": "default", "template": "x"}]}"#
        ))
        .is_some());
        for j in [
            None,
            Some("{}"),
            Some(r#"{"chat_template": ""}"#),
            Some(r#"{"chat_template": [{"name": "tool_use", "template": "x"}]}"#),
            Some("not json"),
        ] {
            assert!(with_config(j).is_none(), "{j:?}");
        }
    }
}