oxinfer 0.0.3

//! Fetching and caching HuggingFace model repositories by name, over rsurl (pure-Rust HTTP).
//!
//! `--model google/gemma-4-E4B-it` resolves to `$OXINFER_CACHE/google--gemma-4-E4B-it` (default
//! `~/.cache/oxinfer/hub`), downloading the config, tokenizer, template and safetensors files that
//! are missing or incomplete. Downloads are resumable. Gated repos need `HF_TOKEN`.
//!
//! Files come from the ATOnline mirror when it carries the repo (`AI/Model:get`, which lists
//! every file with a download URL), otherwise from huggingface.co; `OXINFER_HUB` overrides.

use anyhow::{Context, Result, anyhow, bail};
use rsurl::{DownloadOptions, Request};
use serde::Deserialize;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::Instant;

pub fn cache_root() -> PathBuf {
    if let Ok(p) = std::env::var("OXINFER_CACHE") {
        return PathBuf::from(p);
    }
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".to_string());
    PathBuf::from(home).join(".cache").join("oxinfer").join("hub")
}

/// A local directory is used as is; otherwise `spec` must be a repo id (`org/name`), which is
/// fetched into the cache (if needed) and its directory returned.
pub fn resolve(spec: &str, cache_dir: Option<&Path>) -> Result<PathBuf> {
    let p = Path::new(spec);
    if p.is_dir() {
        return Ok(p.to_path_buf());
    }
    let parts: Vec<&str> = spec.split('/').collect();
    if parts.len() != 2 || parts.iter().any(|s| s.is_empty() || s.starts_with('.')) || spec.starts_with('/') {
        bail!("{spec}: not an existing directory and not a HuggingFace repo id of the form org/name");
    }
    let root = cache_dir.map(Path::to_path_buf).unwrap_or_else(cache_root);
    let dir = root.join(spec.replace('/', "--"));
    std::fs::create_dir_all(&dir).with_context(|| format!("create {}", dir.display()))?;
    fetch_repo(spec, &dir)?;
    Ok(dir)
}

fn token() -> Option<String> {
    std::env::var("HF_TOKEN")
        .ok()
        .or_else(|| std::env::var("HUGGING_FACE_HUB_TOKEN").ok())
        .filter(|t| !t.is_empty())
}

/// A GET request; `auth` adds the HuggingFace token (never sent to other hosts).
fn request(url: &str, auth: bool) -> Result<Request> {
    let mut r = Request::get(url)
        .map_err(|e| anyhow!("{url}: {e}"))?
        .follow_redirects(true)
        .max_redirs(10)
        .header("User-Agent", "oxinfer/0.1")
        // The resumable downloader stores raw bytes; refuse compressed transfer encodings.
        .header("Accept-Encoding", "identity");
    if auth && let Some(t) = token() {
        r = r.header("Authorization", &format!("Bearer {t}"));
    }
    Ok(r)
}

#[derive(Deserialize)]
struct Sibling {
    rfilename: String,
    #[serde(default)]
    size: Option<u64>,
}

#[derive(Deserialize)]
struct RepoInfo {
    #[serde(default)]
    siblings: Vec<Sibling>,
}

const SMALL_FILES: [&str; 8] = [
    "config.json",
    "generation_config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "chat_template.jinja",
    "processor_config.json",
    "preprocessor_config.json",
];

/// Where a repository file comes from.
enum Source {
    Url(String),
    /// Small files the mirror inlines as `data:` URLs.
    Inline(Vec<u8>),
}

/// A repository's file listing: name -> (size, source).
type Listing = HashMap<String, (Option<u64>, Source)>;

/// The ATOnline model mirror: `AI/Model:get` lists a repo's files with download URLs.
pub const MIRROR_API: &str = "https://ws.atonline.com/_special/rest/AI/Model:get";

#[derive(Deserialize)]
struct MirrorFile {
    #[serde(rename = "Filename")]
    filename: String,
    #[serde(rename = "Size")]
    size: String,
    #[serde(rename = "Url")]
    url: String,
}

#[derive(Deserialize)]
struct MirrorModel {
    #[serde(rename = "Files")]
    files: Vec<MirrorFile>,
}

#[derive(Deserialize)]
struct MirrorReply {
    data: MirrorModel,
}

/// Lists `repo` on the mirror; Ok(None) when the mirror does not carry it.
fn list_mirror(repo: &str) -> Result<Option<Listing>> {
    let api = format!("{MIRROR_API}?key={repo}");
    let resp = request(&api, false)?.send().map_err(|e| anyhow!("{api}: {e}"))?;
    match resp.status {
        200 => {}
        404 => return Ok(None),
        s => bail!("{api}: HTTP {s}"),
    }
    let reply: MirrorReply = serde_json::from_slice(&resp.body).with_context(|| format!("parse {api}"))?;
    let mut out = Listing::new();
    for f in reply.data.files {
        let src = match f.url.strip_prefix("data:") {
            Some(d) => {
                let b64 = d
                    .split_once(";base64,")
                    .map(|(_, b)| b)
                    .with_context(|| format!("{}: unsupported data URL", f.filename))?;
                Source::Inline(base64_decode(b64).with_context(|| format!("{}: bad base64", f.filename))?)
            }
            None => Source::Url(f.url),
        };
        out.insert(f.filename, (f.size.parse().ok(), src));
    }
    Ok(Some(out))
}

/// Lists `repo` on huggingface.co.
fn list_hf(repo: &str) -> Result<Listing> {
    let api = format!("https://huggingface.co/api/models/{repo}?blobs=true");
    let resp = request(&api, true)?.send().map_err(|e| anyhow!("{api}: {e}"))?;
    match resp.status {
        200 => {}
        401 | 403 => bail!(
            "{repo}: access denied (HTTP {}). If the repo is gated, accept its license on huggingface.co and set HF_TOKEN.",
            resp.status
        ),
        404 => bail!("{repo}: repository not found on huggingface.co"),
        s => bail!("{api}: HTTP {s}"),
    }
    let info: RepoInfo = serde_json::from_slice(&resp.body).with_context(|| format!("parse {api}"))?;
    Ok(info
        .siblings
        .into_iter()
        .map(|s| {
            let url = format!("https://huggingface.co/{repo}/resolve/main/{}", s.rfilename);
            (s.rfilename, (s.size, Source::Url(url)))
        })
        .collect())
}

/// Which hub to fetch from: `OXINFER_HUB=auto` (default: the mirror, then huggingface.co for
/// repos it does not carry), `mirror` or `hf`.
fn listing(repo: &str) -> Result<Listing> {
    let hub = std::env::var("OXINFER_HUB").unwrap_or_default();
    match hub.as_str() {
        "hf" => list_hf(repo),
        "mirror" => list_mirror(repo)?.with_context(|| format!("{repo}: not on the mirror ({MIRROR_API})")),
        "" | "auto" => match list_mirror(repo) {
            Ok(Some(l)) => Ok(l),
            Ok(None) => list_hf(repo),
            Err(e) => {
                eprintln!("[mirror unavailable ({e:#}); trying huggingface.co]");
                list_hf(repo)
            }
        },
        h => bail!("OXINFER_HUB={h}: expected auto, mirror or hf"),
    }
}

fn base64_decode(s: &str) -> Option<Vec<u8>> {
    let val = |c: u8| -> Option<u32> {
        Some(match c {
            b'A'..=b'Z' => c - b'A',
            b'a'..=b'z' => c - b'a' + 26,
            b'0'..=b'9' => c - b'0' + 52,
            b'+' | b'-' => 62,
            b'/' | b'_' => 63,
            _ => return None,
        } as u32)
    };
    let digits: Vec<u8> = s.bytes().filter(|c| !c.is_ascii_whitespace() && *c != b'=').collect();
    let mut out = Vec::with_capacity(digits.len() * 3 / 4);
    for chunk in digits.chunks(4) {
        let mut acc = 0u32;
        for &c in chunk {
            acc = acc << 6 | val(c)?;
        }
        acc <<= 6 * (4 - chunk.len() as u32);
        let n = match chunk.len() {
            4 => 3,
            3 => 2,
            2 => 1,
            _ => return None,
        };
        out.extend_from_slice(&acc.to_be_bytes()[1..1 + n]);
    }
    Some(out)
}

fn fetch_repo(repo: &str, dir: &Path) -> Result<()> {
    let list = listing(repo)?;
    let mut files: Vec<String> = SMALL_FILES
        .iter()
        .filter(|f| list.contains_key(**f))
        .map(|f| f.to_string())
        .collect();
    if list.contains_key("model.safetensors") {
        files.push("model.safetensors".into());
    } else if let Some((size, src)) = list.get("model.safetensors.index.json") {
        fetch_file(repo, dir, "model.safetensors.index.json", *size, src, 0, 0)?;
        let idx: serde_json::Value = serde_json::from_slice(&std::fs::read(dir.join("model.safetensors.index.json"))?)?;
        let mut shards: Vec<String> = idx["weight_map"]
            .as_object()
            .context("weight_map")?
            .values()
            .filter_map(|v| v.as_str().map(str::to_string))
            .collect();
        shards.sort();
        shards.dedup();
        files.extend(shards);
    } else {
        bail!("{repo}: no safetensors weights found");
    }
    let count = files.len();
    for (i, f) in files.iter().enumerate() {
        let (size, src) = list.get(f).with_context(|| format!("{repo}: {f} is not in the listing"))?;
        fetch_file(repo, dir, f, *size, src, i + 1, count)?;
    }
    Ok(())
}

fn human_gb(b: u64) -> String {
    if b >= 1_000_000_000 {
        format!("{:.2} GB", b as f64 / 1e9)
    } else {
        format!("{:.1} MB", b as f64 / 1e6)
    }
}

/// Renders a progress line: `[=====>     ]  42.0%  3.21/7.64 GB  31 MB/s  eta 2m10s`.
fn progress_line(label: &str, done: u64, total: Option<u64>, secs: f64) -> String {
    let rate = done as f64 / secs.max(1e-3);
    let speed = format!("{:5.0} MB/s", rate / 1e6);
    match total {
        Some(t) if t > 0 => {
            let frac = (done as f64 / t as f64).clamp(0.0, 1.0);
            let width = 30usize;
            let filled = (frac * width as f64) as usize;
            let bar: String = (0..width)
                .map(|i| {
                    if i < filled {
                        '='
                    } else if i == filled {
                        '>'
                    } else {
                        ' '
                    }
                })
                .collect();
            let eta = if rate > 0.0 { ((t - done) as f64 / rate) as u64 } else { 0 };
            // The rate is noise for the first second; show no estimate yet.
            let eta_s = if secs < 1.0 {
                "--".to_string()
            } else if eta >= 3600 {
                format!("{}h{:02}m", eta / 3600, (eta % 3600) / 60)
            } else if eta >= 60 {
                format!("{}m{:02}s", eta / 60, eta % 60)
            } else {
                format!("{eta}s")
            };
            format!(
                "  {label} [{bar}] {:5.1}%  {}/{}  {speed}  eta {eta_s}",
                frac * 100.0,
                human_gb(done),
                human_gb(t)
            )
        }
        _ => format!("  {label} {}  {speed}", human_gb(done)),
    }
}

fn fetch_file(repo: &str, dir: &Path, name: &str, expected: Option<u64>, src: &Source, index: usize, count: usize) -> Result<()> {
    use std::io::IsTerminal;
    let path = dir.join(name);
    if let Ok(md) = std::fs::metadata(&path)
        && md.len() > 0
        && expected.is_none_or(|s| s == md.len())
    {
        return Ok(());
    }
    let url = match src {
        Source::Inline(b) => {
            std::fs::write(&path, b).with_context(|| format!("write {}", path.display()))?;
            return Ok(());
        }
        Source::Url(u) => u.as_str(),
    };
    let t0 = Instant::now();
    let label = format!("[{index}/{count}] {name}");
    let tty = std::io::stderr().is_terminal();
    let mut last = Instant::now();
    let mut last_pct = -1i64;
    let plabel = label.clone();
    let opts = DownloadOptions {
        progress: Some(Box::new(move |done: u64, total: Option<u64>| {
            let secs = t0.elapsed().as_secs_f64();
            if tty {
                if last.elapsed().as_millis() >= 200 {
                    last = Instant::now();
                    eprint!("\r\x1b[K{}", progress_line(&plabel, done, total, secs));
                }
            } else {
                // Not a terminal: one log line per 5% (or every 30 s when the size is unknown).
                let pct = total.filter(|&t| t > 0).map(|t| (done * 100 / t) as i64 / 5 * 5);
                let due = match pct {
                    Some(p) => p != last_pct,
                    None => last.elapsed().as_secs() >= 30,
                };
                if due {
                    last = Instant::now();
                    last_pct = pct.unwrap_or(-1);
                    eprintln!("{}", progress_line(&plabel, done, total, secs));
                }
            }
        })),
        ..Default::default()
    };
    eprintln!(
        "[fetching {repo}/{name}{}]",
        expected.map(|s| format!(" ({})", human_gb(s))).unwrap_or_default()
    );
    let out = request(url, url.starts_with("https://huggingface.co/"))?
        .download_resumable(&path, opts)
        .map_err(|e| anyhow!("download {repo}/{name}: {e}"))?;
    let secs = t0.elapsed().as_secs_f64().max(1e-3);
    if tty {
        eprint!("\r\x1b[K");
    }
    eprintln!(
        "  {label}: done, {} in {:.0}s ({:.0} MB/s){}",
        human_gb(out.bytes_written + out.resumed_from),
        secs,
        out.bytes_written as f64 / secs / 1e6,
        if out.resumed_from > 0 {
            format!(", resumed from {}", human_gb(out.resumed_from))
        } else {
            String::new()
        }
    );
    if let (Some(exp), Ok(md)) = (expected, std::fs::metadata(&path))
        && md.len() != exp
    {
        bail!("{}: size {} after download, expected {exp}", path.display(), md.len());
    }
    Ok(())
}