liboxen 0.50.1

use futures::future;
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use crate::constants::{AVG_CHUNK_SIZE, OXEN_HIDDEN_DIR};
use crate::core::refs::with_ref_manager;
use crate::error::OxenError;
use crate::model::merkle_tree::node::{EMerkleTreeNode, MerkleTreeNode};
use crate::model::{Branch, Commit, CommitEntry, MerkleHash};
use crate::model::{LocalRepository, RemoteBranch, RemoteRepository};
use crate::repositories;
use crate::storage::VersionStore;
use crate::util::concurrency;
use crate::{api, util};

use crate::core::progress::pull_progress::PullProgress;
use crate::opts::fetch_opts::FetchOpts;

#[tracing::instrument(skip(repo, remote_repo, fetch_opts), fields(repo_path = %repo.path.display(), branch = %fetch_opts.branch))]
pub async fn fetch_remote_branch(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    fetch_opts: &FetchOpts,
) -> Result<Branch, OxenError> {
    log::debug!("fetching remote branch with opts {fetch_opts:?}");

    // Start the timer
    let start = std::time::Instant::now();

    // Keep track of how many bytes we have downloaded
    let pull_progress = Arc::new(PullProgress::new());
    pull_progress.set_message(format!("Fetching remote branch {}", fetch_opts.branch));

    // Find the head commit on the remote branch
    let Some(remote_branch) =
        api::client::branches::get_by_name(remote_repo, &fetch_opts.branch).await?
    else {
        return Err(OxenError::remote_branch_not_found(&fetch_opts.branch));
    };

    // We may not have a head commit if the repo is empty (initial clone)
    if let Some(head_commit) = repositories::commits::head_commit_maybe(repo)? {
        log::debug!("Head commit: {head_commit}");
        log::debug!("Remote branch commit: {}", remote_branch.commit_id);
        if head_commit.id == remote_branch.commit_id {
            if !fetch_opts.missing_files && !fetch_opts.all {
                // Already up to date, nothing to do
                println!("Repository is up to date.");
                with_ref_manager(repo, |manager| {
                    manager.set_branch_commit_id(&remote_branch.name, &remote_branch.commit_id)
                })?;
                return Ok(remote_branch);
            }

            if fetch_opts.all {
                fetch_full_tree_and_hashes(repo, remote_repo, &remote_branch, &pull_progress)
                    .await?;
            }
            // missing_files: skip tree sync (trees are already local), fall through to entry scan
        } else if fetch_opts.all {
            fetch_full_tree_and_hashes(repo, remote_repo, &remote_branch, &pull_progress).await?;
        } else {
            // Download the nodes from the commits between the head and the remote head
            sync_from_head(
                repo,
                remote_repo,
                fetch_opts,
                &remote_branch,
                &head_commit,
                &pull_progress,
            )
            .await?;
        }
    } else {
        // If there is no head commit, we are fetching all commits from the remote branch commit
        log::debug!(
            "Fetching all commits from remote branch {}",
            remote_branch.commit_id
        );
        if fetch_opts.all {
            fetch_full_tree_and_hashes(repo, remote_repo, &remote_branch, &pull_progress).await?;
        } else {
            sync_tree_from_commit(
                repo,
                remote_repo,
                &remote_branch.commit_id,
                fetch_opts,
                &pull_progress,
            )
            .await?;
        }
    }

    // Early exit for remote repo
    if repo.is_remote_mode() {
        // Write the new branch commit id to the local repo
        if fetch_opts.should_update_branch_head {
            repositories::branches::update(repo, &fetch_opts.branch, &remote_branch.commit_id)?;
        }
        return Ok(remote_branch);
    }

    // Scan for missing commits. `collect_missing_entries` walks the target commit's
    // merkle tree but prunes any subtree whose root-hash matches one we already have
    // from HEAD's tree (its `shared_hashes` set is seeded from HEAD), so it only
    // descends into the parts of the tree that actually changed between HEAD and the
    // target. `get_missing_entries_for_pull` then stats each remaining entry's blob
    // and only re-downloads the ones that aren't already in the local version store.
    // Cost is proportional to the HEAD-target diff, not repo size.
    let commits: HashSet<Commit> = if fetch_opts.all {
        repositories::commits::list_from(repo, &remote_branch.commit_id)?
            .into_iter()
            .collect()
    } else {
        let hash = remote_branch.commit_id.parse()?;
        let commit_node = repositories::tree::get_node_by_id(repo, &hash)?
            .ok_or_else(|| OxenError::basic_str("Commit node not found"))?;
        HashSet::from([commit_node.commit()?.to_commit()])
    };

    log::debug!("Fetch got {} commits", commits.len());

    let mut total_bytes = 0;
    let missing_entries = collect_missing_entries(
        repo,
        &commits,
        &fetch_opts.subtree_paths,
        &fetch_opts.depth,
        fetch_opts.missing_files,
        &mut total_bytes,
    )?;
    log::debug!(
        "Fetch got {} potentially missing entries",
        missing_entries.len()
    );
    let missing_entries: Vec<CommitEntry> = missing_entries.into_iter().collect();
    pull_progress.finish();
    let pull_progress = Arc::new(PullProgress::new_with_totals(
        missing_entries.len() as u64,
        total_bytes,
    ));
    pull_entries_to_versions_dir(repo, remote_repo, &missing_entries, &pull_progress).await?;

    // If we fetched the data, we're no longer shallow
    repo.write_is_shallow(false)?;

    // Write the new branch commit id to the local repo
    if fetch_opts.should_update_branch_head {
        repositories::branches::update(repo, &fetch_opts.branch, &remote_branch.commit_id)?;
    }

    pull_progress.finish();
    let duration = std::time::Duration::from_millis(start.elapsed().as_millis() as u64);

    println!(
        "🐂 oxen downloaded {} ({} files) in {}",
        bytesize::ByteSize::b(pull_progress.get_num_bytes()),
        pull_progress.get_num_files(),
        humantime::format_duration(duration)
    );

    Ok(remote_branch)
}

#[tracing::instrument(skip(repo, remote_repo, fetch_opts, pull_progress), fields(repo_path = %repo.path.display(), branch = %branch.name, head_commit_id = %head_commit.id))]
async fn sync_from_head(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    fetch_opts: &FetchOpts,
    branch: &Branch,
    head_commit: &Commit,
    pull_progress: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    let repo_hidden_dir = util::fs::oxen_hidden_dir(&repo.path);
    log::debug!("sync_from_head head_commit: {head_commit}");
    log::debug!("sync_from_head branch: {branch}");

    // If HEAD commit IS on the remote server, that means we are behind the remote branch
    if api::client::tree::has_node(remote_repo, head_commit.id.parse()?).await? {
        log::debug!("sync_from_head has head commit: {head_commit}");
        pull_progress.set_message(format!(
            "Downloading commits from {} to {}",
            head_commit.id, branch.commit_id
        ));
        api::client::tree::download_trees_between(
            repo,
            remote_repo,
            &head_commit.id,
            &branch.commit_id,
            fetch_opts,
        )
        .await?;
        api::client::commits::download_base_head_dir_hashes(
            remote_repo,
            &head_commit.id,
            &branch.commit_id,
            &repo_hidden_dir,
        )
        .await?;
    } else {
        // If HEAD commit is NOT on the remote server, that means we are ahead of the remote branch
        // If the node does not exist on the remote server,
        // we need to sync all the commits from the commit id and their parents
        // TODO: This logic doesn't seem correct. Revisit this.
        sync_tree_from_commit(
            repo,
            remote_repo,
            &branch.commit_id,
            fetch_opts,
            pull_progress,
        )
        .await?;
    }
    Ok(())
}

// Sync all the commits from the commit (and their parents)
#[tracing::instrument(skip(repo, remote_repo, commit_id, fetch_opts, pull_progress), fields(repo_path = %repo.path.display()))]
async fn sync_tree_from_commit(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    commit_id: impl AsRef<str>,
    fetch_opts: &FetchOpts,
    pull_progress: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    let repo_hidden_dir = util::fs::oxen_hidden_dir(&repo.path);

    pull_progress.set_message(format!("Downloading commits from {}", commit_id.as_ref()));
    api::client::tree::download_trees_from(repo, remote_repo, &commit_id.as_ref(), fetch_opts)
        .await?;
    api::client::commits::download_dir_hashes_from_commit(
        remote_repo,
        commit_id.as_ref(),
        &repo_hidden_dir,
    )
    .await?;
    Ok(())
}

fn collect_missing_entries(
    repo: &LocalRepository,
    commits: &HashSet<Commit>,
    subtree_paths: &Option<Vec<PathBuf>>,
    depth: &Option<i32>,
    skip_shared_hashes: bool,
    total_bytes: &mut u64,
) -> Result<HashSet<CommitEntry>, OxenError> {
    let mut missing_entries = HashSet::new();

    // When skip_shared_hashes is true (e.g. --missing-files), start with an empty set
    // so all entries in the commit tree are considered candidates for download.
    let mut shared_hashes = if skip_shared_hashes {
        HashSet::new()
    } else if let Some(head_commit) = repositories::commits::head_commit_maybe(repo)? {
        let mut starting_node_hashes = HashSet::new();
        repositories::tree::populate_starting_hashes(
            repo,
            &head_commit,
            &None,
            &None,
            &mut starting_node_hashes,
        )?;
        starting_node_hashes
    } else {
        HashSet::new()
    };

    let mut file_hashes_seen = HashSet::new();

    for commit in commits {
        if let Some(subtree_paths) = subtree_paths {
            log::debug!(
                "collect_missing_entries for {subtree_paths:?} subtree paths and depth {depth:?}"
            );
            for subtree_path in subtree_paths {
                let mut unique_hashes = HashSet::new();
                let Some(tree) = repositories::tree::get_subtree_by_depth_with_unique_children(
                    repo,
                    commit,
                    subtree_path,
                    Some(&shared_hashes),
                    Some(&mut unique_hashes),
                    None,
                    depth.unwrap_or(-1),
                )?
                else {
                    log::warn!("get_subtree_by_depth returned None for path: {subtree_path:?}");
                    continue;
                };

                shared_hashes.extend(unique_hashes);

                collect_missing_entries_for_subtree(
                    &tree,
                    subtree_path,
                    &mut missing_entries,
                    &mut file_hashes_seen,
                    total_bytes,
                )?;
            }
        } else {
            let mut unique_hashes = HashSet::new();
            let Some(tree) = repositories::tree::get_subtree_by_depth_with_unique_children(
                repo,
                commit,
                PathBuf::from("."),
                Some(&shared_hashes),
                Some(&mut unique_hashes),
                None,
                depth.unwrap_or(-1),
            )?
            else {
                log::warn!("get_subtree_by_depth returned None for commit: {commit:?}");
                continue;
            };

            shared_hashes.extend(unique_hashes);

            collect_missing_entries_for_subtree(
                &tree,
                Path::new(""),
                &mut missing_entries,
                &mut file_hashes_seen,
                total_bytes,
            )?;
        }
    }
    Ok(missing_entries)
}

fn collect_missing_entries_for_subtree(
    node: &MerkleTreeNode,
    current_path: &Path,
    missing_entries: &mut HashSet<CommitEntry>,
    file_hashes_seen: &mut HashSet<MerkleHash>,
    total_bytes: &mut u64,
) -> Result<(), OxenError> {
    use crate::model::MerkleTreeNodeType;

    match node.node.node_type() {
        MerkleTreeNodeType::File | MerkleTreeNodeType::FileChunk => {
            let file_hash = *node.node.hash();
            if file_hashes_seen.insert(file_hash) {
                let mut commit_entry = CommitEntry::from_node(&node.node);
                commit_entry.path = current_path.join(&commit_entry.path);
                *total_bytes += commit_entry.num_bytes;
                missing_entries.insert(commit_entry);
            }
        }
        MerkleTreeNodeType::Dir => {
            let dir_path = if let EMerkleTreeNode::Directory(dir_node) = &node.node {
                let name = dir_node.name();
                if name.is_empty() {
                    current_path.to_path_buf()
                } else {
                    current_path.join(name)
                }
            } else {
                current_path.to_path_buf()
            };
            for child in &node.children {
                collect_missing_entries_for_subtree(
                    child,
                    &dir_path,
                    missing_entries,
                    file_hashes_seen,
                    total_bytes,
                )?;
            }
        }
        // Commits and VNodes don't change the path
        MerkleTreeNodeType::Commit | MerkleTreeNodeType::VNode => {
            for child in &node.children {
                collect_missing_entries_for_subtree(
                    child,
                    current_path,
                    missing_entries,
                    file_hashes_seen,
                    total_bytes,
                )?;
            }
        }
    }
    Ok(())
}

pub async fn fetch_tree_and_hashes_for_commit_id(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    commit_id: &str,
) -> Result<(), OxenError> {
    let repo_hidden_dir = repo.path.join(OXEN_HIDDEN_DIR);
    api::client::commits::download_dir_hashes_db_to_path(remote_repo, commit_id, &repo_hidden_dir)
        .await?;

    let hash = commit_id.parse()?;
    api::client::tree::download_tree_from(repo, remote_repo, &hash).await?;

    api::client::commits::download_dir_hashes_from_commit(remote_repo, commit_id, &repo_hidden_dir)
        .await?;

    Ok(())
}

#[tracing::instrument(skip(repo, remote_repo, pull_progress), fields(repo_path = %repo.path.display(), branch = %remote_branch.name))]
pub async fn fetch_full_tree_and_hashes(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    remote_branch: &Branch,
    pull_progress: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    // Download the latest merkle tree
    // Must do this before downloading the commit node
    // because the commit node references the merkle tree
    let repo_hidden_dir = repo.path.join(OXEN_HIDDEN_DIR);
    api::client::commits::download_dir_hashes_db_to_path(
        remote_repo,
        &remote_branch.commit_id,
        &repo_hidden_dir,
    )
    .await?;

    pull_progress.set_message(format!(
        "Downloading all commits from {}",
        remote_branch.commit_id
    ));
    // Download the latest merkle tree
    // let hash = MerkleHash::from_str(&remote_branch.commit_id)?;
    api::client::tree::download_tree(repo, remote_repo).await?;
    // let commit_node = CommitMerkleTree::read_node(repo, &hash, true)?.unwrap();

    // Download the commit history
    // Check what our HEAD commit is locally
    if let Some(head_commit) = repositories::commits::head_commit_maybe(repo)? {
        // Remote is not guaranteed to have our head commit
        // If it doesn't, we will download all commit dir hashes from the remote branch commit
        if api::client::tree::has_node(remote_repo, head_commit.id.parse()?).await? {
            // Download the dir_hashes between the head commit and the remote branch commit
            let base_commit_id = head_commit.id;
            let head_commit_id = &remote_branch.commit_id;

            api::client::commits::download_base_head_dir_hashes(
                remote_repo,
                &base_commit_id,
                head_commit_id,
                &repo_hidden_dir,
            )
            .await?;
        } else {
            // Download the dir hashes from the remote branch commit
            api::client::commits::download_dir_hashes_from_commit(
                remote_repo,
                &remote_branch.commit_id,
                &repo_hidden_dir,
            )
            .await?;
        }
    } else {
        // Download the dir hashes from the remote branch commit
        api::client::commits::download_dir_hashes_from_commit(
            remote_repo,
            &remote_branch.commit_id,
            &repo_hidden_dir,
        )
        .await?;
    };
    Ok(())
}

/// Fetch missing entries for a commit
/// If there is no remote, or we can't find the remote, this will *not* error
pub async fn maybe_fetch_missing_entries(
    repo: &LocalRepository,
    commit: &Commit,
) -> Result<(), OxenError> {
    // If we don't have a remote, there are no missing entries, so return
    let rb = RemoteBranch::default();
    let remote = repo.get_remote(&rb.remote);
    let Some(remote) = remote else {
        log::debug!("No remote, no missing entries to fetch");
        return Ok(());
    };

    let Some(commit_merkle_tree) = repositories::tree::get_root_with_children(repo, commit)? else {
        log::warn!("get_root_with_children returned None for commit: {commit:?}");
        return Ok(());
    };

    let remote_repo = match api::client::repositories::get_by_remote(&remote).await {
        Ok(repo) => repo,
        Err(OxenError::RemoteRepoNotFound(_)) => {
            log::debug!("No remote repo found, skipping fetch");
            return Ok(());
        }
        Err(err) => {
            log::warn!("Error getting remote repo: {err}");
            return Ok(());
        }
    };

    // TODO: what should we print here? If there is nothing to pull, we
    // shouldn't show the PullProgress
    log::debug!("Fetching missing entries for commit {commit}");

    // Keep track of how many bytes we have downloaded
    let pull_progress = Arc::new(PullProgress::new());

    // Recursively download the entries
    let directory = PathBuf::from("");
    r_download_entries(
        repo,
        &remote_repo,
        &commit_merkle_tree,
        &directory,
        &pull_progress,
    )
    .await?;

    Ok(())
}

async fn r_download_entries(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    node: &MerkleTreeNode,
    directory: &Path,
    pull_progress: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    log::debug!(
        "fetch r_download_entries ({}) {:?} {:?}",
        node.children.len(),
        node.hash,
        node.node
    );
    for child in &node.children {
        let mut new_directory = directory.to_path_buf();
        if let EMerkleTreeNode::Directory(dir_node) = &child.node {
            new_directory.push(dir_node.name());
        }

        if child.has_children() {
            Box::pin(r_download_entries(
                repo,
                remote_repo,
                child,
                &new_directory,
                pull_progress,
            ))
            .await?;
        }
    }

    if let EMerkleTreeNode::VNode(_) = &node.node {
        // Figure out which entries need to be downloaded
        let mut missing_entries: Vec<CommitEntry> = vec![];
        let missing_hashes = repositories::tree::list_missing_file_hashes(repo, &node.hash).await?;

        for child in &node.children {
            if let EMerkleTreeNode::File(file_node) = &child.node {
                if !missing_hashes.contains(&child.hash) {
                    continue;
                }

                missing_entries.push(CommitEntry {
                    commit_id: file_node.last_commit_id().to_string(),
                    path: directory.join(file_node.name()),
                    hash: child.hash.to_string(),
                    num_bytes: file_node.num_bytes(),
                    last_modified_seconds: file_node.last_modified_seconds(),
                    last_modified_nanoseconds: file_node.last_modified_nanoseconds(),
                });
            }
        }

        pull_entries_to_versions_dir(repo, remote_repo, &missing_entries, pull_progress).await?;
    }

    Ok(())
}

// pull entries from remote repo to versions dir
#[tracing::instrument(skip(repo, remote_repo, entries, progress_bar), fields(repo_path = %repo.path.display(), num_entries = entries.len()))]
pub async fn pull_entries_to_versions_dir(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    entries: &[CommitEntry],
    progress_bar: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    log::debug!("entries.len() {}", entries.len());
    if entries.is_empty() {
        return Ok(());
    }

    let version_store = repo.version_store();
    let missing_entries = get_missing_entries_for_pull(&version_store, entries).await?;
    log::debug!("Pulling {} missing entries", missing_entries.len());

    if missing_entries.is_empty() {
        return Ok(());
    }

    // Some files may be much larger than others....so we can't just download them within a single body
    // Hence we chunk and send the big ones, and bundle and download the small ones

    // For files smaller than AVG_CHUNK_SIZE, we are going to group them, zip them up, and transfer them
    let smaller_entries: Vec<CommitEntry> = missing_entries
        .iter()
        .filter(|e| e.num_bytes <= AVG_CHUNK_SIZE)
        .map(|e| e.to_owned())
        .collect();

    // For files larger than AVG_CHUNK_SIZE, we are going break them into chunks and download the chunks in parallel
    let larger_entries: Vec<CommitEntry> = missing_entries
        .iter()
        .filter(|e| e.num_bytes > AVG_CHUNK_SIZE)
        .map(|e| e.to_owned())
        .collect();

    let large_entries_sync = pull_large_entries(repo, remote_repo, larger_entries, progress_bar);

    let small_entries_sync = pull_small_entries(repo, remote_repo, smaller_entries, progress_bar);

    match tokio::join!(large_entries_sync, small_entries_sync) {
        (Ok(_), Ok(_)) => {
            log::debug!("Successfully synced entries!");
        }
        (Err(err), Ok(_)) | (Ok(_), Err(err)) => return Err(err),
        (Err(large_err), Err(small_err)) => {
            log::error!("Large entries sync also failed: {large_err}");
            return Err(small_err);
        }
    }

    Ok(())
}

async fn pull_large_entries(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    entries: Vec<CommitEntry>,
    progress_bar: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    if entries.is_empty() {
        return Ok(());
    }
    // Pull the large entries in parallel
    type PieceOfWork = (LocalRepository, RemoteRepository, CommitEntry);
    type TaskQueue = deadqueue::limited::Queue<PieceOfWork>;

    log::debug!("Chunking and sending {} larger files", entries.len());
    let entries: Vec<PieceOfWork> = entries
        .iter()
        .map(|e| (repo.to_owned(), remote_repo.to_owned(), e.to_owned()))
        .collect();

    let queue = Arc::new(TaskQueue::new(entries.len()));
    for entry in entries.iter() {
        queue.try_push(entry.to_owned()).unwrap();
    }

    let worker_count = concurrency::num_threads_for_items(entries.len());
    log::debug!(
        "worker_count {} entries len {}",
        worker_count,
        entries.len()
    );
    let mut handles = vec![];

    for worker in 0..worker_count {
        let queue = queue.clone();
        let progress_bar = Arc::clone(progress_bar);
        let handle: tokio::task::JoinHandle<Result<(), OxenError>> = tokio::spawn(async move {
            loop {
                let Some((repo, remote_repo, commit_entry)) = queue.try_pop() else {
                    // reached end of queue
                    break;
                };
                log::debug!("worker[{worker}] processing task...");

                // Chunk and individual files
                let remote_path = &commit_entry.path;

                // Download to the tmp path, then copy over to the entries dir
                api::client::entries::pull_large_entry(
                    &repo,
                    &remote_repo,
                    remote_path,
                    &commit_entry,
                )
                .await?;

                log::debug!("Pulled large entry {remote_path:?} to versions dir");
                progress_bar.add_bytes(commit_entry.num_bytes);
                progress_bar.add_files(1);
            }
            Ok(())
        });
        handles.push(handle);
    }
    let join_results = future::join_all(handles).await;
    for res in join_results {
        match res {
            Err(e) => return Err(OxenError::basic_str(format!("worker task panicked: {e}"))),
            Ok(Err(e)) => return Err(e),
            Ok(Ok(())) => {}
        }
    }

    log::debug!("All large file tasks done. :-)");

    Ok(())
}

async fn pull_small_entries(
    repo: &LocalRepository,
    remote_repo: &RemoteRepository,
    entries: Vec<CommitEntry>,
    progress_bar: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    if entries.is_empty() {
        return Ok(());
    }

    let total_size = repositories::entries::compute_entries_size(&entries)?;

    // Compute num chunks
    let num_chunks = ((total_size / AVG_CHUNK_SIZE) + 1) as usize;

    let mut chunk_size = entries.len() / num_chunks;
    if num_chunks > entries.len() {
        chunk_size = entries.len();
    }

    log::debug!(
        "pull_small_entries got {} missing content IDs",
        entries.len()
    );

    // Split into chunks, zip up, and post to server. We carry `(hash, path)` pairs through
    // the queue (rather than just hashes) so the bulk-download client can put the file paths
    // into its retry-exhausted error message — paths are far more recognizable to end users
    // than content hashes.
    type PieceOfWork = (RemoteRepository, Vec<(String, PathBuf)>, LocalRepository);
    type TaskQueue = deadqueue::limited::Queue<PieceOfWork>;

    log::debug!(
        "pull_small_entries creating {num_chunks} chunks from {total_size} bytes with size {chunk_size}"
    );

    let chunks: Vec<PieceOfWork> = entries
        .chunks(chunk_size)
        .map(|chunk| {
            let entries = chunk
                .iter()
                .map(|commit_entry| (commit_entry.hash.clone(), commit_entry.path.clone()))
                .collect();
            (remote_repo.to_owned(), entries, repo.to_owned())
        })
        .collect();

    let worker_count = concurrency::num_threads_for_items(entries.len());
    let queue = Arc::new(TaskQueue::new(chunks.len()));
    for chunk in chunks {
        queue.try_push(chunk).unwrap();
    }
    let mut handles = vec![];

    for worker in 0..worker_count {
        let queue = queue.clone();
        let progress_bar = Arc::clone(progress_bar);
        let handle: tokio::task::JoinHandle<Result<(), OxenError>> = tokio::spawn(async move {
            loop {
                let Some((remote_repo, chunk, local_repo)) = queue.try_pop() else {
                    // reached end of queue
                    break;
                };
                log::debug!("worker[{worker}] processing task...");

                let download_size = api::client::versions::download_data_from_version_paths(
                    &remote_repo,
                    &chunk,
                    &local_repo,
                )
                .await?;

                progress_bar.add_bytes(download_size);
                progress_bar.add_files(chunk.len() as u64);
            }
            Ok(())
        });
        handles.push(handle);
    }
    let join_results = future::join_all(handles).await;
    for res in join_results {
        match res {
            Err(e) => return Err(OxenError::basic_str(format!("worker task panicked: {e}"))),
            Ok(Err(e)) => return Err(e),
            Ok(Ok(())) => {}
        }
    }

    log::debug!("All tasks done. :-)");

    Ok(())
}

// download entries to working dir
pub async fn download_entries_to_working_dir(
    remote_repo: &RemoteRepository,
    entries: &[CommitEntry],
    dst: &Path,
    progress_bar: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    log::debug!(
        "download_entries_to_working_dir entries.len() {}",
        entries.len()
    );

    if entries.is_empty() {
        return Ok(());
    }

    let missing_entries = get_missing_entries_for_download(entries, dst);
    log::debug!("Pulling {} missing entries", missing_entries.len());

    if missing_entries.is_empty() {
        return Ok(());
    }

    // Some files may be much larger than others....so we can't just download them within a single body
    // Hence we chunk and send the big ones, and bundle and download the small ones

    // For files smaller than AVG_CHUNK_SIZE, we are going to group them, zip them up, and transfer them
    let smaller_entries: Vec<CommitEntry> = missing_entries
        .iter()
        .filter(|e| e.num_bytes <= AVG_CHUNK_SIZE)
        .map(|e| e.to_owned())
        .collect();

    // For files larger than AVG_CHUNK_SIZE, we are going break them into chunks and download the chunks in parallel
    let larger_entries: Vec<CommitEntry> = missing_entries
        .iter()
        .filter(|e| e.num_bytes > AVG_CHUNK_SIZE)
        .map(|e| e.to_owned())
        .collect();

    let large_entries_sync =
        download_large_entries(remote_repo, larger_entries, &dst, progress_bar);
    log::info!("Downloaded large entries");
    let small_entries_sync =
        download_small_entries(remote_repo, smaller_entries, &dst, progress_bar);
    log::info!("Downloaded small entries");

    match tokio::join!(large_entries_sync, small_entries_sync) {
        (Ok(_), Ok(_)) => {
            log::debug!("Successfully synced entries!");
        }
        (Err(err), Ok(_)) | (Ok(_), Err(err)) => return Err(err),
        (Err(large_err), Err(small_err)) => {
            log::error!("Large entries sync also failed: {large_err}");
            return Err(small_err);
        }
    }
    log::info!("Finished download_entries_to_working_dir");
    Ok(())
}

async fn download_large_entries(
    remote_repo: &RemoteRepository,
    entries: Vec<CommitEntry>,
    dst: impl AsRef<Path>,
    progress_bar: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    if entries.is_empty() {
        return Ok(());
    }
    // Pull the large entries in parallel
    type PieceOfWork = (RemoteRepository, CommitEntry, PathBuf, PathBuf);
    type TaskQueue = deadqueue::limited::Queue<PieceOfWork>;

    log::debug!("Chunking and sending {} larger files", entries.len());
    let large_entry_paths = working_dir_paths_from_large_entries(&entries, dst.as_ref());
    let entries: Vec<PieceOfWork> = entries
        .iter()
        .zip(large_entry_paths.iter())
        .map(|(e, path)| {
            (
                remote_repo.to_owned(),
                e.to_owned(),
                dst.as_ref().to_owned(),
                path.to_owned(),
            )
        })
        .collect();

    let queue = Arc::new(TaskQueue::new(entries.len()));
    for entry in entries.iter() {
        queue.try_push(entry.to_owned()).unwrap();
    }

    let worker_count = concurrency::num_threads_for_items(entries.len());
    log::debug!(
        "worker_count {} entries len {}",
        worker_count,
        entries.len()
    );
    let mut handles = vec![];
    let tmp_dir = util::fs::oxen_hidden_dir(dst).join("tmp").join("pulled");
    log::debug!("Backing up pulls to tmp dir: {:?}", &tmp_dir);
    for worker in 0..worker_count {
        let queue = queue.clone();
        let progress_bar = Arc::clone(progress_bar);
        let handle: tokio::task::JoinHandle<Result<(), OxenError>> = tokio::spawn(async move {
            while let Some((remote_repo, commit_entry, _dst, download_path)) = queue.try_pop() {
                log::debug!("worker[{worker}] processing task...");

                // Chunk and individual files
                let remote_path = &commit_entry.path;

                // Download to the tmp path, then copy over to the entries dir
                api::client::entries::download_large_entry(
                    &remote_repo,
                    remote_path,
                    &download_path,
                    &commit_entry.commit_id,
                    commit_entry.num_bytes,
                )
                .await?;

                progress_bar.add_bytes(commit_entry.num_bytes);
                progress_bar.add_files(1);
            }
            Ok(())
        });
        handles.push(handle);
    }
    let join_results = future::join_all(handles).await;
    for res in join_results {
        match res {
            Err(e) => return Err(OxenError::basic_str(format!("worker task panicked: {e}"))),
            Ok(Err(e)) => return Err(e),
            Ok(Ok(())) => {}
        }
    }

    log::debug!("All large file tasks done. :-)");

    Ok(())
}

async fn download_small_entries(
    remote_repo: &RemoteRepository,
    entries: Vec<CommitEntry>,
    dst: impl AsRef<Path>,
    progress_bar: &Arc<PullProgress>,
) -> Result<(), OxenError> {
    if entries.is_empty() {
        return Ok(());
    }

    let total_size = repositories::entries::compute_entries_size(&entries)?;

    // Compute num chunks
    let num_chunks = ((total_size / AVG_CHUNK_SIZE) + 1) as usize;

    let mut chunk_size = entries.len() / num_chunks;
    if num_chunks > entries.len() {
        chunk_size = entries.len();
    }

    log::debug!(
        "download_small_entries got {} missing entries",
        entries.len()
    );

    // Split into chunks, zip up, and post to server
    type PieceOfWork = (RemoteRepository, Vec<(String, PathBuf)>, PathBuf);
    type TaskQueue = deadqueue::limited::Queue<PieceOfWork>;

    let chunks: Vec<PieceOfWork> = entries
        .chunks(chunk_size)
        .map(|chunk| {
            let mut content_ids: Vec<(String, PathBuf)> = vec![];
            for e in chunk {
                content_ids.push((e.hash.clone(), e.path.clone()));
            }
            (remote_repo.to_owned(), content_ids, dst.as_ref().to_owned())
        })
        .collect();

    let worker_count = concurrency::num_threads_for_items(entries.len());
    let queue = Arc::new(TaskQueue::new(chunks.len()));
    for chunk in chunks {
        queue.try_push(chunk).unwrap();
    }
    let mut handles = vec![];

    for worker in 0..worker_count {
        let queue = queue.clone();
        let progress_bar = Arc::clone(progress_bar);
        let handle: tokio::task::JoinHandle<Result<(), OxenError>> = tokio::spawn(async move {
            while let Some((remote_repo, chunk, path)) = queue.try_pop() {
                log::debug!("worker[{worker}] processing task...");

                let download_size = api::client::entries::download_data_from_version_paths(
                    &remote_repo,
                    &chunk,
                    &path,
                )
                .await?;

                progress_bar.add_bytes(download_size);
                progress_bar.add_files(chunk.len() as u64);
            }
            Ok(())
        });
        handles.push(handle);
    }

    let join_results = future::join_all(handles).await;
    for res in join_results {
        match res {
            Err(e) => return Err(OxenError::basic_str(format!("worker task panicked: {e}"))),
            Ok(Err(e)) => return Err(e),
            Ok(Ok(())) => {}
        }
    }
    log::debug!("All tasks done. :-)");

    Ok(())
}

fn get_missing_entries_for_download(entries: &[CommitEntry], dst: &Path) -> Vec<CommitEntry> {
    let mut missing_entries: Vec<CommitEntry> = vec![];
    for commit_entry in entries {
        let working_path = dst.join(&commit_entry.path);
        if !working_path.exists() {
            missing_entries.push(commit_entry.to_owned())
        }
    }
    missing_entries
}

async fn get_missing_entries_for_pull(
    version_store: &Arc<dyn VersionStore>,
    entries: &[CommitEntry],
) -> Result<Vec<CommitEntry>, OxenError> {
    let mut missing_entries: Vec<CommitEntry> = vec![];
    // TODO: parallelize for S3
    for commit_entry in entries {
        if !version_store.version_exists(&commit_entry.hash).await? {
            missing_entries.push(commit_entry.to_owned())
        }
    }

    Ok(missing_entries)
}

fn working_dir_paths_from_large_entries(entries: &[CommitEntry], dst: &Path) -> Vec<PathBuf> {
    let mut paths: Vec<PathBuf> = vec![];
    for commit_entry in entries.iter() {
        let working_path = dst.join(&commit_entry.path);
        paths.push(working_path);
    }
    paths
}