candding 0.1.0

//! One listed file into the cache: streamed into a partial file under `blobs/` while hashed, renamed onto
//! `blobs/<hash>` and linked from the snapshot by a relative path, the layout huggingface_hub writes.

use std::fs;
use std::fs::File;
use std::io;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;

use futures_util::StreamExt;
use hf_hub::HFError;
use hf_hub::HFRepository;
use hf_hub::RepoTypeModel;
use sha1::Digest;
use sha1::Sha1;
use sha2::Sha256;

use super::listing::ListedFile;
use super::verify::hex;
use crate::error::CanddingError;
use crate::error::Result;

/// The repository lock keeps a second candding fetch from writing this file at the same time.
const PARTIAL: &str = "candding.incomplete";

/// Streams `file` at `commit` into the cache and links `pointer` to its blob; nothing is kept when it does not match.
pub(super) async fn store(
  repo: &HFRepository<RepoTypeModel>,
  file: &ListedFile,
  commit: &str,
  repo_dir: &Path,
  pointer: &Path,
  failed: impl Fn(HFError) -> CanddingError,
) -> Result<()> {
  let mismatch = |reason: String| CanddingError::DownloadFailed {
    id: repo.repo_path(),
    reason: format!(
      "`{}` did not match the hub's listing: {reason}; nothing was kept",
      file.path
    ),
  };
  let (_, mut stream) = repo
    .download_file_stream()
    .filename(file.path.as_str())
    .revision(commit)
    .send()
    .await
    .map_err(&failed)?;
  let blobs = repo_dir.join("blobs");
  fs::create_dir_all(&blobs)?;
  let mut partial = Partial::create(blobs.join(PARTIAL))?;
  let mut hashes = Hashes::new(file.size);
  while let Some(chunk) = stream.next().await {
    let chunk = chunk.map_err(&failed)?;
    hashes.update(&chunk);
    if hashes.len > file.size {
      return Err(mismatch(format!(
        "more than its {} bytes arrived",
        file.size
      )));
    }
    partial.write(&chunk)?;
  }
  let name = hashes.name(file).map_err(mismatch)?;
  partial.persist(&blobs.join(&name))?;
  Ok(link(pointer, &name, &file.path)?)
}

/// The content's SHA-256 and git's blob SHA-1, which hashes `blob <size>\0` first; the listed size stands in until
/// the length is checked.
struct Hashes {
  sha256: Sha256,
  git: Sha1,
  len: u64,
}

impl Hashes {
  fn new(size: u64) -> Self {
    let mut git = Sha1::new();
    git.update(format!("blob {size}\0"));
    Self {
      sha256: Sha256::new(),
      git,
      len: 0,
    }
  }

  fn update(&mut self, bytes: &[u8]) {
    self.sha256.update(bytes);
    self.git.update(bytes);
    self.len += bytes.len() as u64;
  }

  /// huggingface_hub names a blob by the hub's ETag: the git blob id of a file git stores, which has to be the
  /// listed one, and the SHA-256 of an LFS or xet file.
  fn name(self, file: &ListedFile) -> std::result::Result<String, String> {
    if self.len != file.size {
      return Err(format!("{} bytes arrived, not {}", self.len, file.size));
    }
    if !file.in_git {
      return Ok(hex(&self.sha256.finalize()));
    }
    let id = hex(&self.git.finalize());
    match id.eq_ignore_ascii_case(&file.oid) {
      true => Ok(id),
      false => Err(format!("its git blob id is {id}, not {}", file.oid)),
    }
  }
}

/// A file that is removed on drop unless it was renamed onto its blob first.
struct Partial {
  path: PathBuf,
  file: Option<File>,
}

impl Partial {
  fn create(path: PathBuf) -> io::Result<Self> {
    let file = File::create(&path)?;
    Ok(Self {
      path,
      file: Some(file),
    })
  }

  fn write(&mut self, bytes: &[u8]) -> io::Result<()> {
    match &mut self.file {
      Some(file) => file.write_all(bytes),
      None => Err(io::Error::other("the partial file is closed")),
    }
  }

  /// Closed first: Windows renames no open file.
  fn persist(mut self, blob: &Path) -> io::Result<()> {
    self.file.take();
    fs::rename(&self.path, blob)
  }
}

impl Drop for Partial {
  fn drop(&mut self) {
    let _ = fs::remove_file(&self.path);
  }
}

/// `listed` passed `listing::problem`, so its `/` count is its depth below the snapshot directory.
fn link(pointer: &Path, name: &str, listed: &str) -> io::Result<()> {
  if let Some(parent) = pointer.parent() {
    fs::create_dir_all(parent)?;
  }
  let mut target = PathBuf::new();
  for _ in 0..listed.matches('/').count() + 2 {
    target.push("..");
  }
  target.push("blobs");
  target.push(name);
  symlink(&target, pointer)
}

#[cfg(unix)]
fn symlink(target: &Path, pointer: &Path) -> io::Result<()> {
  std::os::unix::fs::symlink(target, pointer)
}

/// A symbolic link needs a privilege on Windows, so the pointer is a copy there, as hf-hub makes it.
#[cfg(not(unix))]
fn symlink(target: &Path, pointer: &Path) -> io::Result<()> {
  let blob = pointer.parent().map(|parent| parent.join(target));
  fs::copy(blob.unwrap_or_default(), pointer).map(drop)
}