semlith 0.35.0

//! Which embedding model a store uses, and how to load it.
//!
//! fastembed ships an enum of models it knows how to fetch and configure. The
//! default is not in that enum: granite-embedding-small-english-r2, quantized
//! to int8, has to be assembled by hand from its ONNX graph, its separate
//! weights file, and its tokenizer. [`Model`] is the union of "something
//! fastembed knows" and "the one we assemble ourselves", so a store can record
//! either and read it back.

use anyhow::{Context, Result, bail};
use fastembed::{
    EmbeddingModel, InitOptionsUserDefined, Pooling, TextEmbedding, TextInitOptions,
    TokenizerFiles, UserDefinedEmbeddingModel,
};
use std::path::{Path, PathBuf};
use std::str::FromStr;

/// The Hugging Face repository holding granite's ONNX export.
pub const GRANITE_REPO: &str = "onnx-community/granite-embedding-small-english-r2-ONNX";

/// int8 rather than fp32: a quarter of the download, and measurably faster on
/// ARM despite the folklore that quantized ONNX never is.
const GRANITE_ONNX: &str = "onnx/model_quantized.onnx";

/// The graph refers to its weights by this exact name. ONNX Runtime matches on
/// the string, so it has to survive the round trip through the loader verbatim.
const GRANITE_WEIGHTS: &str = "onnx/model_quantized.onnx_data";

/// What a store writes into its `model` meta row for granite. Changing this
/// string orphans every store that recorded the old one.
pub const GRANITE_NAME: &str = "granite-embedding-small-english-r2-int8";

const GRANITE_DIM: usize = 384;

/// Cap on ONNX Runtime's intra-op thread pool, overriding the derived value.
pub const THREADS_ENV: &str = "SEMLITH_EMBED_THREADS";

/// A model a store can be built with.
#[derive(Debug, Clone, PartialEq, Default)]
pub enum Model {
    /// One of fastembed's own; it handles fetching and configuration.
    Builtin(EmbeddingModel),
    /// granite-embedding-small-english-r2, int8. Assembled here because
    /// fastembed has no entry for it. This is the default for new stores.
    #[default]
    Granite,
}

impl std::fmt::Display for Model {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Model::Granite => f.write_str(GRANITE_NAME),
            Model::Builtin(m) => write!(f, "{m}"),
        }
    }
}

impl FromStr for Model {
    type Err = String;

    /// Stores written by 0.1.0 hold a fastembed enum name such as
    /// `BGESmallENV15`, so that spelling has to keep parsing exactly as before.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == GRANITE_NAME {
            return Ok(Model::Granite);
        }
        EmbeddingModel::from_str(s).map(Model::Builtin)
    }
}

impl Model {
    pub fn dim(&self) -> Result<usize> {
        match self {
            Model::Granite => Ok(GRANITE_DIM),
            Model::Builtin(m) => Ok(TextEmbedding::get_model_info(m)
                .map_err(|e| anyhow::anyhow!("{e}"))?
                .dim),
        }
    }

    /// BGE English models were trained with an asymmetric query instruction;
    /// omitting it measurably costs recall. granite was trained without one,
    /// and adding a prefix it never saw would only add noise.
    pub fn query_text(&self, query: &str) -> String {
        match self {
            Model::Builtin(m) => {
                let name = m.to_string();
                if name.starts_with("BGE") && name.contains("EN") {
                    format!("Represent this sentence for searching relevant passages: {query}")
                } else {
                    query.to_string()
                }
            }
            Model::Granite => query.to_string(),
        }
    }

    /// The tokenizer this model embeds with, for counting tokens rather than
    /// estimating them.
    ///
    /// The same `tokenizer.json` [`Self::load`] hands to fastembed, read from
    /// the same cache and verified against the same pinned digest — so a token
    /// count here and the model's own segmentation are the same arithmetic,
    /// not two approximations of it.
    ///
    /// `None` when the file is not cached, which is the airgapped machine that
    /// has never fetched a model. The caller falls back to four characters per
    /// token and says so in the row it writes.
    pub fn tokenizer(&self, cache_dir: &Path) -> Option<tokenizers::Tokenizer> {
        let name = match self {
            Model::Granite => GRANITE_REPO,
            // fastembed lays a builtin model's files out under its own
            // directory name, which is the model's `model_code`.
            Model::Builtin(m) => return builtin_tokenizer(cache_dir, m),
        };
        let dir = snapshot_dir(cache_dir, name, GRANITE_REVISION)?;
        tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).ok()
    }

    pub fn load(
        &self,
        cache_dir: PathBuf,
        max_length: usize,
        quiet: bool,
    ) -> Result<TextEmbedding> {
        self.load_with_threads(cache_dir, max_length, quiet, embed_threads())
    }

    /// [`Self::load`] with the intra-op thread count named, for a writer that
    /// builds its session with the `threads each` value in force rather than
    /// with the derived one.
    pub fn load_with_threads(
        &self,
        cache_dir: PathBuf,
        max_length: usize,
        quiet: bool,
        threads: usize,
    ) -> Result<TextEmbedding> {
        self.load_variant(cache_dir, max_length, quiet, threads, Variant::Int8)
    }

    /// [`Self::load_with_threads`] for one of granite's variants. A builtin
    /// model has one and ignores the argument.
    pub fn load_variant(
        &self,
        cache_dir: PathBuf,
        max_length: usize,
        quiet: bool,
        threads: usize,
        variant: Variant,
    ) -> Result<TextEmbedding> {
        let threads = threads.max(1);
        // Checked here, in the one place weights are ever fetched, rather than
        // at each call site: an airgapped machine's whole claim is that this
        // process cannot have been the one that reached the network, and a
        // check that lives anywhere else is a check a later caller can miss.
        if airgap() && !is_cached(&cache_dir) {
            bail!(
                "{AIRGAP_ENV} is set and no model is cached at {} — \
                 pre-seed it with SEMLITH_MODEL_CACHE on a connected machine, \
                 or drop --airgap to let this run download it",
                cache_dir.display()
            );
        }
        match self {
            Model::Builtin(m) => {
                link_runtime()?;
                let opts = TextInitOptions::new(m.clone())
                    .with_show_download_progress(!quiet)
                    .with_max_length(max_length)
                    .with_intra_threads(threads)
                    .with_cache_dir(cache_dir);
                TextEmbedding::try_new(opts).map_err(|e| anyhow::anyhow!("{e}"))
            }
            Model::Granite => load_granite(cache_dir, max_length, quiet, threads, variant),
        }
    }
}

/// Which of granite's pinned exports a CPU session loads.
///
/// int8 is the default and the only one a store is built with unless asked:
/// a quarter of fp32's download and faster on ARM. fp16 is what the GPU lanes
/// run, and fp32 is what the known-answer fixture was made with. The CPU can
/// load all three, which is how the retrieval harness measures a store holding
/// a mix of them (item 1.16) on a machine with no GPU in the loop.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Variant {
    Int8,
    Fp16,
    Fp32,
}

impl Variant {
    /// The graph and its weights, each with the SHA-256 at [`GRANITE_REVISION`].
    fn files(self) -> [(&'static str, &'static str); 2] {
        match self {
            Variant::Int8 => [GRANITE_FILES[4], GRANITE_FILES[5]],
            Variant::Fp16 => [
                (
                    "onnx/model_fp16.onnx",
                    "ee200de55cb2f94e858aabca54be7697a9c0805a14c858ee26ad0922b05f57d7",
                ),
                (
                    "onnx/model_fp16.onnx_data",
                    "28d16e29cd623f25cc6fa0968700c5bc31036466091a5fa06d1353c1777f050e",
                ),
            ],
            Variant::Fp32 => [
                (
                    "onnx/model.onnx",
                    "cddb145cd1147ec24a3908b2ca2602b98b20a3d198365cff270b7cb26c98179e",
                ),
                (
                    "onnx/model.onnx_data",
                    "86a3a705d4598615894d89540ea71a3d9bbdb17a315e79edcd5dfc737222834b",
                ),
            ],
        }
    }

    /// What a store's variant counts call vectors this session made.
    pub fn name(self) -> &'static str {
        match self {
            Variant::Int8 => "int8-cpu",
            Variant::Fp16 => "fp16-cpu",
            Variant::Fp32 => "fp32-cpu",
        }
    }

    fn parse(text: &str) -> Option<Self> {
        match text {
            "int8" => Some(Variant::Int8),
            "fp16" => Some(Variant::Fp16),
            "fp32" => Some(Variant::Fp32),
            _ => None,
        }
    }
}

/// The variant an index pass embeds with, and the second one a `mix`
/// alternates with batch by batch. A harness knob, not part of the documented
/// environment: `int8` (the default), `fp16`, `fp32` or `mix`.
pub const VARIANT_ENV: &str = "SEMLITH_EMBED_VARIANT";

/// The variant a query is embedded with. Also a harness knob.
pub const QUERY_VARIANT_ENV: &str = "SEMLITH_QUERY_VARIANT";

pub fn index_variant() -> (Variant, Option<Variant>) {
    match std::env::var(VARIANT_ENV).ok().as_deref() {
        Some("mix") => (Variant::Int8, Some(Variant::Fp16)),
        Some(other) => (Variant::parse(other).unwrap_or(Variant::Int8), None),
        None => (Variant::Int8, None),
    }
}

pub fn query_variant() -> Variant {
    std::env::var(QUERY_VARIANT_ENV)
        .ok()
        .and_then(|v| Variant::parse(&v))
        .unwrap_or(Variant::Int8)
}

/// granite's tokenizer, prepared exactly as fastembed prepares it for the
/// CPU session: padding to the batch's longest with the model's pad token,
/// truncation at `max_length`, and the special tokens registered. The GPU
/// lanes run ONNX Runtime directly, and a tokenizer set up any other way
/// would give their vectors a different input from the CPU lane's.
pub fn granite_tokenizer(cache_dir: &Path, max_length: usize) -> Result<tokenizers::Tokenizer> {
    use tokenizers::{AddedToken, PaddingParams, PaddingStrategy, TruncationParams};
    let dir = snapshot_dir(cache_dir, GRANITE_REPO, GRANITE_REVISION)
        .context("granite is not cached yet; the CPU lane fetches it first")?;
    let read = |name: &str| -> Result<Vec<u8>> {
        let bytes = std::fs::read(dir.join(name)).with_context(|| format!("reading {name}"))?;
        let expected = GRANITE_FILES
            .iter()
            .find(|(file, _)| *file == name)
            .map(|(_, digest)| *digest)
            .with_context(|| format!("{name} is not a file semlith pins"))?;
        verify(name, &bytes, expected)?;
        Ok(bytes)
    };
    let config: serde_json::Value = serde_json::from_slice(&read("config.json")?)?;
    let special: serde_json::Value = serde_json::from_slice(&read("special_tokens_map.json")?)?;
    let tokenizer_config: serde_json::Value =
        serde_json::from_slice(&read("tokenizer_config.json")?)?;
    let mut tokenizer = tokenizers::Tokenizer::from_bytes(read("tokenizer.json")?)
        .map_err(|e| anyhow::anyhow!("reading the tokenizer: {e}"))?;
    let model_max = tokenizer_config["model_max_length"]
        .as_f64()
        .unwrap_or(max_length as f64) as usize;
    tokenizer
        .with_padding(Some(PaddingParams {
            strategy: PaddingStrategy::BatchLongest,
            pad_token: tokenizer_config["pad_token"]
                .as_str()
                .unwrap_or("[PAD]")
                .to_string(),
            pad_id: config["pad_token_id"].as_u64().unwrap_or(0) as u32,
            ..Default::default()
        }))
        .with_truncation(Some(TruncationParams {
            max_length: max_length.min(model_max),
            ..Default::default()
        }))
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    if let serde_json::Value::Object(map) = special {
        for value in map.values() {
            let token = match value {
                serde_json::Value::String(content) => AddedToken {
                    content: content.clone(),
                    special: true,
                    ..Default::default()
                },
                serde_json::Value::Object(_) => match (
                    value["content"].as_str(),
                    value["single_word"].as_bool(),
                    value["lstrip"].as_bool(),
                    value["rstrip"].as_bool(),
                    value["normalized"].as_bool(),
                ) {
                    (
                        Some(content),
                        Some(single_word),
                        Some(lstrip),
                        Some(rstrip),
                        Some(normalized),
                    ) => AddedToken {
                        content: content.into(),
                        special: true,
                        single_word,
                        lstrip,
                        rstrip,
                        normalized,
                    },
                    _ => continue,
                },
                _ => continue,
            };
            let _ = tokenizer.add_special_tokens([token]);
        }
    }
    Ok(tokenizer)
}

/// Point ONNX Runtime at the library shipped beside this binary.
///
/// Only in a `dynamic-ort` build, which is the two Linux release binaries and
/// nothing else. Every other build — `cargo install`, macOS, Windows, a
/// developer's own `cargo build` — links the library `ort` downloads at build
/// time and this function does nothing.
///
/// The reason the Linux artifacts differ is issue #57: the library `ort`
/// downloads references glibc 2.38 symbols, so the binary built against it
/// cannot start on Debian 12, Ubuntu 22.04 LTS, RHEL 9 or Amazon Linux 2023,
/// and building on an older runner cannot lower a floor the vendored library
/// sets. Microsoft's own ONNX Runtime release needs glibc 2.27, so the release
/// jobs pack that beside the binary and this loads it.
///
/// Run once, before anything asks for a model.
pub fn link_runtime() -> Result<()> {
    #[cfg(feature = "dynamic-ort")]
    let linked = {
        use std::sync::OnceLock;
        static ONCE: OnceLock<Result<(), String>> = OnceLock::new();
        ONCE.get_or_init(|| {
            let Ok(exe) = std::env::current_exe() else {
                return Err("the running binary could not be located".to_string());
            };
            let beside = exe
                .parent()
                .map(|dir| dir.join(RUNTIME_FILE))
                .unwrap_or_else(|| PathBuf::from(RUNTIME_FILE));
            if !beside.exists() {
                return Err(format!(
                    "{} is not beside the semlith binary. The Linux release \
                         archive carries it next to `semlith`, and `install.sh` \
                         puts both in the same directory — a binary copied out of \
                         the archive on its own cannot embed anything. Unpack the \
                         archive again, or install with the one-liner in the \
                         README.",
                    beside.display()
                ));
            }
            ort::init_from(beside.to_string_lossy().as_ref())
                .map_err(|e| format!("loading {}: {e}", beside.display()))?
                .commit();
            // `commit` answers whether this call was the one that
            // installed the environment; a second caller getting `false`
            // is the OnceLock doing its job, not a failure.
            Ok(())
        })
        .clone()
        .map_err(anyhow::Error::msg)
    };
    #[cfg(not(feature = "dynamic-ort"))]
    let linked = Ok(());
    linked
}

/// What a `dynamic-ort` build looks for beside itself.
#[cfg(feature = "dynamic-ort")]
const RUNTIME_FILE: &str = if cfg!(target_os = "macos") {
    "libonnxruntime.dylib"
} else if cfg!(windows) {
    "onnxruntime.dll"
} else {
    "libonnxruntime.so"
};

/// How long a model download may receive nothing before it is given up on.
///
/// A stalled connection used to leave `semlith setup` at 0 % CPU for as long
/// as anyone waited, with an empty `.part` file and no message (#155): the
/// Hugging Face client sets no read timeout. A minute is far past any pause a
/// working connection makes, and short enough that a person is still there.
/// `SEMLITH_DOWNLOAD_STALL` (seconds) overrides it.
const DOWNLOAD_STALL: std::time::Duration = std::time::Duration::from_secs(60);

fn download_stall() -> std::time::Duration {
    std::env::var("SEMLITH_DOWNLOAD_STALL")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .filter(|n| *n > 0)
        .map(std::time::Duration::from_secs)
        .unwrap_or(DOWNLOAD_STALL)
}

/// A pinned model repository, read from the cache or fetched file by file
/// under a watchdog that fails a download which stops receiving bytes.
pub(crate) struct Pinned {
    api: hf_hub::api::sync::Api,
    cache: hf_hub::Cache,
    repo: hf_hub::Repo,
    quiet: bool,
    stall: std::time::Duration,
}

impl Pinned {
    pub(crate) fn new(cache_dir: &Path, repo: &str, revision: &str, quiet: bool) -> Result<Self> {
        Self::at("https://huggingface.co", cache_dir, repo, revision, quiet)
    }

    fn at(
        endpoint: &str,
        cache_dir: &Path,
        repo: &str,
        revision: &str,
        quiet: bool,
    ) -> Result<Self> {
        let cache = hf_hub::Cache::new(cache_dir.to_path_buf());
        let api = hf_hub::api::sync::ApiBuilder::from_cache(cache.clone())
            .with_endpoint(endpoint.to_string())
            .with_progress(false)
            .build()
            .context("building the Hugging Face client")?;
        Ok(Self {
            api,
            cache,
            repo: hf_hub::Repo::with_revision(
                repo.to_string(),
                hf_hub::RepoType::Model,
                revision.to_string(),
            ),
            quiet,
            stall: download_stall(),
        })
    }

    /// The file's path in the cache, fetching it first if it is not there.
    pub(crate) fn get(&self, file: &str) -> Result<PathBuf> {
        if let Some(path) = self.cache.repo(self.repo.clone()).get(file) {
            return Ok(path);
        }
        let url = self.api.repo(self.repo.clone()).url(file);
        crate::add::note_outbound("model", &url);
        let watch = Watch::default();
        let (done, answer) = std::sync::mpsc::channel();
        let (api, repo, name, seen) = (
            self.api.clone(),
            self.repo.clone(),
            file.to_string(),
            watch.clone(),
        );
        // ponytail: a download that stalls is abandoned, not cancelled -- its
        // thread stays blocked in the client's read until the process exits,
        // which for `setup` is the next line. A cancellable read needs a
        // client with a read timeout.
        std::thread::spawn(move || {
            let _ = done.send(api.repo(repo).download_with_progress(&name, seen));
        });
        let mut said = std::time::Instant::now();
        loop {
            match answer.recv_timeout(std::time::Duration::from_millis(250)) {
                Ok(fetched) => return fetched.with_context(|| format!("fetching {url}")),
                Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
                    bail!("fetching {url}: the download ended without an answer")
                }
                Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
            }
            let (received, total, quiet_for) = watch.read();
            if quiet_for >= self.stall {
                bail!(
                    "the model download stalled: no bytes for {} s while fetching {url} \
                     ({} of {} received). Check this machine can reach huggingface.co, then \
                     run `semlith setup` again; it resumes into the same cache ({})",
                    self.stall.as_secs(),
                    crate::human_bytes(received as i64),
                    if total > 0 {
                        crate::human_bytes(total as i64)
                    } else {
                        "an unknown size".to_string()
                    },
                    self.cache.path().display(),
                );
            }
            if !self.quiet && total > 0 && said.elapsed() >= std::time::Duration::from_secs(2) {
                said = std::time::Instant::now();
                eprintln!(
                    "  {file}: {} of {}",
                    crate::human_bytes(received as i64),
                    crate::human_bytes(total as i64)
                );
            }
        }
    }
}

/// What a download has received and when it last received anything, shared
/// between the client's thread and the watchdog.
#[derive(Clone)]
struct Watch(std::sync::Arc<std::sync::Mutex<(u64, u64, std::time::Instant)>>);

impl Default for Watch {
    fn default() -> Self {
        Self(std::sync::Arc::new(std::sync::Mutex::new((
            0,
            0,
            std::time::Instant::now(),
        ))))
    }
}

impl Watch {
    /// `(received, total, time since the last byte)`.
    fn read(&self) -> (u64, u64, std::time::Duration) {
        let g = self.0.lock().unwrap_or_else(|e| e.into_inner());
        (g.0, g.1, g.2.elapsed())
    }
}

impl hf_hub::api::Progress for Watch {
    fn init(&mut self, size: usize, _filename: &str) {
        let mut g = self.0.lock().unwrap_or_else(|e| e.into_inner());
        g.1 = size as u64;
        g.2 = std::time::Instant::now();
    }
    fn update(&mut self, size: usize) {
        let mut g = self.0.lock().unwrap_or_else(|e| e.into_inner());
        g.0 += size as u64;
        g.2 = std::time::Instant::now();
    }
    fn finish(&mut self) {}
}

fn load_granite(
    cache_dir: PathBuf,
    max_length: usize,
    quiet: bool,
    threads: usize,
    variant: Variant,
) -> Result<TextEmbedding> {
    link_runtime()?;
    let paths = granite_paths(&cache_dir, quiet, variant)?;
    let read = |name: &str| -> Result<Vec<u8>> {
        let path = paths
            .iter()
            .find(|(file, _)| *file == name)
            .map(|(_, path)| path)
            .with_context(|| format!("{name} is not a file semlith pins a digest for"))?;
        std::fs::read(path).with_context(|| format!("reading {}", path.display()))
    };
    let [graph, weights] = variant.files();
    let tokenizer_files = TokenizerFiles {
        tokenizer_file: read("tokenizer.json")?,
        config_file: read("config.json")?,
        special_tokens_map_file: read("special_tokens_map.json")?,
        tokenizer_config_file: read("tokenizer_config.json")?,
    };

    // fastembed does not export ExternalInitializerFile, so the weights can
    // only be attached through this builder — a struct literal will not compile.
    let model = UserDefinedEmbeddingModel::new(read(graph.0)?, tokenizer_files)
        // 1_Pooling/config.json in the source repo sets pooling_mode_cls_token.
        .with_pooling(Pooling::Cls)
        .with_external_initializer(
            Path::new(weights.0)
                .file_name()
                .expect("weights constant has a file name")
                .to_string_lossy()
                .into_owned(),
            read(weights.0)?,
        );

    let opts = InitOptionsUserDefined::new()
        .with_max_length(max_length)
        .with_intra_threads(threads);

    TextEmbedding::try_new_from_user_defined(model, opts)
        .map_err(|e| anyhow::anyhow!("loading {GRANITE_NAME}: {e}"))
}

/// The graph file of one of granite's variants, fetched and verified, for a
/// session built on ONNX Runtime directly. Its weights file sits beside it
/// under the name the graph refers to, which is how ONNX Runtime finds them.
pub fn granite_graph(cache_dir: &Path, quiet: bool, variant: Variant) -> Result<PathBuf> {
    let [graph, _] = variant.files();
    granite_paths(cache_dir, quiet, variant)?
        .into_iter()
        .find(|(file, _)| *file == graph.0)
        .map(|(_, path)| path)
        .context("the graph is missing from its own pinned file list")
}

/// Every file one variant of granite needs — the tokenizer's four and the
/// variant's two — each in the cache and checked against its pinned digest.
///
/// The one place granite's weights are fetched, so the airgap refusal and the
/// digest check cannot be missed by a caller that loads them another way.
fn granite_paths(
    cache_dir: &Path,
    quiet: bool,
    variant: Variant,
) -> Result<Vec<(&'static str, PathBuf)>> {
    if airgap() && !is_cached(cache_dir) {
        bail!(
            "{AIRGAP_ENV} is set and no model is cached at {} — \
             pre-seed it with SEMLITH_MODEL_CACHE on a connected machine, \
             or drop --airgap to let this run download it",
            cache_dir.display()
        );
    }
    check_cache_dir(cache_dir)?;
    let cache = cache_dir.to_path_buf();
    // The tokenizer's four files and this variant's two. int8 keeps the stamp
    // it always had, so no cache written before 0.28.0 is checked again.
    let [graph, weights] = variant.files();
    let files: Vec<(&str, &str)> = GRANITE_FILES[..4]
        .iter()
        .copied()
        .chain([graph, weights])
        .collect();
    let stamp = match variant {
        Variant::Int8 => STAMP.to_string(),
        other => format!("{STAMP}-{}", other.name()),
    };

    // Whether this cache has already been checked against this pin, with every
    // file still the size and age it was. Asked once here rather than per file,
    // because the answer is about the set.
    let checked = snapshot_dir(&cache, GRANITE_REPO, GRANITE_REVISION)
        .is_some_and(|dir| already_verified_as(&dir, &stamp, GRANITE_REVISION, &files));

    // A revision rather than a branch. `main` is a name somebody else controls;
    // a commit is the bytes this release was built against.
    let repo = Pinned::new(&cache, GRANITE_REPO, GRANITE_REVISION, quiet)?;

    let mut paths = Vec::with_capacity(files.len());
    for (name, expected) in &files {
        let path = repo
            .get(name)
            .with_context(|| format!("fetching {name} from {GRANITE_REPO}"))?;
        // Verified whether it was just fetched or was already in the cache: a
        // cache is a directory on disk, and the point of a digest is that it
        // does not matter how the bytes got there.
        if !checked {
            let bytes =
                std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
            verify(name, &bytes, expected)?;
        }
        paths.push((*name, path));
    }

    // Recorded after every file has been read and checked, and only when this
    // run did the checking — the fetch above may have created the snapshot
    // directory that did not exist when `checked` was read.
    if !checked && let Some(dir) = snapshot_dir(&cache, GRANITE_REPO, GRANITE_REVISION) {
        record_verified_as(&dir, &stamp, GRANITE_REVISION, &files);
    }
    Ok(paths)
}

/// Refuse to download model weights, so an air-gapped machine can prove this
/// process never reached the network. Set by `--airgap` and readable directly.
pub const AIRGAP_ENV: &str = "SEMLITH_AIRGAP";

/// Find a builtin model's cached `tokenizer.json`.
///
/// fastembed does not publish where it put the files, so this looks for the
/// one file that matters under the cache root. Cheap and bounded: the cache
/// holds a handful of model directories, and a miss costs a directory walk and
/// returns `None`.
fn builtin_tokenizer(cache_dir: &Path, model: &EmbeddingModel) -> Option<tokenizers::Tokenizer> {
    let wanted = format!("{model}").to_lowercase().replace('/', "--");
    let entries = std::fs::read_dir(cache_dir).ok()?;
    for entry in entries.flatten() {
        let name = entry.file_name().to_string_lossy().to_lowercase();
        if !name.contains(&wanted) {
            continue;
        }
        for candidate in walk_for(&entry.path(), "tokenizer.json") {
            if let Ok(tokenizer) = tokenizers::Tokenizer::from_file(&candidate) {
                return Some(tokenizer);
            }
        }
    }
    None
}

/// Every `name` at most two directories below `root`.
fn walk_for(root: &Path, name: &str) -> Vec<PathBuf> {
    let mut found = Vec::new();
    let direct = root.join(name);
    if direct.is_file() {
        found.push(direct);
    }
    let Ok(entries) = std::fs::read_dir(root) else {
        return found;
    };
    for entry in entries.flatten() {
        if !entry.path().is_dir() {
            continue;
        }
        let nested = entry.path().join(name);
        if nested.is_file() {
            found.push(nested);
        }
    }
    found
}

pub fn airgap() -> bool {
    matches!(
        std::env::var(AIRGAP_ENV).ok().as_deref(),
        Some("1" | "true" | "yes")
    ) || crate::add::runtime_airgap()
}

/// Whether the cache already holds something to load.
///
/// Deliberately "is there anything here", not "are these the right files":
/// the loader below fails with its own, far more specific error if what is
/// cached is wrong, and duplicating its file list here would be a second
/// definition of the model to keep in step.
/// Refuse a model fetch on an airgapped machine, naming the model.
///
/// The same rule the text model is loaded under, in a form the image models can
/// use: an airgapped run's whole claim is that this process cannot have been
/// the one that reached the network, and a caller that fetches without asking
/// is a caller that breaks the claim.
///
/// This one asks about *that* model rather than about the cache as a whole. A
/// machine with the text model pre-seeded and no CLIP has a non-empty cache and
/// still cannot fetch, so the coarse "is anything cached" test would have let
/// the first image a store indexed open a socket on a machine whose whole point
/// is that it does not.
pub fn refuse_if_airgapped(model: &str) -> Result<()> {
    let cache = crate::model_cache_dir()?;
    if airgap() && !model_is_cached(&cache, model) {
        bail!(
            "{AIRGAP_ENV} is set and {model} is not cached at {} — \
             pre-seed it with SEMLITH_MODEL_CACHE on a connected machine, \
             or drop --airgap to let this run download it",
            cache.display()
        );
    }
    Ok(())
}

/// Whether one model's weights are already in the cache.
///
/// hf-hub names a model's directory `models--<org>--<name>`, so the repository
/// id is matched against the directory name with the separators normalised
/// rather than reconstructed — a naming change upstream should read as "not
/// cached", which refuses, rather than as "cached", which would not.
pub fn model_is_cached(cache_dir: &Path, model: &str) -> bool {
    let want = model.replace('/', "--").to_ascii_lowercase();
    std::fs::read_dir(cache_dir)
        .map(|entries| {
            entries.flatten().any(|entry| {
                let name = entry.file_name().to_string_lossy().to_ascii_lowercase();
                name.ends_with(&want) && entry.path().join("snapshots").exists()
            })
        })
        .unwrap_or(false)
}

/// Whether the default model is cached, with the right bytes.
///
/// "Is there anything here" until 0.14.0, which answered yes for a cache
/// holding half a download or somebody else's files. It now asks the question
/// the airgap check needs answered: can this run load the model without
/// reaching the network, and is what it would load the model this release pins.
///
/// A digest per file is the expensive form and this is called on every airgap
/// check, so the size is the cheap first pass: a file of the right length whose
/// digest is wrong is the case `load_granite` refuses by name a moment later.
pub fn is_cached(cache_dir: &Path) -> bool {
    let snapshot = match snapshot_dir(cache_dir, GRANITE_REPO, GRANITE_REVISION) {
        Some(dir) => dir,
        None => return false,
    };
    GRANITE_FILES
        .iter()
        .all(|(name, _)| snapshot.join(name).exists())
}

/// Where hf-hub puts one revision of one repository.
///
/// `<cache>/models--<org>--<name>/snapshots/<revision>`. Derived rather than
/// guessed at: hf-hub writes it, and a naming change upstream should read as
/// "not cached", which refuses, rather than as "cached", which would not.
fn snapshot_dir(cache_dir: &Path, repo: &str, revision: &str) -> Option<PathBuf> {
    let dir = cache_dir
        .join(format!("models--{}", repo.replace('/', "--")))
        .join("snapshots")
        .join(revision);
    dir.is_dir().then_some(dir)
}

/// Verify every pinned file of a repository that fastembed fetched itself.
///
/// `load_granite` verifies as it reads, because semlith does that fetching. The
/// image models are fastembed's own built-ins: it resolves and caches them
/// through its own client, so the check happens against what landed in the
/// cache rather than against bytes passing through semlith's hands. The
/// property is the same — a file that is not the one this release pins is
/// refused by name — and it is checked before the model is used.
pub fn verify_cached(
    cache_dir: &Path,
    repo: &str,
    revision: &str,
    files: &[(&str, &str)],
) -> Result<()> {
    check_cache_dir(cache_dir)?;
    let _ = revision;
    // Every snapshot under this repository, not only the pinned one. fastembed
    // resolves its built-in models at `main`, so a repository that moved leaves
    // its new commit here under a different name — and checking only the pinned
    // directory would pass by finding nothing. What is checked is what is on
    // disk to be loaded.
    let snapshots = cache_dir
        .join(format!("models--{}", repo.replace('/', "--")))
        .join("snapshots");
    let Ok(entries) = std::fs::read_dir(&snapshots) else {
        // Nothing cached. fastembed will fetch, and the next call through this
        // function is the one that checks what it fetched.
        return Ok(());
    };
    for entry in entries.flatten() {
        let dir = entry.path();
        // The same stamp the text model uses: the CLIP weights are 350 MB and
        // 254 MB, so hashing them on every load would cost more than loading
        // them does.
        if already_verified(&dir, revision, files) {
            continue;
        }
        let mut read_all = true;
        for (name, expected) in files {
            let path = dir.join(name);
            let Ok(bytes) = std::fs::read(&path) else {
                read_all = false;
                continue;
            };
            verify(&format!("{repo}/{name}"), &bytes, expected)?;
        }
        if read_all {
            record_verified(&dir, revision, files);
        }
    }
    Ok(())
}

/// The `threads each` value the daemon has in force, or 0 outside one.
static THREADS_IN_FORCE: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);

/// Set the value an index run's session is built with. The daemon calls this
/// at start and on every save of the setting; a writer whose session was built
/// with a different count rebuilds it at its next batch.
pub fn set_threads_in_force(threads: usize) {
    THREADS_IN_FORCE.store(threads, std::sync::atomic::Ordering::Relaxed);
}

/// Index passes embedding in this process right now.
static WRITERS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);

/// Held for the length of an index pass, so a derived thread count can be
/// split between the passes actually running rather than the most there
/// could be.
pub struct Writer(());

impl Drop for Writer {
    fn drop(&mut self) {
        WRITERS.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
    }
}

pub fn writer() -> Writer {
    WRITERS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    Writer(())
}

/// What an index run's session is built with: a count saved or set in the
/// environment, else the derived count split between the passes running now.
///
/// Before 0.28.0 every session was built with [`embed_threads`], so the
/// saved setting was shown on the page and never reached a session, and
/// `semlith start` printed "1 embedder thread(s) each" while every session ran
/// four. The first 0.28.0 build then split the derived count by `runs at
/// once` whether or not the other runs existed, and a lone run in the daemon
/// indexed at 54 % of the same run in a terminal on a four-core runner.
pub fn threads_in_force() -> usize {
    match THREADS_IN_FORCE.load(std::sync::atomic::Ordering::Relaxed) {
        0 => split_threads(WRITERS.load(std::sync::atomic::Ordering::Relaxed)),
        n => n,
    }
}

/// The derived count for `writers` passes at once: all of it for one, and for
/// more, an equal share held inside the cores less one kept free.
pub fn split_threads(writers: usize) -> usize {
    let cores = std::thread::available_parallelism().map_or(1, |n| n.get());
    split(embed_threads(), writers, cores)
}

fn split(all: usize, writers: usize, cores: usize) -> usize {
    if writers <= 1 {
        return all;
    }
    let budget = cores.saturating_sub(1).max(1);
    (all / writers).clamp(1, (budget / writers).max(1))
}

/// How many threads ONNX Runtime should use inside one operator.
///
/// ORT synchronises its threads at every operator boundary, so the slowest
/// thread paces the whole batch. On a CPU with both performance and efficiency
/// cores, a thread scheduled onto an efficiency core drags everything with it:
/// measured on a 4P+4E M1, four threads indexed at 16.5 chunks/s while eight
/// managed only 13.9, and one managed 5.1. Undersubscribing costs far more than
/// oversubscribing, so only heterogeneous machines get a reduced count.
pub fn embed_threads() -> usize {
    if let Ok(raw) = std::env::var(THREADS_ENV)
        && let Ok(n) = raw.parse::<usize>()
        && n > 0
    {
        return n;
    }
    performance_cores().unwrap_or_else(|| {
        std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1)
    })
}

/// Performance-core count on a hybrid CPU: Apple silicon's P-cores, and
/// Intel's P-cores on Linux and Windows. `None` where every core is equal, or
/// where the platform's answer is not one this recognises, and the caller then
/// uses the total, which is the right answer on a symmetric machine.
#[cfg(target_os = "macos")]
fn performance_cores() -> Option<usize> {
    let mut out: i32 = 0;
    let mut len = std::mem::size_of::<i32>();
    let name = c"hw.perflevel0.logicalcpu";
    // SAFETY: name is a NUL-terminated C string, and out/len describe a live
    // i32 whose size sysctlbyname is told about and will not exceed.
    let rc = unsafe {
        libc::sysctlbyname(
            name.as_ptr(),
            &mut out as *mut i32 as *mut libc::c_void,
            &mut len,
            std::ptr::null_mut(),
            0,
        )
    };
    (rc == 0 && out > 0).then_some(out as usize)
}

/// The kernel registers a `cpu_core` PMU beside `cpu_atom` only on Intel hybrid
/// parts, so the file existing is itself the test for a hybrid CPU; a symmetric
/// machine has a plain `cpu` PMU and no such directory.
#[cfg(target_os = "linux")]
fn performance_cores() -> Option<usize> {
    cpulist_len(&std::fs::read_to_string("/sys/devices/cpu_core/cpus").ok()?)
}

/// How many CPUs a kernel cpulist names: `0-15` is 16, `0-7,16-19` is 12.
///
/// Strict on purpose. A shape this does not know (a stride, a reversed range,
/// an empty list) is `None`, and a wrong P-core count is worse than none: none
/// falls back to every core, a wrong one starves or oversubscribes the pool.
#[cfg(any(target_os = "linux", test))]
fn cpulist_len(list: &str) -> Option<usize> {
    let mut count = 0;
    for part in list.trim().split(',') {
        let (first, last) = part.split_once('-').unwrap_or((part, part));
        let (first, last) = (first.parse::<usize>().ok()?, last.parse::<usize>().ok()?);
        count += last.checked_sub(first)? + 1;
    }
    (count > 0).then_some(count)
}

/// Windows ranks cores by `EfficiencyClass`, higher meaning faster, and gives
/// every core the same class on a symmetric CPU.
#[cfg(windows)]
fn performance_cores() -> Option<usize> {
    use windows_sys::Win32::System::SystemInformation::{
        CpuSetInformation, GetSystemCpuSetInformation, SYSTEM_CPU_SET_INFORMATION,
    };

    let mut len = 0u32;
    // SAFETY: a null buffer of length zero only asks for the size needed, which
    // the call writes to len, a live u32. It fails by design here.
    unsafe {
        GetSystemCpuSetInformation(std::ptr::null_mut(), 0, &mut len, std::ptr::null_mut(), 0)
    };
    let mut buf = vec![0u8; len as usize];
    // SAFETY: buf is live and exactly len bytes long, which is what the call is
    // told; it writes no more than that and reports what it wrote in len.
    let ok = unsafe {
        GetSystemCpuSetInformation(
            buf.as_mut_ptr().cast(),
            len,
            &mut len,
            std::ptr::null_mut(),
            0,
        )
    } != 0;
    if !ok || len as usize > buf.len() {
        return None;
    }

    // Entries are variable-length, each led by its own size, so the walk goes by
    // that field rather than by the struct's size.
    let mut classes = Vec::new();
    let mut at = 0;
    let whole = std::mem::size_of::<SYSTEM_CPU_SET_INFORMATION>();
    while at + whole <= len as usize {
        // SAFETY: at + whole is inside the bytes the call wrote, read_unaligned
        // needs no alignment, and every field is an integer, so any bytes are a
        // valid value; the union is read as the CpuSet variant only when Type
        // says that is what it holds.
        let entry: SYSTEM_CPU_SET_INFORMATION =
            unsafe { std::ptr::read_unaligned(buf.as_ptr().add(at).cast()) };
        if entry.Size == 0 {
            return None;
        }
        if entry.Type == CpuSetInformation {
            // SAFETY: Type says the union holds CpuSet, and a u8 has no
            // invalid values.
            classes.push(unsafe { entry.Anonymous.CpuSet.EfficiencyClass });
        }
        at += entry.Size as usize;
    }
    fastest_class_len(&classes)
}

/// How many cores sit in the highest efficiency class, when there is more than
/// one class. One class is a symmetric CPU, where the total is the answer.
#[cfg(any(windows, test))]
fn fastest_class_len(classes: &[u8]) -> Option<usize> {
    let fastest = *classes.iter().max()?;
    let count = classes.iter().filter(|c| **c == fastest).count();
    (count < classes.len()).then_some(count)
}

#[cfg(not(any(target_os = "macos", target_os = "linux", windows)))]
fn performance_cores() -> Option<usize> {
    None
}

#[cfg(test)]
mod performance_cores_tests {
    use super::*;

    /// A 12th-gen Core i9-12900K as Linux lists it: 8 P-cores with two threads
    /// each, then 8 E-cores in `cpu_atom`.
    #[test]
    fn a_hybrid_intel_cpulist_counts_its_performance_threads() {
        assert_eq!(cpulist_len("0-15\n"), Some(16));
        assert_eq!(cpulist_len("16-23\n"), Some(8));
        assert_eq!(cpulist_len("0-7,16-19"), Some(12));
        assert_eq!(cpulist_len("3"), Some(1));
    }

    #[test]
    fn a_cpulist_this_does_not_know_is_none() {
        for list in ["", "\n", "0-15:2/4", "15-0", "0-", "a-b", "0,,2"] {
            assert_eq!(cpulist_len(list), None, "{list:?}");
        }
    }

    /// Eight cores at class 1 and eight at class 0, as `GetSystemCpuSetInformation`
    /// lists a hybrid part.
    #[test]
    fn two_efficiency_classes_count_the_higher() {
        let classes = [[1u8; 8], [0u8; 8]].concat();
        assert_eq!(fastest_class_len(&classes), Some(8));
    }

    #[test]
    fn one_efficiency_class_is_a_symmetric_cpu() {
        assert_eq!(fastest_class_len(&[0; 16]), None);
        assert_eq!(fastest_class_len(&[]), None);
    }
}

// ---------------------------------------------------------------- pinned weights

/// The commit each model repository is fetched at.
///
/// A Hugging Face repository is a git repository somebody else can push to, and
/// before 0.14.0 semlith fetched from whatever `main` pointed at. The weights
/// are what computes every vector in every store: a model that changed under a
/// user would change what their corpus means without changing anything they can
/// see, and a model that was replaced would do it on purpose.
pub const GRANITE_REVISION: &str = "1dc7835ba0cb9c76a3618d0bf0c427c97671b3c8";

/// Every file semlith fetches from [`GRANITE_REPO`], with its SHA-256.
///
/// Resolved against the repository at [`GRANITE_REVISION`] and recorded in
/// `docs/models.md` with the commit URL, so the number in this table can be
/// checked against the one Hugging Face publishes without reading this file.
pub const GRANITE_FILES: &[(&str, &str)] = &[
    (
        "tokenizer.json",
        "feeb83348dcb033bc6b9d2e1f7906ca9eb2d122845000c9416d894d7c2927149",
    ),
    (
        "config.json",
        "1a1710c20911da8c96179716bf44058e54cca6fa7952cce77be83ef05edae3ee",
    ),
    (
        "special_tokens_map.json",
        "ea97ecdbcc73713039d8d64dbb05e3689495c96657fbd9a18f5bed381be81049",
    ),
    (
        "tokenizer_config.json",
        "ce06781b38bb393db68c9e0709bddd31ef5d88f2c6fbb3fd9f369778fb85e451",
    ),
    (
        GRANITE_ONNX,
        "a3fad524afc3f060216a8ddbb1ac89c9b6498fba8995b5718bde879076a2e9ba",
    ),
    (
        GRANITE_WEIGHTS,
        "1f4cf47e4adec7f7ae09db03d071ba8667e07f9a4203142c7efa8d37fe453597",
    ),
];

/// What a cache directory records about the files this release already verified.
///
/// Hashing every pinned file on every load is 52 MB of SHA-256 before a process
/// can answer anything — measured at 276 ms on an M1, which is most of what a
/// `semlith search` from the command line costs. So the digests are checked
/// once and the result recorded here, keyed by the pin and by each file's size
/// and modification time: if all three still match, the bytes are the ones that
/// were verified.
///
/// This does not weaken the check it replaces. The attacker it would have to
/// let through is one who can write this cache *and* preserve each file's size
/// and mtime — and a cache anybody but this user can write to is already
/// refused by [`check_cache_dir`], so that attacker is this user. What the
/// digests are actually for — a corrupted download, and an upstream repository
/// whose bytes changed — moves a file's size or mtime every time.
#[derive(serde::Serialize, serde::Deserialize)]
struct Verified {
    revision: String,
    /// File name to `(size, mtime nanoseconds, the digest it was checked
    /// against)`.
    files: std::collections::BTreeMap<String, (u64, u128, String)>,
}

const STAMP: &str = ".semlith-verified";

/// What a file looks like on disk right now, for the stamp.
fn fingerprint(path: &Path) -> Option<(u64, u128)> {
    let meta = std::fs::metadata(path).ok()?;
    let modified = meta
        .modified()
        .ok()?
        .duration_since(std::time::UNIX_EPOCH)
        .ok()?
        .as_nanos();
    Some((meta.len(), modified))
}

/// Whether every pinned file was verified at this revision and has not moved.
fn already_verified(dir: &Path, revision: &str, files: &[(&str, &str)]) -> bool {
    already_verified_as(dir, STAMP, revision, files)
}

fn already_verified_as(dir: &Path, stamp: &str, revision: &str, files: &[(&str, &str)]) -> bool {
    let Ok(text) = std::fs::read_to_string(dir.join(stamp)) else {
        return false;
    };
    let Ok(stamp) = serde_json::from_str::<Verified>(&text) else {
        return false;
    };
    if stamp.revision != revision {
        return false;
    }
    files.iter().all(|(name, expected)| {
        stamp
            .files
            .get(*name)
            .zip(fingerprint(&dir.join(name)))
            .is_some_and(|((size, mtime, digest), (now_size, now_mtime))| {
                digest == expected && *size == now_size && *mtime == now_mtime
            })
    })
}

/// Record that every pinned file has been checked, so the next process need not.
fn record_verified(dir: &Path, revision: &str, files: &[(&str, &str)]) {
    record_verified_as(dir, STAMP, revision, files)
}

fn record_verified_as(dir: &Path, stamp_name: &str, revision: &str, files: &[(&str, &str)]) {
    let mut stamp = Verified {
        revision: revision.to_string(),
        files: std::collections::BTreeMap::new(),
    };
    for (name, expected) in files {
        let Some((size, mtime)) = fingerprint(&dir.join(name)) else {
            // A file that cannot be measured is one the next run should check
            // again, so the stamp is not written at all.
            return;
        };
        stamp
            .files
            .insert((*name).to_string(), (size, mtime, (*expected).to_string()));
    }
    let Ok(body) = serde_json::to_vec(&stamp) else {
        return;
    };
    // Best effort: a cache that cannot be written to is one that gets verified
    // every time, which is slower and not wrong.
    let _ = crate::home::write_private(&dir.join(stamp_name), &body);
}

/// The SHA-256 of some bytes, as lowercase hex.
pub fn digest(bytes: &[u8]) -> String {
    use sha2::{Digest, Sha256};
    // sha2 0.11 finalises to a `hybrid-array` value, which does not implement
    // `LowerHex` the way the old `GenericArray` did. Thirty-two bytes is not
    // worth a hex crate.
    Sha256::digest(bytes)
        .iter()
        .map(|b| format!("{b:02x}"))
        .collect()
}

/// Check bytes against the digest recorded for them, or say which file failed.
pub fn verify(name: &str, bytes: &[u8], expected: &str) -> Result<()> {
    let actual = digest(bytes);
    if actual != expected {
        bail!(
            "{name} does not match the digest semlith pins for it.\n  \
             expected {expected}\n  got      {actual}\n\
             The model cache holds a file that is not the one this release was \
             built against. Delete the cache and let semlith fetch it again; if \
             it happens twice, the file upstream has changed and semlith needs a \
             release rather than a retry."
        );
    }
    Ok(())
}

/// Refuse a model cache that somebody else on this machine can write to.
///
/// The cache is weights, and the weights are what every vector is computed by.
/// A directory another account can write to is a model another account
/// chooses — and the choosing is invisible, because a corpus embedded by a
/// different model still answers, just differently.
pub fn check_cache_dir(cache: &Path) -> Result<()> {
    #[cfg(unix)]
    {
        use std::os::unix::fs::MetadataExt;
        use std::os::unix::fs::PermissionsExt;
        let Ok(meta) = std::fs::metadata(cache) else {
            return Ok(());
        };
        // SAFETY: `getuid` reads this process's own id and cannot fail.
        let me = unsafe { libc::getuid() };
        if meta.uid() != me {
            bail!(
                "{} is owned by uid {} rather than by you. semlith will not load \
                 model weights from a directory somebody else owns — run `chown -R \
                 {me} {}`, or point {} somewhere you own.",
                cache.display(),
                meta.uid(),
                cache.display(),
                crate::MODEL_CACHE_ENV,
            );
        }
        let mode = meta.permissions().mode() & 0o777;
        if mode & 0o022 != 0 {
            bail!(
                "{} is mode {mode:o}, so other users on this machine can write to \
                 it. semlith will not load model weights from it — run `chmod 700 \
                 {}`.",
                cache.display(),
                cache.display(),
            );
        }
    }
    #[cfg(not(unix))]
    let _ = cache;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn granite_round_trips_through_its_stored_name() {
        let parsed: Model = GRANITE_NAME.parse().unwrap();
        assert_eq!(parsed, Model::Granite);
        assert_eq!(parsed.to_string(), GRANITE_NAME);
        assert_eq!(parsed.dim().unwrap(), 384);
    }

    #[test]
    fn stores_written_by_0_1_0_still_parse() {
        // The exact spelling a 0.1.0 store holds in its `model` meta row.
        let parsed: Model = "BGESmallENV15".parse().unwrap();
        assert_eq!(parsed, Model::Builtin(EmbeddingModel::BGESmallENV15));
        assert_eq!(parsed.to_string(), "BGESmallENV15");
        assert_eq!(parsed.dim().unwrap(), 384);
    }

    #[test]
    fn nonsense_model_names_are_rejected() {
        assert!("not-a-model".parse::<Model>().is_err());
    }

    #[test]
    fn only_bge_english_gets_the_instruction_prefix() {
        let bge = Model::Builtin(EmbeddingModel::BGESmallENV15);
        assert!(bge.query_text("hi").starts_with("Represent this sentence"));
        assert!(bge.query_text("hi").ends_with("hi"));
        assert_eq!(Model::Granite.query_text("hi"), "hi");
    }

    #[test]
    fn thread_count_honours_the_override() {
        // Serialised against the other env-reading test by using a distinct value.
        unsafe { std::env::set_var(THREADS_ENV, "3") };
        assert_eq!(embed_threads(), 3);
        unsafe { std::env::remove_var(THREADS_ENV) };
        assert!(embed_threads() >= 1);
    }

    /// A run going alone gets every derived thread, as in a terminal; only
    /// runs embedding at the same time share them.
    #[test]
    fn a_lone_writer_gets_every_derived_thread() {
        assert_eq!(split(4, 0, 4), 4);
        assert_eq!(split(4, 1, 4), 4, "a four-core runner, one run");
        assert_eq!(split(4, 3, 8), 1, "the M1, three runs");
        assert_eq!(split(4, 2, 8), 2);
        for writers in 2..=8 {
            let each = split(4, writers, 2);
            assert!((1..=4).contains(&each), "{writers} writers got {each}");
        }
    }

    /// The digest a store's content hashes and the upgrade checksum are built
    /// from. These three strings are SHA-256's own published vectors, so they
    /// are the same under sha2 0.10 and 0.11 -- which is the point: the 0.11
    /// bump changed the return type, never the bytes, and no store's hashes
    /// move under it.
    #[test]
    fn the_digest_is_lowercase_hex_of_the_published_vectors() {
        assert_eq!(
            digest(b""),
            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
        );
        assert_eq!(
            digest(b"abc"),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );
        assert_eq!(
            digest(b"semlith"),
            digest("semlith".as_bytes()),
            "the same bytes hash the same way whatever produced them"
        );
        assert_eq!(
            digest(b"").len(),
            64,
            "sixty-four hex characters, never a shorter form"
        );
    }

    /// #155: a server that accepts and never answers fails the fetch once the
    /// stall interval passes, naming the URL, the bytes received and the
    /// command that repairs it -- where it used to wait for ever.
    #[test]
    fn a_stalled_download_fails_with_a_reason_instead_of_waiting() {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        // Accept and hold every connection open without a byte of answer.
        std::thread::spawn(move || {
            let mut held = Vec::new();
            for stream in listener.incoming().flatten() {
                held.push(stream);
            }
        });
        let cache = std::env::temp_dir().join(format!("semlith-stall-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&cache);
        let mut pinned = Pinned::at(
            &format!("http://127.0.0.1:{port}"),
            &cache,
            "someone/some-model",
            "main",
            true,
        )
        .unwrap();
        pinned.stall = std::time::Duration::from_secs(1);
        let started = std::time::Instant::now();
        let error = format!("{:#}", pinned.get("model.onnx").unwrap_err());
        assert!(
            started.elapsed() < std::time::Duration::from_secs(10),
            "{error}"
        );
        assert!(error.contains("stalled"), "{error}");
        assert!(error.contains("someone/some-model"), "{error}");
        assert!(error.contains("0 B of an unknown size received"), "{error}");
        assert!(error.contains("semlith setup"), "{error}");
        let _ = std::fs::remove_dir_all(&cache);
    }
}