modelshelf 0.1.0

//! Update checks: compare local content hashes against the current state of
//! each model's Hugging Face source — without downloading anything.

use std::collections::BTreeMap;

use crate::download::hf::HfClient;
use crate::registry::{ModelEntry, ModelId};
use crate::Result;

/// Options for [`crate::Shelf::check_updates`].
#[derive(Debug, Clone, Default)]
pub struct UpdateOptions {
    /// Alternative hub endpoint (tests, mirrors). `None` = huggingface.co.
    pub endpoint: Option<String>,
}

/// A newer revision of a locally present model.
#[derive(Debug, Clone)]
pub struct UpdateInfo {
    /// The local (old) model.
    pub id: ModelId,
    /// Display name of the local model.
    pub display_name: String,
    /// Source repo.
    pub repo: String,
    /// File within the repo.
    pub filename: String,
    /// The revision (commit sha) offering the newer file.
    pub latest_revision: String,
    /// sha256 of the newer content, when the hub reports it (LFS oid).
    pub new_sha256: Option<String>,
    /// Size of the newer content, when reported.
    pub new_size: Option<u64>,
}

/// A repo whose update check failed (network, deleted repo, auth).
#[derive(Debug, Clone)]
pub struct UpdateCheckError {
    /// The repo that could not be checked.
    pub repo: String,
    /// Why.
    pub reason: String,
}

/// Result of an update check.
#[derive(Debug, Clone, Default)]
pub struct UpdateReport {
    /// Models with a newer version available.
    pub updates: Vec<UpdateInfo>,
    /// Repos that could not be checked.
    pub errors: Vec<UpdateCheckError>,
}

/// Check every model with a Hugging Face source. One API call per distinct
/// repo, never more.
pub(crate) fn check_updates(client: &HfClient, models: &[ModelEntry]) -> Result<UpdateReport> {
    let mut by_repo: BTreeMap<String, Vec<&ModelEntry>> = BTreeMap::new();
    for m in models {
        if let Some(src) = &m.source {
            if src.kind == "huggingface" {
                by_repo.entry(src.repo.clone()).or_default().push(m);
            }
        }
    }

    let mut report = UpdateReport::default();
    for (repo, entries) in by_repo {
        let info = match client.repo_info(&repo) {
            Ok(i) => i,
            Err(e) => {
                report.errors.push(UpdateCheckError {
                    repo,
                    reason: e.to_string(),
                });
                continue;
            }
        };
        for m in entries {
            let src = m.source.as_ref().expect("filtered above");
            let Some(remote) = info.file(&src.filename) else {
                // File no longer exists at this path; not an "update".
                continue;
            };
            let remote_sha = remote.sha256();
            let local_sha = m.id.sha256_hex().map(str::to_lowercase);
            let changed = match (&remote_sha, &local_sha) {
                (Some(r), Some(l)) => r != l,
                // No LFS oid (e.g. non-LFS storage): fall back to comparing
                // revisions when we recorded one; otherwise assume unchanged
                // rather than nagging about unverifiable "updates".
                _ => src
                    .revision
                    .as_ref()
                    .is_some_and(|local_rev| local_rev != &info.sha && remote_sha.is_none()),
            };
            if changed {
                report.updates.push(UpdateInfo {
                    id: m.id.clone(),
                    display_name: m.display_name.clone(),
                    repo: src.repo.clone(),
                    filename: src.filename.clone(),
                    latest_revision: info.sha.clone(),
                    new_sha256: remote_sha.clone(),
                    new_size: remote.lfs.as_ref().and_then(|l| l.size).or(remote.size),
                });
            }
        }
    }
    Ok(report)
}