liboxen 0.50.6

use crate::api;
use crate::api::client;
use crate::constants::{max_retries, stream_segment_size};
use crate::error::OxenError;
use crate::model::{CommitEntry, LocalRepository, MerkleHash, RemoteRepository};
use crate::util::{self, concurrency};
use crate::view::versions::{
    CleanCorruptedVersionsResponse, CompleteVersionUploadRequest, CompletedFileUpload,
    CreateVersionUploadRequest, MultipartLargeFileUpload, MultipartLargeFileUploadStatus,
    VersionFile, VersionFileResponse,
};
use crate::view::{ErrorFileInfo, ErrorFilesResponse};

use crate::core::progress::push_progress::PushProgress;
use async_compression::tokio::bufread::GzipDecoder;
use flate2::Compression;
use flate2::write::GzEncoder;
use futures_util::StreamExt;
use futures_util::stream::FuturesUnordered;
use http::Method;
use http::header::CONTENT_LENGTH;
use rand::{Rng, thread_rng};
use tokio_tar::Archive;
use tokio_util::codec::{BytesCodec, FramedRead};

use std::collections::{HashMap, HashSet};
use std::io::{SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tokio::fs::OpenOptions;
use tokio::io::{AsyncReadExt, AsyncSeekExt};
use tokio::sync::Semaphore;
use tokio::time::sleep;

// Multipart upload strategy, based off of AWS S3 Multipart Upload and huggingface hf_transfer
// https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html
// https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs#L104
const BASE_WAIT_TIME: usize = 300;
const MAX_WAIT_TIME: usize = 10_000;
const PARALLEL_FAILURES: usize = 63;

/// Get the size of a version
pub async fn get(
    repository: &RemoteRepository,
    version_id: MerkleHash,
) -> Result<Option<VersionFile>, OxenError> {
    let uri = format!("/versions/{version_id}/metadata");
    let url = api::endpoint::url_from_repo(repository, &uri)?;
    log::debug!("api::client::versions::get {url}");

    let client = client::new_for_url(&url)?;
    let res = client.get(&url).send().await?;
    if res.status() == 404 {
        return Ok(None);
    }

    let body = client::parse_json_body(&url, res).await?;
    let response: Result<VersionFileResponse, serde_json::Error> = serde_json::from_str(&body);
    match response {
        Ok(version_file) => Ok(Some(version_file.version)),
        Err(err) => Err(OxenError::basic_str(format!(
            "api::client::versions::get() Could not deserialize response [{err}]\n{body}"
        ))),
    }
}

pub async fn clean(
    remote_repo: &RemoteRepository,
) -> Result<CleanCorruptedVersionsResponse, OxenError> {
    let uri = "/versions";
    let url = api::endpoint::url_from_repo(remote_repo, uri)?;
    log::debug!("api::client::versions::clean {url}");

    let client = client::new_for_url(&url)?;
    let res = client.delete(&url).send().await?;
    let body = client::parse_json_body(&url, res).await?;
    let response: Result<CleanCorruptedVersionsResponse, serde_json::Error> =
        serde_json::from_str(&body);
    match response {
        Ok(response) => Ok(response),
        Err(err) => Err(OxenError::basic_str(format!(
            "api::client::versions::clean() Could not deserialize response [{err}]\n{body}"
        ))),
    }
}

/// Uploads a large file to the server in parallel and unpacks it in the versions directory
/// Returns the `MultipartLargeFileUpload` struct for the created upload
pub async fn parallel_large_file_upload(
    remote_repo: &RemoteRepository,
    file_path: impl AsRef<Path>,
    dst_dir: Option<impl AsRef<Path>>, // dst_dir is provided for workspace add workflow
    workspace_id: Option<String>,
    commit_entry: Option<CommitEntry>, // entry is provided for push workflow
    progress: Option<&Arc<PushProgress>>, // for push workflow
) -> Result<MultipartLargeFileUpload, OxenError> {
    log::debug!("multipart_large_file_upload path: {:?}", file_path.as_ref());

    let mut upload =
        create_multipart_large_file_upload(remote_repo, file_path, dst_dir, commit_entry).await?;

    log::debug!("multipart_large_file_upload upload: {:?}", upload.hash);

    let max_retries = max_retries();
    let results = upload_chunks(
        remote_repo,
        &mut upload,
        stream_segment_size(),
        PARALLEL_FAILURES,
        max_retries,
        progress,
    )
    .await?;
    let num_chunks = results.len();
    log::debug!("multipart_large_file_upload num_chunks: {num_chunks:?}");
    complete_multipart_large_file_upload(remote_repo, upload, num_chunks, workspace_id).await
}

/// Creates a new multipart large file upload
/// Will reject the upload if the hash already exists on the server.
/// The rejection helps prevent duplicate uploads or parallel uploads of the same file.
/// Returns the `MultipartLargeFileUpload` struct for the created upload
async fn create_multipart_large_file_upload(
    remote_repo: &RemoteRepository,
    file_path: impl AsRef<Path>,
    dst_dir: Option<impl AsRef<Path>>,
    commit_entry: Option<CommitEntry>,
) -> Result<MultipartLargeFileUpload, OxenError> {
    let file_path = file_path.as_ref();
    let dst_dir = dst_dir.as_ref();

    let (file_size, hash) = match commit_entry {
        Some(commit_entry) => (commit_entry.num_bytes, commit_entry.hash.clone()),
        None => {
            // Figure out how many parts we need to upload
            let Ok(metadata) = file_path.metadata() else {
                return Err(OxenError::path_does_not_exist(file_path));
            };
            let file_size = metadata.len();
            let hash = util::hasher::hash_file_contents(file_path)?;
            (file_size, hash)
        }
    };

    let uri = format!("/versions/{hash}/create");
    let url = api::endpoint::url_from_repo(remote_repo, &uri)?;
    let client = client::new_for_url(&url)?;

    let body = CreateVersionUploadRequest {
        hash: hash.to_string(),
        file_name: file_path.file_name().unwrap().to_string_lossy().to_string(),
        size: file_size,
        dst_dir: dst_dir.map(|d| d.as_ref().to_path_buf()),
    };

    let body = serde_json::to_string(&body)?;
    let response = client
        .post(&url)
        .header("Content-Type", "application/json")
        .body(body)
        .send()
        .await?;
    response.error_for_status()?;

    Ok(MultipartLargeFileUpload {
        local_path: file_path.to_path_buf(),
        dst_dir: dst_dir.map(|d| d.as_ref().to_path_buf()),
        hash: hash.parse()?,
        size: file_size,
        status: MultipartLargeFileUploadStatus::Pending,
        reason: None,
    })
}

/// Batch download.
///
/// `entries` is a list of `(content_hash, source_path)` pairs — only the hash is sent over the wire
/// (the bulk endpoint identifies blobs by hash); the path is preserved purely for diagnostic
/// surfaces (errors, logs) so end users can identify which file(s) failed.
#[tracing::instrument(skip_all)]
pub async fn download_data_from_version_paths(
    remote_repo: &RemoteRepository,
    entries: &[(String, PathBuf)],
    local_repo: &LocalRepository,
) -> Result<u64, OxenError> {
    let total_retries = max_retries().try_into().unwrap_or(max_retries() as u64);
    let mut num_retries = 0;
    let mut last_err: Option<OxenError> = None;

    while num_retries < total_retries {
        match try_download_data_from_version_paths(remote_repo, entries, local_repo).await {
            Ok(val) => return Ok(val),
            // Short-circuit on errors that won't change on retry (auth failures, 4xx
            // responses, server-confirmed missing blobs). Without this, a doomed pull
            // pays the full exponential backoff before surfacing the diagnostic.
            Err(err) if err.is_fatal_for_retry() => return Err(err),
            Err(err) => {
                num_retries += 1;
                // Exponentially back off
                let sleep_time = num_retries * num_retries;
                log::warn!("Could not download content {err:?} sleeping {sleep_time}");
                last_err = Some(err);
                tokio::time::sleep(std::time::Duration::from_secs(sleep_time)).await;
            }
        }
    }

    // Preserve the failing batch's entries (path + hash) and last underlying error so callers
    // and server logs can identify which file(s) the server couldn't serve. Without this, the
    // symptom surfaces as a bare "failed to download N files" with no way to map back to
    // specific files.
    let last_error = last_err
        .as_ref()
        .map(|e| format!("{e}"))
        .unwrap_or_else(|| "(no attempts made)".to_string());
    let formatted_entries = entries
        .iter()
        .map(|(h, p)| format!("{} (hash: {h})", p.display()))
        .collect::<Vec<_>>()
        .join(", ");
    log::error!(
        "Bulk download exhausted {} retries for {} entries: [{}]. Last error: {}",
        total_retries,
        entries.len(),
        formatted_entries,
        last_error,
    );
    Err(OxenError::DownloadBatchExhausted {
        num_files: entries.len(),
        num_retries: total_retries,
        entries: entries.to_vec(),
        last_error,
    })
}

#[tracing::instrument(skip_all)]
pub async fn try_download_data_from_version_paths(
    remote_repo: &RemoteRepository,
    entries: &[(String, PathBuf)],
    local_repo: &LocalRepository,
) -> Result<u64, OxenError> {
    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
    for (hash, _path) in entries.iter() {
        let line = format!("{hash}\n");
        // log::debug!("download_data_from_version_paths encoding line: {} path: {:?}", line, path);
        encoder.write_all(line.as_bytes())?;
    }
    let body = encoder.finish()?;
    log::debug!("download_data_from_version_paths body len: {}", body.len());

    let url = api::endpoint::url_from_repo(remote_repo, "/versions")?;
    let client = client::new_for_url(&url)?;
    let query_method = Method::from_bytes(b"QUERY").unwrap();
    if let Ok(res) = client.request(query_method, &url).body(body).send().await {
        let res = client::handle_non_json_response(&url, res).await?;

        let stream = res.bytes_stream();
        let reader = tokio_util::io::StreamReader::new(
            stream.map(|result| result.map_err(std::io::Error::other)),
        );
        let buf_reader = tokio::io::BufReader::new(reader);
        let decoder = GzipDecoder::new(buf_reader);
        let mut archive = Archive::new(decoder);

        let version_store = local_repo.version_store();
        let mut size: u64 = 0;

        // Iterate over archive entries and stream them to version store
        let mut entries = archive.entries()?;
        while let Some(file) = entries.next().await {
            let file = match file {
                Ok(file) => file,
                Err(err) => {
                    let err = format!("Could not unwrap file -> {err:?}");
                    return Err(OxenError::basic_str(err));
                }
            };

            let file_hash = file
                .path()
                .map_err(|e| OxenError::basic_str(format!("Failed to get entry path: {e}")))?
                .to_string_lossy()
                .to_string();

            // Get file size from tar entry header
            let file_size = file.header().size()?;
            size += file_size;

            // Stream the file content directly to version store without loading into memory
            match version_store
                .store_version_from_reader(&file_hash, Box::new(file), file_size)
                .await
            {
                Ok(_) => {
                    log::debug!(
                        "Successfully stored file {file_hash} ({file_size} bytes) to version store"
                    );
                }
                Err(err) => {
                    let err =
                        format!("Could not store file {file_hash} to version store -> {err:?}");
                    return Err(OxenError::basic_str(err));
                }
            }
        }

        Ok(size)
    } else {
        let err =
            format!("api::entries::download_data_from_version_paths Err request failed: {url}");
        Err(OxenError::basic_str(err))
    }
}

async fn upload_chunks(
    remote_repo: &RemoteRepository,
    upload: &mut MultipartLargeFileUpload,
    chunk_size: u64,
    parallel_failures: usize,
    max_retries: usize,
    progress: Option<&Arc<PushProgress>>,
) -> Result<Vec<HashMap<String, String>>, OxenError> {
    let client = Arc::new(api::client::new_for_remote_repo(remote_repo)?);

    // Figure out how many parts we need to upload
    let file_size = upload.size;
    let num_chunks = file_size.div_ceil(chunk_size);

    let max_files = concurrency::num_threads_for_items(num_chunks as usize);
    let mut handles = FuturesUnordered::new();
    let semaphore = Arc::new(Semaphore::new(max_files));
    let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures));

    for chunk_number in 0..num_chunks {
        let remote_repo = remote_repo.clone();
        let upload = upload.clone();
        let client = Arc::clone(&client);

        let start = chunk_number * chunk_size;
        let semaphore = semaphore.clone();
        let parallel_failures_semaphore = parallel_failures_semaphore.clone();
        handles.push(tokio::spawn(async move {
                    let permit = semaphore
                        .clone()
                        .acquire_owned()
                        .await
                        .map_err(|err| OxenError::basic_str(format!("Error acquiring semaphore: {err}")))?;
                    let mut chunk = upload_chunk(&client, &remote_repo, &upload, start, chunk_size).await;
                    let mut i = 0;
                    if parallel_failures > 0 {
                        while let Err(ul_err) = chunk {
                            if i >= max_retries {
                                return Err(OxenError::basic_str(format!(
                                    "Failed after too many retries ({max_retries}): {ul_err}"
                                )));
                            }

                            let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
                                OxenError::basic_str(format!(
                                    "Failed too many failures in parallel ({parallel_failures}): {ul_err} ({err})"
                                ))
                            })?;

                            let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
                            sleep(Duration::from_millis(wait_time as u64)).await;

                            chunk = upload_chunk(&client, &remote_repo, &upload, start, chunk_size).await;
                            i += 1;
                            drop(parallel_failure_permit);
                        }
                    }
                    drop(permit);
                    chunk
                    .map_err(|e| OxenError::basic_str(format!("Upload error {e}")))
                    .map(|chunk| (chunk_number, chunk, chunk_size))
                }));
    }

    let mut results: Vec<HashMap<String, String>> = vec![HashMap::default(); num_chunks as usize];

    while let Some(result) = handles.next().await {
        match result {
            Ok(Ok((chunk_number, headers, size))) => {
                log::debug!("Uploaded part {chunk_number} with size {size}");
                results[chunk_number as usize] = headers;
                if let Some(p) = progress {
                    p.add_bytes(size);
                }
            }
            Ok(Err(py_err)) => {
                return Err(py_err);
            }
            Err(err) => {
                return Err(OxenError::basic_str(format!(
                    "Error occurred while uploading: {err}"
                )));
            }
        }
    }
    if let Some(p) = progress {
        p.add_files(1);
    }
    Ok(results)
}

async fn upload_chunk(
    client: &reqwest::Client,
    remote_repo: &RemoteRepository,
    upload: &MultipartLargeFileUpload,
    start: u64,
    chunk_size: u64,
) -> Result<HashMap<String, String>, OxenError> {
    let path = &upload.local_path;
    let mut options = OpenOptions::new();
    let mut file = options.read(true).open(path).await?;

    let file_size = upload.size;
    let bytes_transferred = std::cmp::min(file_size - start, chunk_size);

    file.seek(SeekFrom::Start(start)).await?;
    let chunk = file.take(chunk_size);

    let file_hash = &upload.hash.to_string();

    let uri = format!("/versions/{file_hash}/chunks?offset={start}");
    let url = api::endpoint::url_from_repo(remote_repo, &uri)?;

    let response = client
        .put(url)
        .header(CONTENT_LENGTH, bytes_transferred)
        .body(reqwest::Body::wrap_stream(FramedRead::new(
            chunk,
            BytesCodec::new(),
        )))
        .send()
        .await?;
    let response = response.error_for_status()?;
    let mut headers = HashMap::new();
    for (name, value) in response.headers().into_iter() {
        headers.insert(
            name.to_string(),
            value
                .to_str()
                .map_err(|e| OxenError::basic_str(format!("Invalid header value: {e}")))?
                .to_owned(),
        );
    }
    Ok(headers)
}

async fn complete_multipart_large_file_upload(
    remote_repo: &RemoteRepository,
    upload: MultipartLargeFileUpload,
    num_chunks: usize,
    workspace_id: Option<String>,
) -> Result<MultipartLargeFileUpload, OxenError> {
    let file_hash = &upload.hash.to_string();

    let uri = format!("/versions/{file_hash}/complete");
    let url = api::endpoint::url_from_repo(remote_repo, &uri)?;
    log::debug!("complete_multipart_large_file_upload {url}");
    let client = client::new_for_url(&url)?;

    let body = CompleteVersionUploadRequest {
        files: vec![CompletedFileUpload {
            hash: file_hash.to_string(),
            file_name: upload
                .local_path
                .file_name()
                .unwrap()
                .to_string_lossy()
                .to_string(),
            dst_dir: upload.dst_dir.clone(),
            num_chunks: Some(num_chunks),
            upload_results: None,
        }],
        workspace_id,
    };

    let body = serde_json::to_string(&body)?;
    let response = client.post(&url).body(body).send().await?;
    let body = client::parse_json_body(&url, response).await?;
    log::debug!("complete_multipart_large_file_upload got body: {body}");
    Ok(upload)
}

/// Multipart batch upload with retry
/// Uploads a batch of small files to the server in parallel and retries on failure
/// Returns a list of files that failed to upload
pub async fn multipart_batch_upload_with_retry(
    local_repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    chunk: &[CommitEntry],
    client: &reqwest::Client,
) -> Result<(), OxenError> {
    let mut files_to_retry: Vec<ErrorFileInfo> = vec![];
    let mut first_try = true;
    let mut retry_count: usize = 0;
    let max_retries = max_retries();

    while (first_try || !files_to_retry.is_empty()) && retry_count < max_retries {
        first_try = false;
        retry_count += 1;

        files_to_retry =
            multipart_batch_upload(local_repo, remote_repo, chunk, client, files_to_retry).await?;

        if !files_to_retry.is_empty() {
            let wait_time = exponential_backoff(BASE_WAIT_TIME, retry_count, MAX_WAIT_TIME);
            sleep(Duration::from_millis(wait_time as u64)).await;
        }
    }
    if files_to_retry.is_empty() {
        Ok(())
    } else {
        Err(OxenError::basic_str(format!(
            "Failed to upload files: {files_to_retry:#?}"
        )))
    }
}

pub async fn multipart_batch_upload(
    local_repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    chunk: &[CommitEntry],
    client: &reqwest::Client,
    files_to_retry: Vec<ErrorFileInfo>,
) -> Result<Vec<ErrorFileInfo>, OxenError> {
    let version_store = local_repo.version_store();
    let mut form = reqwest::multipart::Form::new();
    let mut err_files: Vec<ErrorFileInfo> = vec![];

    // if it's the first try, we don't have any files to retry
    let retry_hashes: HashSet<String> = if files_to_retry.is_empty() {
        HashSet::new()
    } else {
        files_to_retry.iter().map(|f| f.hash.clone()).collect()
    };

    for commit_entry in chunk {
        let file_hash = &commit_entry.hash;

        // if it's not the first try and the file is not in the retry list, skip
        if !files_to_retry.is_empty() && !retry_hashes.contains(file_hash) {
            continue;
        }

        let data = version_store.get_version(file_hash).await?;
        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        std::io::copy(&mut data.as_slice(), &mut encoder)?;
        let compressed_bytes = match encoder.finish() {
            Ok(bytes) => bytes,
            Err(e) => {
                log::error!("Failed to finish gzip for file {}: {}", file_hash, e);
                err_files.push(ErrorFileInfo {
                    hash: file_hash.clone(),
                    path: None,
                    error: format!("Failed to finish gzip for file {}: {}", file_hash, e),
                });
                continue;
            }
        };

        let file_part = reqwest::multipart::Part::bytes(compressed_bytes)
            .file_name(commit_entry.hash.clone())
            .mime_str("application/gzip")?;
        form = form.part("file[]", file_part);
    }

    // If there are nodes to mark as synced, re-route API call
    let uri = ("/versions").to_string();

    let url = api::endpoint::url_from_repo(remote_repo, &uri)?;

    // Post the node hashes to sync on the first chunk upload

    let response = client.post(&url).multipart(form).send().await?;
    let body = client::parse_json_body(&url, response).await?;
    let response: ErrorFilesResponse = serde_json::from_str(&body)?;

    err_files.extend(response.err_files);

    Ok(err_files)
}

pub fn exponential_backoff(base_wait_time: usize, n: usize, max: usize) -> usize {
    log::debug!(
        "Exponential backoff got called with base_wait_time {base_wait_time}. n {n}, and max {max}"
    );
    (base_wait_time + n.pow(2) + jitter()).min(max)
}

fn jitter() -> usize {
    thread_rng().gen_range(0..=500)
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use crate::api;
    use crate::error::OxenError;
    use crate::test;

    #[tokio::test]
    async fn test_upload_large_file_in_chunks() -> Result<(), OxenError> {
        test::run_remote_repo_test_bounding_box_csv_pushed(|_local_repo, remote_repo| async move {
            let path = test::test_30k_parquet();

            // Get original file size
            let metadata = path.metadata().unwrap();
            let original_file_size = metadata.len();

            // Just testing upload, not adding to workspace
            let workspace_id = None;
            let dst_dir: Option<PathBuf> = None;
            let result = api::client::versions::parallel_large_file_upload(
                &remote_repo,
                path,
                dst_dir,
                workspace_id,
                None,
                None,
            )
            .await;
            assert!(result.is_ok());

            let version = api::client::versions::get(&remote_repo, result.unwrap().hash).await?;
            assert!(version.is_some());
            assert_eq!(version.unwrap().size, original_file_size);

            Ok(remote_repo)
        })
        .await
    }

    /// When the bulk versions download endpoint is asked for hashes that don't exist
    /// on the server, the server's pre-flight check returns a structured 404 and the
    /// client's retry loop short-circuits — instead of paying multiple rounds of
    /// exponential backoff that won't change the outcome.
    #[tokio::test]
    async fn test_bulk_download_short_circuits_on_missing_blob_on_server() -> Result<(), OxenError>
    {
        test::run_remote_repo_test_bounding_box_csv_pushed(|local_repo, remote_repo| async move {
            // A well-formed 32-char hex string that can't possibly exist on the server.
            let bogus_hash = "deadbeefdeadbeefdeadbeefdeadbeef".to_string();
            let entries = vec![(bogus_hash.clone(), PathBuf::from("does-not-exist.txt"))];

            let result = api::client::versions::download_data_from_version_paths(
                &remote_repo,
                &entries,
                &local_repo,
            )
            .await;

            let err = result.expect_err("expected error for missing hash");
            // The short-circuit returns the underlying fatal error directly — *not*
            // DownloadBatchExhausted, which is only emitted after the retry loop runs
            // out of attempts. Seeing DownloadBatchExhausted here would mean the loop
            // retried on a 4xx and slept its way through backoff.
            assert!(
                !matches!(&err, OxenError::DownloadBatchExhausted { .. }),
                "should have short-circuited on the 4xx instead of exhausting retries: {err:?}"
            );
            assert!(
                err.is_fatal_for_retry(),
                "missing-blob error should classify as fatal: {err:?}"
            );

            // The rendered error names the missing hash so the user can map it back to
            // the broken blob. We assert on the surface message rather than the variant
            // shape since the server's wire-format may evolve.
            let rendered = err.to_string();
            assert!(
                rendered.contains(&bogus_hash),
                "error should name the missing hash; got: {rendered}"
            );

            Ok(remote_repo)
        })
        .await
    }
}