candding 0.1.0

use std::fs;
use std::path::Path;

use sha2::Digest;
use sha2::Sha256;
use tempfile::tempdir;

use super::ID;
use super::REPO;
use super::SHA;
use super::body;
use super::config;
use super::hub;
use super::server::Reply;
use super::server::Server;
use crate::error::CanddingError;
use crate::hub::HubConfig;
use crate::hub::download::fetch;
use crate::hub::download::verify::hex;
use crate::hub::resolve_hub;
use crate::registry::WeightsFormat;

/// `hub`, except that every GET of `file` answers `status` with `served`.
fn answering(
  file: &'static str,
  status: u16,
  served: String,
) -> impl Fn(&str, &str) -> Reply + Send + Sync + 'static {
  move |method, path| match method == "GET" && path.ends_with(&format!("/resolve/{SHA}/{file}")) {
    true => Reply::new(status, served.clone()),
    false => hub(method, path),
  }
}

fn gets(server: &Server, file: &str) -> usize {
  let request = format!("GET /{ID}/resolve/{SHA}/{file}");
  server.requests().iter().filter(|r| **r == request).count()
}

fn blob_names(cache: &Path) -> Vec<String> {
  let blobs = fs::read_dir(cache.join(REPO).join("blobs"));
  let names = blobs.into_iter().flatten().flatten();
  names
    .map(|entry| entry.file_name().to_string_lossy().into_owned())
    .collect()
}

fn a_blob_holds(cache: &Path, bytes: &str) -> bool {
  let blobs = cache.join(REPO).join("blobs");
  blob_names(cache)
    .iter()
    .any(|name| fs::read_to_string(blobs.join(name)).is_ok_and(|text| text == bytes))
}

/// No pointer, blob, partial file or ref was kept, and the next resolve goes back to the hub.
fn nothing_kept(cache: &Path, config: &HubConfig, server: &Server, file: &str, served: &str) {
  let pointer = cache.join(REPO).join("snapshots").join(SHA).join(file);
  assert!(fs::symlink_metadata(pointer).is_err(), "{file} stayed");
  assert!(!a_blob_holds(cache, served), "{file}'s body stayed");
  assert!(
    !blob_names(cache)
      .iter()
      .any(|name| name.ends_with(".incomplete"))
  );
  assert!(!cache.join(REPO).join("refs/main").exists());
  let before = server.requests().len();
  assert!(resolve_hub(config, ID, "main", WeightsFormat::Safetensors).is_err());
  assert!(
    server.requests().len() > before,
    "{file}: taken as complete"
  );
}

#[test]
fn a_non_success_get_is_an_error_and_never_a_file() {
  let served = "Entry not found here";
  for (file, status) in [("model.safetensors", 403), ("config.json", 404)] {
    let server = Server::start(answering(file, status, served.into()));
    let cache = tempdir().unwrap();
    let config = config(cache.path(), &server.url);
    match (
      status,
      fetch(&config, ID, "main", WeightsFormat::Safetensors),
    ) {
      (403, Err(CanddingError::GatedModel { id })) => assert_eq!(id, ID),
      (404, Err(CanddingError::DownloadFailed { id, reason })) => {
        assert_eq!(id, ID);
        assert!(reason.contains(file), "{reason}");
      }
      (_, other) => panic!("{file}: {other:?}"),
    }
    assert_eq!(gets(&server, file), 1);
    nothing_kept(cache.path(), &config, &server, file, served);
  }
}

#[test]
fn a_body_that_does_not_match_the_listing_is_refused_and_never_kept() {
  let model = body("model.safetensors");
  for (file, served, reason) in [
    (
      "config.json",
      body("config.json").to_uppercase(),
      "its git blob id is ",
    ),
    (
      "model.safetensors",
      "short".to_string(),
      "5 bytes arrived, not 26",
    ),
    (
      "model.safetensors",
      format!("{model}!"),
      "more than its 26 bytes arrived",
    ),
  ] {
    let server = Server::start(answering(file, 200, served.clone()));
    let cache = tempdir().unwrap();
    let config = config(cache.path(), &server.url);
    match fetch(&config, ID, "main", WeightsFormat::Safetensors) {
      Err(CanddingError::DownloadFailed { id, reason: text }) => {
        assert_eq!(id, ID);
        let expected = format!("`{file}` did not match the hub's listing: {reason}");
        assert!(text.contains(&expected), "{text}");
      }
      other => panic!("{file}: {other:?}"),
    }
    nothing_kept(cache.path(), &config, &server, file, &served);
  }
}

#[test]
fn a_wrong_file_under_the_right_blob_name_is_replaced() {
  let server = Server::start(hub);
  let cache = tempdir().unwrap();
  let model = body("model.safetensors");
  let blob = cache
    .path()
    .join(REPO)
    .join("blobs")
    .join(hex(&Sha256::digest(&model)));
  fs::create_dir_all(blob.parent().unwrap()).unwrap();
  fs::write(&blob, "stale").unwrap();
  let config = config(cache.path(), &server.url);
  let snapshot = fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
  assert_eq!(fs::read_to_string(&blob).unwrap(), model);
  let linked = fs::read_to_string(snapshot.join("model.safetensors")).unwrap();
  assert_eq!(linked, model);
  assert_eq!(gets(&server, "model.safetensors"), 1);
}

/// The state an older fetch could leave: every file linked to an ETag-named blob, the ref written, one blob damaged.
#[cfg(unix)]
#[test]
fn a_fetch_repairs_a_snapshot_whose_file_holds_an_error_body() {
  let server = Server::start(hub);
  let cache = tempdir().unwrap();
  let repo = cache.path().join(REPO);
  let snapshot = repo.join("snapshots").join(SHA);
  for file in ["1_Pooling/config.json", "config.json", "model.safetensors"] {
    let blob = repo
      .join("blobs")
      .join(format!("blob-{}", file.replace('/', "-")));
    let pointer = snapshot.join(file);
    fs::create_dir_all(blob.parent().unwrap()).unwrap();
    fs::create_dir_all(pointer.parent().unwrap()).unwrap();
    let damaged = file == "model.safetensors";
    fs::write(
      &blob,
      if damaged {
        "interrupted".into()
      } else {
        body(file)
      },
    )
    .unwrap();
    std::os::unix::fs::symlink(&blob, &pointer).unwrap();
  }
  fs::create_dir_all(repo.join("refs")).unwrap();
  fs::write(repo.join("refs/main"), SHA).unwrap();
  let config = config(cache.path(), &server.url);
  fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
  let model = "model.safetensors";
  assert_eq!(
    fs::read_to_string(snapshot.join(model)).unwrap(),
    body(model)
  );
  let name = hex(&Sha256::digest(body(model)));
  let target = fs::read_link(snapshot.join(model)).unwrap();
  assert_eq!(target, Path::new("../../blobs").join(name));
  let resolves = |server: &Server| {
    server
      .requests()
      .iter()
      .filter(|r| r.contains("/resolve/"))
      .count()
  };
  assert_eq!(resolves(&server), 2, "{:?}", server.requests());
  fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
  assert_eq!(
    resolves(&server),
    2,
    "a second fetch found every file intact"
  );
}

/// A copied pointer, as hf-hub writes on Windows, is replaced like a link.
#[test]
fn a_fetch_repairs_a_copied_file() {
  let server = Server::start(hub);
  let cache = tempdir().unwrap();
  let snapshot = cache.path().join(REPO).join("snapshots").join(SHA);
  fs::create_dir_all(snapshot.join("1_Pooling")).unwrap();
  for file in ["1_Pooling/config.json", "config.json"] {
    fs::write(snapshot.join(file), body(file)).unwrap();
  }
  fs::write(snapshot.join("model.safetensors"), "interrupted").unwrap();
  let config = config(cache.path(), &server.url);
  fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
  let model = "model.safetensors";
  assert_eq!(
    fs::read_to_string(snapshot.join(model)).unwrap(),
    body(model)
  );
  assert_eq!(gets(&server, "config.json"), 0);
  assert_eq!(gets(&server, model), 1);
}