modelshelf 0.1.0

//! Hugging Face Hub API client.
//!
//! Secrets policy: the token is read from the environment (`HF_TOKEN` or
//! `HUGGING_FACE_HUB_TOKEN`), attached only to outgoing requests, and never
//! appears in `Debug` output, logs, or error messages. Tests assert this.

use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use serde::Deserialize;
use sha2::{Digest, Sha256};

use super::{Progress, ProgressFn};
use crate::{Error, Result};

/// Default API endpoint.
pub const DEFAULT_ENDPOINT: &str = "https://huggingface.co";

/// Environment variables consulted for the access token, in order.
pub const TOKEN_ENV_VARS: [&str; 2] = ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"];

/// Metadata for one file in a repo, from the hub API.
#[derive(Debug, Clone, Deserialize)]
pub struct RepoFile {
    /// Path of the file within the repo.
    pub rfilename: String,
    /// Size in bytes, when reported.
    #[serde(default)]
    pub size: Option<u64>,
    /// LFS pointer info carrying the content sha256, when reported.
    #[serde(default)]
    pub lfs: Option<LfsInfo>,
}

impl RepoFile {
    /// The content sha256 as lowercase hex, when the hub reports one.
    pub fn sha256(&self) -> Option<String> {
        let oid = self.lfs.as_ref()?.oid.as_deref()?;
        let hex = oid.strip_prefix("sha256:").unwrap_or(oid);
        Some(hex.to_lowercase())
    }
}

/// LFS object info.
///
/// The hub has spelled the content hash differently over time (`oid`
/// historically, `sha256` since the Xet migration), and some entries carry
/// neither — parsing must never fail over this field.
#[derive(Debug, Clone, Deserialize)]
pub struct LfsInfo {
    /// sha256 hex of the file content (possibly `sha256:`-prefixed).
    #[serde(default, alias = "sha256")]
    pub oid: Option<String>,
    /// Content size in bytes.
    #[serde(default)]
    pub size: Option<u64>,
}

/// Repo-level metadata from `GET /api/models/{repo}?blobs=true`.
#[derive(Debug, Clone, Deserialize)]
pub struct RepoInfo {
    /// The commit sha of the resolved revision.
    pub sha: String,
    /// Files in the repo.
    #[serde(default)]
    pub siblings: Vec<RepoFile>,
}

impl RepoInfo {
    /// Find a sibling by exact rfilename.
    pub fn file(&self, rfilename: &str) -> Option<&RepoFile> {
        self.siblings.iter().find(|f| f.rfilename == rfilename)
    }

    /// All `.gguf` siblings.
    pub fn gguf_files(&self) -> Vec<&RepoFile> {
        self.siblings
            .iter()
            .filter(|f| f.rfilename.to_lowercase().ends_with(".gguf"))
            .collect()
    }
}

/// Minimal Hugging Face Hub client.
pub struct HfClient {
    endpoint: String,
    agent: ureq::Agent,
    token: Option<String>,
}

// Manual Debug: the token must never leak through debug formatting.
impl std::fmt::Debug for HfClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HfClient")
            .field("endpoint", &self.endpoint)
            .field("token", &self.token.as_ref().map(|_| "<redacted>"))
            .finish()
    }
}

/// Environment variable overriding the hub endpoint when none is given
/// explicitly (mirrors, tests).
pub const ENDPOINT_ENV: &str = "MODELSHELF_HF_ENDPOINT";

impl HfClient {
    /// Client for `endpoint` (defaulting to [`ENDPOINT_ENV`], then
    /// huggingface.co) reading the token from the environment.
    pub fn from_env(endpoint: Option<&str>) -> Self {
        let env_endpoint = std::env::var(ENDPOINT_ENV).ok().filter(|s| !s.is_empty());
        let endpoint = endpoint.map(str::to_owned).or(env_endpoint);
        let token = TOKEN_ENV_VARS
            .iter()
            .find_map(|k| std::env::var(k).ok())
            .filter(|t| !t.is_empty());
        Self::new(endpoint.as_deref(), token)
    }

    /// Client with an explicit token (or none).
    pub fn new(endpoint: Option<&str>, token: Option<String>) -> Self {
        HfClient {
            endpoint: endpoint
                .unwrap_or(DEFAULT_ENDPOINT)
                .trim_end_matches('/')
                .to_owned(),
            agent: ureq::AgentBuilder::new()
                .user_agent(concat!("modelshelf/", env!("CARGO_PKG_VERSION")))
                .build(),
            token,
        }
    }

    fn with_auth(&self, req: ureq::Request) -> ureq::Request {
        match &self.token {
            Some(t) => req.set("Authorization", &format!("Bearer {t}")),
            None => req,
        }
    }

    /// `ureq::Error` may echo request details; keep only status/transport
    /// info so a token can never surface in an error message.
    fn sanitize_err(context: &str, e: ureq::Error) -> Error {
        match e {
            ureq::Error::Status(code, _) => {
                Error::Network(format!("{context}: HTTP status {code}"))
            }
            ureq::Error::Transport(t) => {
                Error::Network(format!("{context}: transport error ({:?})", t.kind()))
            }
        }
    }

    /// Fetch repo metadata including per-file LFS sha256 oids.
    pub fn repo_info(&self, repo: &str) -> Result<RepoInfo> {
        let url = format!("{}/api/models/{}?blobs=true", self.endpoint, repo);
        let resp = self
            .with_auth(self.agent.get(&url))
            .call()
            .map_err(|e| Self::sanitize_err(&format!("repo info for {repo}"), e))?;
        resp.into_json::<RepoInfo>()
            .map_err(|e| Error::Network(format!("repo info for {repo}: invalid JSON ({e})")))
    }

    /// Where a (possibly partial) download of this file lives inside
    /// `dest_dir`. Stable across runs — this is what makes resume work.
    pub fn part_path(dest_dir: &Path, repo: &str, revision: &str, filename: &str) -> PathBuf {
        dest_dir.join(format!("{}.part", job_key(repo, revision, filename)))
    }

    /// Download `filename` at `revision` into `dest_dir`, resuming a partial
    /// download if one exists. The content sha256 is computed while
    /// streaming; when `expected_sha256` is given, a mismatch removes the
    /// file and fails. Returns the path of the completed `.part` file, which
    /// the caller moves into its final location.
    pub fn download_file(
        &self,
        repo: &str,
        revision: &str,
        filename: &str,
        expected_sha256: Option<&str>,
        dest_dir: &Path,
        progress: Option<&ProgressFn>,
    ) -> Result<DownloadedFile> {
        std::fs::create_dir_all(dest_dir).map_err(|e| Error::io(dest_dir, e))?;
        let part_path = Self::part_path(dest_dir, repo, revision, filename);
        let url = format!(
            "{}/{}/resolve/{}/{}",
            self.endpoint, repo, revision, filename
        );

        let mut offset = std::fs::metadata(&part_path).map(|m| m.len()).unwrap_or(0);
        let mut req = self.with_auth(self.agent.get(&url));
        if offset > 0 {
            req = req.set("Range", &format!("bytes={offset}-"));
        }
        let resp = match req.call() {
            Ok(r) => r,
            // 416: the .part already covers the whole file (e.g. the final
            // rename was interrupted). Fall through to verification.
            Err(ureq::Error::Status(416, _)) if offset > 0 => {
                return self.finalize(&part_path, offset, expected_sha256);
            }
            Err(e) => return Err(Self::sanitize_err(&format!("download {filename}"), e)),
        };

        // Server honored the range only if it answered 206.
        let (mut file, resumed) = if offset > 0 && resp.status() == 206 {
            let f = std::fs::OpenOptions::new()
                .append(true)
                .open(&part_path)
                .map_err(|e| Error::io(&part_path, e))?;
            (f, true)
        } else {
            offset = 0;
            let f = std::fs::File::create(&part_path).map_err(|e| Error::io(&part_path, e))?;
            (f, false)
        };

        let total_bytes = if resumed {
            content_range_total(resp.header("Content-Range"))
        } else {
            resp.header("Content-Length").and_then(|v| v.parse().ok())
        };

        let mut reader = resp.into_reader();
        let mut buf = vec![0u8; 256 * 1024];
        let mut downloaded = offset;
        loop {
            let n = reader
                .read(&mut buf)
                .map_err(|e| Error::Network(format!("download {filename}: read failed ({e})")))?;
            if n == 0 {
                break;
            }
            file.write_all(&buf[..n])
                .map_err(|e| Error::io(&part_path, e))?;
            downloaded += n as u64;
            if let Some(cb) = progress {
                cb(Progress {
                    downloaded_bytes: downloaded,
                    total_bytes: total_bytes.map(|t: u64| t + if resumed { offset } else { 0 }),
                });
            }
        }
        file.sync_all().map_err(|e| Error::io(&part_path, e))?;
        drop(file);

        self.finalize(&part_path, downloaded, expected_sha256)
    }

    /// Hash the completed part file and verify it, if a hash was expected.
    fn finalize(
        &self,
        part_path: &Path,
        size: u64,
        expected_sha256: Option<&str>,
    ) -> Result<DownloadedFile> {
        let sha256 = hash_file(part_path)?;
        if let Some(expected) = expected_sha256 {
            if !expected.eq_ignore_ascii_case(&sha256) {
                let _ = std::fs::remove_file(part_path);
                return Err(Error::HashMismatch {
                    expected: expected.to_owned(),
                    actual: sha256,
                });
            }
        }
        Ok(DownloadedFile {
            path: part_path.to_path_buf(),
            sha256,
            size_bytes: size,
        })
    }
}

/// A verified, completed download still sitting in the downloads directory.
#[derive(Debug, Clone)]
pub struct DownloadedFile {
    /// Path of the completed `.part` file.
    pub path: PathBuf,
    /// sha256 hex of the content.
    pub sha256: String,
    /// Size in bytes.
    pub size_bytes: u64,
}

/// Stable per-(repo, revision, file) key for resume files.
fn job_key(repo: &str, revision: &str, filename: &str) -> String {
    let digest = Sha256::digest(format!("{repo}\n{revision}\n{filename}").as_bytes());
    let mut s = String::with_capacity(32);
    for b in &digest[..16] {
        use std::fmt::Write;
        let _ = write!(s, "{b:02x}");
    }
    s
}

/// Parse the total size out of `Content-Range: bytes 100-499/500`.
fn content_range_total(header: Option<&str>) -> Option<u64> {
    header?.rsplit('/').next()?.trim().parse().ok()
}

fn hash_file(path: &Path) -> Result<String> {
    crate::identity::full_sha256(path)
}

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

    #[test]
    fn debug_never_shows_the_token() {
        let client = HfClient::new(None, Some("hf_SUPER_SECRET_VALUE".into()));
        let dbg = format!("{client:?}");
        assert!(!dbg.contains("SUPER_SECRET"), "token leaked: {dbg}");
        assert!(dbg.contains("<redacted>"));
    }

    #[test]
    fn content_range_parsing() {
        assert_eq!(content_range_total(Some("bytes 100-499/500")), Some(500));
        assert_eq!(content_range_total(Some("bytes */1234")), Some(1234));
        assert_eq!(content_range_total(None), None);
    }

    #[test]
    fn job_key_is_stable_and_distinct() {
        let a = job_key("org/repo", "main", "f.gguf");
        assert_eq!(a, job_key("org/repo", "main", "f.gguf"));
        assert_ne!(a, job_key("org/repo", "main", "g.gguf"));
        assert_eq!(a.len(), 32);
    }
}