zc2 0.0.25

//! `zc dataset get` — fetch a marketplace dataset's files.
//!
//! The models side of the marketplace has had `zc infer` for a while; datasets
//! had nothing, so the hub's dataset page printed a `curl` invocation with a
//! bearer token in it. This is the verb that page can point at instead.
//!
//! A noun namespace (`zc dataset <verb>`) rather than a bare top-level verb:
//! it leaves room for `zc dataset ls`/`info` without spending more top-level
//! words, and `zc pull` already means "pull Docker images".
//!
//! Public datasets only, deliberately. The marketplace's `auth/deps.py`
//! accepts a browser session and nothing else — a `zc login` CLI token is
//! refused there on purpose — so there is no credential this command could
//! send that would unlock a private dataset. Rather than pretend otherwise, a
//! 404 says what is actually true.

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

use crate::model_uri::{address_to_ref, parse_model_address, ModelAddress};

/// One file inside a published version.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatasetFile {
    pub path: String,
    pub sha256: String,
    pub size_bytes: u64,
}

/// What a reference resolved to: concrete bytes to fetch.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Resolved {
    pub dataset_id: String,
    pub digest: String,
    pub name: String,
    pub files: Vec<DatasetFile>,
}

/// How a `zc://` reference reaches a dataset id.
///
/// Two shapes, and they are **not** interchangeable: the marketplace's
/// `_REF_RE` requires an `owner/name` pair, so handing it the bare-uuid form
/// (which the hub emits for an owner with no handle) answers 422, not a
/// dataset. They have to take different routes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Lookup {
    /// `GET /api/datasets/resolve?ref=zc://owner/name[@sha256:…]`
    Resolve(String),
    /// `GET /api/datasets/{id}` — the reference already names the id.
    Direct(String),
}

/// Decide how to look a reference up, or explain why it is not one.
pub fn plan_lookup(reference: &str) -> Result<Lookup, String> {
    match parse_model_address(reference) {
        Some(ModelAddress::Uuid(id)) => Ok(Lookup::Direct(id)),
        Some(addr @ ModelAddress::Named { .. }) => Ok(Lookup::Resolve(address_to_ref(&addr))),
        None => Err(format!(
            "not a dataset reference: {reference:?}\n\
             Expected zc://<owner>/<name> or zc://<uuid>."
        )),
    }
}

/// Parsed `zc dataset get` arguments.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GetArgs {
    pub reference: String,
    pub out: Option<String>,
    pub api_url: Option<String>,
}

pub fn parse_get_args(args: &[String]) -> Result<GetArgs, String> {
    let mut reference: Option<String> = None;
    let mut out: Option<String> = None;
    let mut api_url: Option<String> = None;
    let mut it = args.iter();
    while let Some(a) = it.next() {
        match a.as_str() {
            "-o" | "--out" => {
                out = Some(
                    it.next()
                        .cloned()
                        .ok_or_else(|| "-o needs a directory".to_string())?,
                )
            }
            "--api-url" => {
                api_url = Some(
                    it.next()
                        .cloned()
                        .ok_or_else(|| "--api-url needs a URL".to_string())?,
                )
            }
            other if other.starts_with('-') => {
                return Err(format!("unknown flag: {other}"));
            }
            other => {
                if reference.is_some() {
                    return Err(format!("unexpected extra argument: {other}"));
                }
                reference = Some(other.to_string());
            }
        }
    }
    match reference {
        Some(reference) => Ok(GetArgs {
            reference,
            out,
            api_url,
        }),
        None => Err("a dataset reference is required".to_string()),
    }
}

/// Hex sha256 of `bytes`, for verifying a download against the API's manifest.
pub fn sha256_hex(bytes: &[u8]) -> String {
    use ring::digest::{digest, SHA256};
    digest(&SHA256, bytes)
        .as_ref()
        .iter()
        .map(|b| format!("{b:02x}"))
        .collect()
}

/// Whether a downloaded file matches the hash the catalogue published.
///
/// The API hands us `sha256` per file for free; not checking it would be
/// choosing not to notice a truncated or corrupted download.
pub fn verify(file: &DatasetFile, bytes: &[u8]) -> Result<(), String> {
    let actual = sha256_hex(bytes);
    if actual == file.sha256.to_ascii_lowercase() {
        return Ok(());
    }
    Err(format!(
        "{}: sha256 mismatch (expected {}, got {})",
        file.path, file.sha256, actual
    ))
}

/// Where a file lands, given the output directory.
///
/// A dataset path is server-controlled, so `..` and absolute paths are treated
/// as hostile: every component that could climb out of `dir` is dropped. A
/// registry entry must not be able to write over `~/.ssh/authorized_keys`.
pub fn safe_destination(dir: &Path, path: &str) -> Option<PathBuf> {
    let mut out = dir.to_path_buf();
    let mut wrote = false;
    for part in path.split('/') {
        if part.is_empty() || part == "." || part == ".." {
            continue;
        }
        if part.contains('\\') || part.contains(':') {
            return None;
        }
        out.push(part);
        wrote = true;
    }
    if wrote {
        Some(out)
    } else {
        None
    }
}

fn api_base(explicit: Option<&str>) -> String {
    explicit
        .map(|s| s.to_string())
        .unwrap_or_else(crate::credentials::default_api_url)
        .trim_end_matches('/')
        .to_string()
}

fn get_json(url: &str) -> Result<(u16, String), String> {
    match ureq::get(url)
        .config()
        .http_status_as_error(false)
        .build()
        .call()
    {
        Ok(r) => {
            let status = r.status().as_u16();
            let body = r.into_body().read_to_string().unwrap_or_default();
            Ok((status, body))
        }
        Err(e) => Err(e.to_string()),
    }
}

fn get_bytes(url: &str) -> Result<(u16, Vec<u8>), String> {
    match ureq::get(url)
        .config()
        .http_status_as_error(false)
        .build()
        .call()
    {
        Ok(r) => {
            let status = r.status().as_u16();
            let mut buf = Vec::new();
            r.into_body()
                .into_reader()
                .read_to_end(&mut buf)
                .map_err(|e| e.to_string())?;
            Ok((status, buf))
        }
        Err(e) => Err(e.to_string()),
    }
}

/// The message a 404 gets. Private datasets are the likely cause and the CLI
/// genuinely cannot reach them, so say that rather than leaving the reader to
/// wonder whether they typed the name wrong.
fn not_found(reference: &str) -> String {
    format!(
        "No public dataset at {reference}.\n\
         It may not exist, or it may be private — `zc` can only fetch\n\
         public datasets. Download a private one from the hub in a browser."
    )
}

/// Read `files` out of a dataset-version JSON body.
pub fn parse_files(body: &str) -> Result<Vec<DatasetFile>, String> {
    let v: serde_json::Value = serde_json::from_str(body).map_err(|e| e.to_string())?;
    let files = v
        .get("files")
        .and_then(|f| f.as_array())
        .ok_or_else(|| "response carried no files".to_string())?;
    Ok(files
        .iter()
        .filter_map(|f| {
            Some(DatasetFile {
                path: f.get("path")?.as_str()?.to_string(),
                sha256: f.get("sha256")?.as_str()?.to_string(),
                size_bytes: f.get("size_bytes").and_then(|s| s.as_u64()).unwrap_or(0),
            })
        })
        .collect())
}

fn resolve(base: &str, reference: &str) -> Result<Resolved, String> {
    match plan_lookup(reference)? {
        Lookup::Resolve(canonical) => {
            let url = format!("{base}/api/datasets/resolve?ref={}", urlencode(&canonical));
            let (status, body) = get_json(&url)?;
            if status == 404 || status == 422 {
                return Err(not_found(reference));
            }
            if status != 200 {
                return Err(format!("resolve failed: HTTP {status}"));
            }
            let v: serde_json::Value = serde_json::from_str(&body).map_err(|e| e.to_string())?;
            let dataset_id = v
                .get("dataset_id")
                .and_then(|s| s.as_str())
                .ok_or("resolve returned no dataset_id")?
                .to_string();
            let digest = v
                .get("digest")
                .and_then(|s| s.as_str())
                .ok_or("resolve returned no digest")?
                .to_string();
            let name = v
                .get("name")
                .and_then(|s| s.as_str())
                .unwrap_or("dataset")
                .to_string();
            let files = version_files(base, &dataset_id, &digest)?;
            Ok(Resolved {
                dataset_id,
                digest,
                name,
                files,
            })
        }
        Lookup::Direct(id) => {
            let (status, body) = get_json(&format!("{base}/api/datasets/{id}"))?;
            if status == 404 {
                return Err(not_found(reference));
            }
            if status != 200 {
                return Err(format!("lookup failed: HTTP {status}"));
            }
            let v: serde_json::Value = serde_json::from_str(&body).map_err(|e| e.to_string())?;
            let latest = v
                .get("latest_version")
                .filter(|l| !l.is_null())
                .ok_or_else(|| format!("dataset {id} has no published version"))?;
            let digest = latest
                .get("digest")
                .and_then(|s| s.as_str())
                .ok_or("version carried no digest")?
                .to_string();
            let name = v
                .get("name")
                .and_then(|s| s.as_str())
                .unwrap_or("dataset")
                .to_string();
            Ok(Resolved {
                dataset_id: id,
                digest,
                name,
                files: parse_files(&latest.to_string())?,
            })
        }
    }
}

fn version_files(base: &str, id: &str, digest: &str) -> Result<Vec<DatasetFile>, String> {
    let (status, body) = get_json(&format!("{base}/api/datasets/{id}/versions/{digest}"))?;
    if status != 200 {
        return Err(format!("could not list files: HTTP {status}"));
    }
    parse_files(&body)
}

/// Percent-encode the few characters a `zc://` ref can carry that a query
/// string must not. Not a general encoder: handles and dataset names are
/// already URL-safe slugs, so this covers `:` and `/` and leaves the rest.
pub fn urlencode(s: &str) -> String {
    s.bytes()
        .map(|b| match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                (b as char).to_string()
            }
            _ => format!("%{b:02X}"),
        })
        .collect()
}

/// `zc dataset get <ref> [-o DIR] [--api-url URL]`
pub fn run_get(args: &[String]) -> i32 {
    let parsed = match parse_get_args(args) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("{e}");
            usage();
            return 2;
        }
    };
    let base = api_base(parsed.api_url.as_deref());
    let resolved = match resolve(&base, &parsed.reference) {
        Ok(r) => r,
        Err(e) => {
            eprintln!("{e}");
            return 1;
        }
    };
    if resolved.files.is_empty() {
        eprintln!("{} has no files to download.", parsed.reference);
        return 1;
    }

    let dir = PathBuf::from(parsed.out.unwrap_or_else(|| resolved.name.clone()));
    if let Err(e) = fs::create_dir_all(&dir) {
        eprintln!("cannot create {}: {e}", dir.display());
        return 1;
    }

    let mut failed = 0;
    for file in &resolved.files {
        let url = format!(
            "{base}/api/datasets/{}/versions/{}/files/{}",
            resolved.dataset_id,
            resolved.digest,
            file.path
                .split('/')
                .map(urlencode)
                .collect::<Vec<_>>()
                .join("/")
        );
        let bytes = match get_bytes(&url) {
            Ok((200, b)) => b,
            Ok((status, _)) => {
                eprintln!("{}: HTTP {status}", file.path);
                failed += 1;
                continue;
            }
            Err(e) => {
                // The presigned redirect points at object storage, which is a
                // different host from the API and may not be reachable from
                // here even when the API is.
                eprintln!("{}: {e}", file.path);
                failed += 1;
                continue;
            }
        };
        if let Err(e) = verify(file, &bytes) {
            eprintln!("{e}");
            failed += 1;
            continue;
        }
        let Some(dest) = safe_destination(&dir, &file.path) else {
            eprintln!("{}: unsafe path, skipped", file.path);
            failed += 1;
            continue;
        };
        if let Some(parent) = dest.parent() {
            if let Err(e) = fs::create_dir_all(parent) {
                eprintln!("{}: {e}", file.path);
                failed += 1;
                continue;
            }
        }
        if let Err(e) = fs::write(&dest, &bytes) {
            eprintln!("{}: {e}", file.path);
            failed += 1;
            continue;
        }
        println!("{} ({} bytes)", dest.display(), bytes.len());
    }

    if failed > 0 {
        eprintln!("{failed} file(s) failed.");
        return 1;
    }
    println!("Downloaded {} to {}", resolved.name, dir.display());
    0
}

fn usage() {
    eprintln!(
        "Usage: zc dataset get <zc://owner/name> [-o DIR] [--api-url URL]\n\
         \n\
         Examples:\n\
         \x20 zc dataset get zc://alice/sentiment-mini\n\
         \x20 zc dataset get zc://alice/sentiment-mini -o ./data\n\
         \n\
         Public datasets only. Private ones download from the hub in a browser."
    );
}

/// Entry point for the `dataset` noun.
pub fn run(args: &[String]) -> i32 {
    match args.first().map(|s| s.as_str()) {
        Some("get") => run_get(&args[1..]),
        Some("-h") | Some("--help") | None => {
            usage();
            0
        }
        Some(other) => {
            eprintln!("Unknown dataset command: {other}");
            usage();
            2
        }
    }
}

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

    #[test]
    fn a_named_reference_resolves_through_the_resolve_route() {
        assert_eq!(
            plan_lookup("zc://alice/sentiment-mini").unwrap(),
            Lookup::Resolve("zc://alice/sentiment-mini".to_string())
        );
    }

    #[test]
    fn a_pinned_digest_survives_into_the_resolve_ref() {
        let d = "a".repeat(64);
        assert_eq!(
            plan_lookup(&format!("zc://alice/mini@sha256:{d}")).unwrap(),
            Lookup::Resolve(format!("zc://alice/mini@sha256:{d}"))
        );
    }

    #[test]
    fn a_uuid_reference_never_goes_to_resolve() {
        // The marketplace's `_REF_RE` demands owner/name and answers 422 for
        // this shape, so routing it to /resolve would break exactly the
        // datasets whose owner has no handle.
        let id = "e629e662-1d32-4e52-88e9-b0e83416c852";
        assert_eq!(
            plan_lookup(&format!("zc://{id}")).unwrap(),
            Lookup::Direct(id.to_string())
        );
    }

    #[test]
    fn a_bare_uuid_is_accepted_too() {
        let id = "e629e662-1d32-4e52-88e9-b0e83416c852";
        assert_eq!(plan_lookup(id).unwrap(), Lookup::Direct(id.to_string()));
    }

    #[test]
    fn nonsense_is_rejected_with_the_expected_shapes_named() {
        let err = plan_lookup("https://example.com/x").unwrap_err();
        assert!(err.contains("zc://<owner>/<name>"), "{err}");
    }

    #[test]
    fn parses_a_reference_with_an_output_directory() {
        let args: Vec<String> = ["zc://a/b", "-o", "./data"]
            .iter()
            .map(|s| s.to_string())
            .collect();
        let got = parse_get_args(&args).unwrap();
        assert_eq!(got.reference, "zc://a/b");
        assert_eq!(got.out.as_deref(), Some("./data"));
    }

    #[test]
    fn an_api_url_flag_overrides_the_default() {
        let args: Vec<String> = ["zc://a/b", "--api-url", "http://localhost:8000"]
            .iter()
            .map(|s| s.to_string())
            .collect();
        assert_eq!(
            parse_get_args(&args).unwrap().api_url.as_deref(),
            Some("http://localhost:8000")
        );
    }

    #[test]
    fn a_missing_reference_is_an_error_not_a_download_of_nothing() {
        assert!(parse_get_args(&[]).is_err());
    }

    #[test]
    fn a_dangling_output_flag_is_an_error() {
        let args = vec!["zc://a/b".to_string(), "-o".to_string()];
        assert!(parse_get_args(&args).is_err());
    }

    #[test]
    fn an_unknown_flag_is_refused_rather_than_ignored() {
        let args = vec!["zc://a/b".to_string(), "--recursive".to_string()];
        assert!(parse_get_args(&args).is_err());
    }

    #[test]
    fn sha256_matches_a_known_vector() {
        assert_eq!(
            sha256_hex(b"abc"),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );
    }

    #[test]
    fn a_matching_hash_verifies() {
        let f = DatasetFile {
            path: "d.csv".into(),
            sha256: sha256_hex(b"hello"),
            size_bytes: 5,
        };
        assert!(verify(&f, b"hello").is_ok());
    }

    #[test]
    fn a_corrupted_download_is_rejected() {
        let f = DatasetFile {
            path: "d.csv".into(),
            sha256: sha256_hex(b"hello"),
            size_bytes: 5,
        };
        let err = verify(&f, b"hello!").unwrap_err();
        assert!(err.contains("sha256 mismatch"), "{err}");
    }

    #[test]
    fn an_uppercase_published_hash_still_matches() {
        let f = DatasetFile {
            path: "d.csv".into(),
            sha256: sha256_hex(b"hello").to_uppercase(),
            size_bytes: 5,
        };
        assert!(verify(&f, b"hello").is_ok());
    }

    #[test]
    fn a_nested_path_lands_under_the_output_directory() {
        let dir = Path::new("/tmp/out");
        assert_eq!(
            safe_destination(dir, "train/part-0.csv").unwrap(),
            PathBuf::from("/tmp/out/train/part-0.csv")
        );
    }

    #[test]
    fn a_traversing_path_cannot_climb_out() {
        // A registry entry must not be able to write outside the target dir.
        let dir = Path::new("/tmp/out");
        assert_eq!(
            safe_destination(dir, "../../etc/passwd").unwrap(),
            PathBuf::from("/tmp/out/etc/passwd")
        );
    }

    #[test]
    fn an_absolute_path_is_reparented_not_honoured() {
        let dir = Path::new("/tmp/out");
        assert_eq!(
            safe_destination(dir, "/etc/passwd").unwrap(),
            PathBuf::from("/tmp/out/etc/passwd")
        );
    }

    #[test]
    fn a_path_that_is_only_dots_yields_nothing_to_write() {
        assert!(safe_destination(Path::new("/tmp/out"), "../..").is_none());
    }

    #[test]
    fn files_are_read_out_of_a_version_body() {
        let body = r#"{"files":[{"path":"d.csv","sha256":"ab","size_bytes":12}]}"#;
        assert_eq!(
            parse_files(body).unwrap(),
            vec![DatasetFile {
                path: "d.csv".into(),
                sha256: "ab".into(),
                size_bytes: 12
            }]
        );
    }

    #[test]
    fn a_body_without_files_is_an_error() {
        assert!(parse_files(r#"{"id":"x"}"#).is_err());
    }

    #[test]
    fn the_reference_is_encoded_for_a_query_string() {
        assert_eq!(urlencode("zc://alice/mini"), "zc%3A%2F%2Falice%2Fmini");
    }

    #[test]
    fn the_404_message_says_private_datasets_are_out_of_reach() {
        let msg = not_found("zc://a/b");
        assert!(msg.contains("private"), "{msg}");
        assert!(msg.contains("browser"), "{msg}");
    }

    #[test]
    fn an_unknown_subcommand_is_refused() {
        assert_eq!(run(&["frobnicate".to_string()]), 2);
    }

    #[test]
    fn help_is_not_an_error() {
        assert_eq!(run(&["--help".to_string()]), 0);
    }
}