modelshelf 0.1.0

//! Ollama repacks model weights into a content-addressed blob store:
//!
//! ```text
//! ~/.ollama/models/
//! ├── manifests/<registry>/<namespace>/<model>/<tag>   # JSON manifest
//! └── blobs/sha256-<64 hex>                            # raw layer bytes
//! ```
//!
//! The manifest's `application/vnd.ollama.image.model` layer digest names
//! the blob holding the GGUF weights, so the sha256 is known without hashing.
//! This entire tree is read-only for modelshelf: we adopt blobs *from* it via
//! hardlink but never modify anything inside it.

use std::path::{Path, PathBuf};

use serde::Deserialize;

use super::{FoundFile, ScanEnv, Scanner};
use crate::registry::Ecosystem;

pub(crate) struct OllamaScanner;

const MODEL_MEDIA_TYPE: &str = "application/vnd.ollama.image.model";

#[derive(Deserialize)]
struct Manifest {
    #[serde(default)]
    layers: Vec<Layer>,
}

#[derive(Deserialize)]
struct Layer {
    #[serde(rename = "mediaType")]
    media_type: String,
    digest: String,
    #[serde(default)]
    size: u64,
}

impl Scanner for OllamaScanner {
    fn ecosystem(&self) -> Ecosystem {
        Ecosystem::Ollama
    }

    fn detect_root(&self, env: &ScanEnv) -> Option<PathBuf> {
        if let Some(root) = env.root_override(self.ecosystem()) {
            return root.is_dir().then(|| root.to_path_buf());
        }
        // Respect a relocated store, like Ollama itself does.
        if let Some(env_root) = std::env::var_os("OLLAMA_MODELS") {
            let p = PathBuf::from(env_root);
            if p.is_dir() {
                return Some(p);
            }
        }
        let candidate = env.home()?.join(".ollama").join("models");
        candidate.is_dir().then_some(candidate)
    }

    fn scan(&self, root: &Path) -> Vec<FoundFile> {
        let manifests_dir = root.join("manifests");
        let blobs_dir = root.join("blobs");
        let mut found = Vec::new();

        for entry in walkdir::WalkDir::new(&manifests_dir)
            .follow_links(false)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.file_type().is_file())
        {
            let Ok(bytes) = std::fs::read(entry.path()) else {
                continue;
            };
            let Ok(manifest) = serde_json::from_slice::<Manifest>(&bytes) else {
                tracing::debug!(path = %entry.path().display(), "skipping unparsable Ollama manifest");
                continue;
            };
            let Some(layer) = manifest
                .layers
                .iter()
                .find(|l| l.media_type == MODEL_MEDIA_TYPE)
            else {
                continue;
            };
            let Some(hex) = layer.digest.strip_prefix("sha256:") else {
                continue;
            };
            if !super::is_sha256_hex(hex) {
                continue;
            }
            let blob = blobs_dir.join(format!("sha256-{hex}"));
            if !blob.is_file() {
                tracing::debug!(blob = %blob.display(), "Ollama manifest references missing blob");
                continue;
            }
            let size = std::fs::metadata(&blob)
                .map(|m| m.len())
                .unwrap_or(layer.size);
            found.push(FoundFile {
                path: blob,
                ecosystem: Ecosystem::Ollama,
                size_bytes: size,
                display_name: display_name_from_manifest_path(&manifests_dir, entry.path()),
                known_sha256: Some(hex.to_owned()),
                source: None,
            });
        }
        found
    }
}

/// `manifests/registry.ollama.ai/library/llama3/latest` → `llama3:latest`.
fn display_name_from_manifest_path(manifests_dir: &Path, manifest: &Path) -> String {
    let rel: Vec<String> = manifest
        .strip_prefix(manifests_dir)
        .map(|r| {
            r.components()
                .map(|c| c.as_os_str().to_string_lossy().into_owned())
                .collect()
        })
        .unwrap_or_default();
    match rel.as_slice() {
        // registry / namespace / model / tag
        [_, ns, model, tag] if ns == "library" => format!("{model}:{tag}"),
        [_, ns, model, tag] => format!("{ns}/{model}:{tag}"),
        _ => manifest.to_string_lossy().into_owned(),
    }
}

#[cfg(any(test, feature = "test-fixtures"))]
pub mod fixtures {
    //! Build a fake-but-structurally-real Ollama store for tests.
    use std::path::Path;

    /// Write `content` as a blob plus a manifest naming it, exactly like
    /// `ollama pull <model>:<tag>` would lay it out.
    pub fn add_model(root: &Path, model: &str, tag: &str, content: &[u8]) -> String {
        use sha2::Digest;
        let hex = {
            let mut s = String::new();
            for b in sha2::Sha256::digest(content) {
                use std::fmt::Write;
                let _ = write!(s, "{b:02x}");
            }
            s
        };
        let blobs = root.join("blobs");
        std::fs::create_dir_all(&blobs).unwrap();
        std::fs::write(blobs.join(format!("sha256-{hex}")), content).unwrap();

        let manifest_dir = root
            .join("manifests")
            .join("registry.ollama.ai")
            .join("library")
            .join(model);
        std::fs::create_dir_all(&manifest_dir).unwrap();
        let manifest = serde_json::json!({
            "schemaVersion": 2,
            "mediaType": "application/vnd.docker.distribution.manifest.v2+json",
            "layers": [
                {
                    "mediaType": "application/vnd.ollama.image.model",
                    "digest": format!("sha256:{hex}"),
                    "size": content.len(),
                },
                {
                    "mediaType": "application/vnd.ollama.image.params",
                    "digest": "sha256:0000000000000000000000000000000000000000000000000000000000000000",
                    "size": 2,
                }
            ]
        });
        std::fs::write(
            manifest_dir.join(tag),
            serde_json::to_vec_pretty(&manifest).unwrap(),
        )
        .unwrap();
        hex
    }
}

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

    #[test]
    fn finds_model_blob_via_manifest_with_free_hash() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        let content = crate::gguf::fixtures::gguf_bytes("Llama Test", "llama", 15, 4096, b"OLLAMA");
        let hex = fixtures::add_model(root, "llama-test", "latest", &content);

        let found = OllamaScanner.scan(root);
        assert_eq!(found.len(), 1);
        assert_eq!(found[0].known_sha256.as_deref(), Some(hex.as_str()));
        assert_eq!(found[0].display_name, "llama-test:latest");
        assert!(found[0].path.ends_with(format!("sha256-{hex}")));
    }

    #[test]
    fn missing_blob_is_skipped_not_fatal() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        fixtures::add_model(root, "ghost", "latest", b"CONTENT");
        // Remove the blob but keep the manifest.
        for e in std::fs::read_dir(root.join("blobs")).unwrap() {
            std::fs::remove_file(e.unwrap().path()).unwrap();
        }
        assert!(OllamaScanner.scan(root).is_empty());
    }
}