void_core/workspace/

checkout.rs

//! Checkout module - restore files from a commit to the working tree
//!
//! Provides functions for checking out files from encrypted commits:
//! - `checkout_tree`: Restore entire tree from a commit
//! - `checkout_paths`: Restore specific paths from a commit
//!
//! Uses parallel file restoration via rayon for performance.

use std::collections::{HashMap, HashSet};
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::SystemTime;

use camino::Utf8PathBuf;
use rayon::prelude::*;

use void_crypto::WrappedKey;

use crate::crypto::{CommitReader, ContentKey, KeyVault, SecretKey};
use crate::staged;
use crate::index::{
    entry_matches_file, read_index, write_workspace_index, IndexEntry, WorkspaceIndex,
};
use crate::metadata::{ManifestEntry, ShardReference};
use crate::metadata::manifest_tree::TreeManifest;
use crate::pathspec::Pathspec;

use crate::store::ObjectStoreExt;
use crate::support::events::{emit_workspace, VoidObserver, WorkspaceEvent};
use crate::{cid, ContentHash, Result, VoidError};

/// Options for checkout operations
#[derive(Clone)]
pub struct CheckoutOptions {
    /// Specific paths to checkout (None = full tree)
    pub paths: Option<Vec<String>>,
    /// Overwrite modified files without prompting
    pub force: bool,
    /// Optional observer for progress events.
    pub observer: Option<Arc<dyn VoidObserver>>,
    /// Per-workspace state directory for index operations.
    /// When `None`, defaults to `workspace.join(".void")` (main workspace).
    pub workspace_dir: Option<PathBuf>,
    /// Include large (chunked) files in full-tree checkout.
    /// When false (default), full-tree checkout skips files with shard_count > 1.
    /// Path-specific checkout always includes large files regardless.
    pub include_large: bool,
}

impl Default for CheckoutOptions {
    fn default() -> Self {
        Self {
            paths: None,
            force: false,
            observer: None,
            workspace_dir: None,
            include_large: false,
        }
    }
}

/// Statistics from checkout operation
#[derive(Clone, Debug, Default)]
pub struct CheckoutStats {
    /// Number of files restored to the workspace
    pub files_restored: usize,
    /// Total bytes written to disk
    pub bytes_written: u64,
    /// Number of files skipped (already up-to-date)
    pub files_skipped: usize,
    /// Number of shards read
    pub shards_read: usize,
    /// Number of large files deferred (not checked out)
    pub files_deferred: usize,
}

/// A file to restore during checkout.
///
/// Each entry maps a file path to its manifest entry and shard info.
#[derive(Clone, Debug)]
pub struct FileToRestore {
    /// The manifest entry with path, content_hash, offset, length, shard_index
    pub entry: ManifestEntry,
    /// Shard CID (from manifest.shards)
    pub shard_cid: void_crypto::ShardCid,
    /// Wrapped shard key (if shard uses per-shard encryption)
    pub wrapped_key: Option<WrappedKey>,
}

/// Helper struct to hold loaded commit/manifest info
struct CommitInfo {
    manifest: TreeManifest,
    /// All files in this commit with their content hashes (path -> content_hash)
    all_files: HashMap<String, ContentHash>,
    /// Commit reader for shard decryption
    reader: CommitReader,
    /// Ancestor content keys for shard decryption fallback
    ancestor_keys: Vec<ContentKey>,
}

struct CheckoutPlan {
    files: Vec<FileToRestore>,
    /// Non-materialized index entries for deferred large files.
    deferred_entries: Vec<IndexEntry>,
    index: Option<WorkspaceIndex>,
    pathspec: Pathspec,
    /// Number of large files skipped during full-tree checkout.
    files_deferred: usize,
}

/// Loads commit and manifest, returning deserialized info.
fn load_commit_info<S: ObjectStoreExt>(
    store: &S,
    vault: &KeyVault,
    commit_cid: &crate::VoidCid,
) -> Result<CommitInfo> {
    let commit_encrypted: void_crypto::EncryptedCommit = store.get_blob(commit_cid)?;
    let (commit_bytes, reader) = CommitReader::open_with_vault(vault, &commit_encrypted)?;
    let commit = commit_bytes.parse()?;

    let manifest = TreeManifest::from_commit(store, &commit, &reader)?
        .ok_or_else(|| VoidError::IntegrityError {
            expected: "manifest_cid present on commit".into(),
            actual: "None".into(),
        })?;

    // Collect all files map from manifest
    let mut all_files = HashMap::new();
    for entry_result in manifest.iter() {
        let entry = entry_result?;
        all_files.insert(entry.path.clone(), entry.content_hash);
    }

    let ancestor_keys = crate::crypto::collect_ancestor_content_keys_vault(vault, store, &commit);

    Ok(CommitInfo {
        manifest,
        all_files,
        reader,
        ancestor_keys,
    })
}

/// Determines which files need to be restored and groups them by shard.
fn plan_checkout(
    vault: &KeyVault,
    commit_info: &CommitInfo,
    workspace: &Path,
    options: &CheckoutOptions,
) -> Result<CheckoutPlan> {
    // Build pathspec matcher if paths are specified
    let pathspec = match &options.paths {
        Some(paths) => {
            let path_refs: Vec<&str> = paths.iter().map(|s| s.as_str()).collect();
            Pathspec::new(&path_refs)?
        }
        None => Pathspec::new(&[])?, // matches all
    };

    // Try to load existing index for dirty detection
    let void_dir = options.workspace_dir.clone().unwrap_or_else(|| workspace.join(".void"));
    let existing_index = if void_dir.exists() {
        read_index(&void_dir, vault.index_key()?).ok()
    } else {
        None
    };
    let base_path = Utf8PathBuf::try_from(workspace.to_path_buf())
        .map_err(|e| VoidError::Io(std::io::Error::new(std::io::ErrorKind::InvalidData, e)))?;

    let shards = commit_info.manifest.shards();
    let mut files_to_restore = Vec::new();
    let mut deferred_entries = Vec::new();
    let mut files_deferred = 0usize;
    let is_path_specific = options.paths.is_some();

    for entry_result in commit_info.manifest.iter() {
        let entry = entry_result?;

        // Check if path matches filter
        if !pathspec.matches(&entry.path) {
            emit_workspace(
                &options.observer,
                WorkspaceEvent::FileSkipped {
                    path: entry.path.clone(),
                    reason: "does not match pathspec".to_string(),
                },
            );
            continue;
        }

        // Check if file is dirty (modified in workspace)
        if !options.force {
            if let Some(ref index) = existing_index {
                if let Some(idx_entry) = index.get(&entry.path) {
                    let file_path_on_disk = crate::util::safe_join(workspace, &entry.path)?;
                    if file_path_on_disk.exists() {
                        let matches = entry_matches_file(idx_entry, &base_path).unwrap_or(false);
                        if !matches {
                            return Err(VoidError::Shard(format!(
                                "file '{}' has local modifications; use --force to overwrite",
                                entry.path
                            )));
                        }
                    }
                }
            }
        }

        // Skip large (chunked) files during full-tree checkout unless explicitly requested
        if entry.shard_count > 1 && !is_path_specific && !options.include_large {
            files_deferred += 1;
            deferred_entries.push(IndexEntry::new_remote(
                entry.path.clone(),
                entry.content_hash,
                entry.size,
            ));
            emit_workspace(
                &options.observer,
                WorkspaceEvent::FileSkipped {
                    path: entry.path.clone(),
                    reason: "large file (use --include-large or checkout by path)".to_string(),
                },
            );
            continue;
        }

        // Get shard info from manifest
        let shard_ref = shards.get(entry.shard_index as usize)
            .ok_or_else(|| VoidError::Shard(format!(
                "shard_index {} out of range for file '{}'", entry.shard_index, entry.path
            )))?;

        files_to_restore.push(FileToRestore {
            entry,
            shard_cid: shard_ref.cid.clone(),
            wrapped_key: shard_ref.wrapped_key.clone(),
        });
    }

    Ok(CheckoutPlan {
        files: files_to_restore,
        deferred_entries,
        index: existing_index,
        pathspec,
        files_deferred,
    })
}

fn prune_extra_files(
    workspace: &Path,
    pathspec: &Pathspec,
    target_set: &HashSet<String>,
    existing_index: Option<&WorkspaceIndex>,
    force: bool,
) -> Result<Vec<String>> {
    let Some(index) = existing_index else {
        return Ok(Vec::new());
    };

    let base_path = Utf8PathBuf::try_from(workspace.to_path_buf())
        .map_err(|e| VoidError::Io(std::io::Error::new(std::io::ErrorKind::InvalidData, e)))?;

    let mut removed = Vec::new();

    for entry in &index.entries {
        if !pathspec.matches(&entry.path) {
            continue;
        }
        if target_set.contains(&entry.path) {
            continue;
        }

        let file_path = crate::util::safe_join(workspace, &entry.path)?;
        if file_path.exists() {
            if !force {
                let matches = entry_matches_file(entry, &base_path).unwrap_or(false);
                if !matches {
                    return Err(VoidError::Shard(format!(
                        "file '{}' has local modifications; use --force to overwrite",
                        entry.path
                    )));
                }
            }
            fs::remove_file(&file_path)?;
        }

        removed.push(entry.path.clone());
    }

    Ok(removed)
}

/// Checkout entire tree from a commit.
pub fn checkout_tree<S: ObjectStoreExt + Sync>(
    store: &S,
    vault: &KeyVault,
    commit_cid: &crate::VoidCid,
    workspace: &Path,
    options: &CheckoutOptions,
) -> Result<CheckoutStats> {
    let commit_info = load_commit_info(store, vault, commit_cid)?;
    let CheckoutPlan { files: plan_files, deferred_entries, index: plan_index, pathspec: plan_pathspec, files_deferred } =
        plan_checkout(vault, &commit_info, workspace, options)?;
    let target_set: HashSet<String> = commit_info
        .all_files
        .keys()
        .filter(|path| plan_pathspec.matches(path))
        .cloned()
        .collect();

    let removed_paths = prune_extra_files(
        workspace,
        &plan_pathspec,
        &target_set,
        plan_index.as_ref(),
        options.force,
    )?;

    let void_dir = options.workspace_dir.clone().unwrap_or_else(|| workspace.join(".void"));

    // Separate small files from chunked files
    let (small_files, chunked_files): (Vec<FileToRestore>, Vec<FileToRestore>) = plan_files
        .into_iter()
        .partition(|f| f.entry.shard_count <= 1);

    let (mut stats, mut restored_entries) =
        restore_files(store, &commit_info.reader, vault.staged_key()?, &commit_info.ancestor_keys, workspace, &small_files, &options.observer, Some(&void_dir))?;

    // Restore chunked files sequentially (one shard in memory at a time)
    let manifest_shards = commit_info.manifest.shards();
    for file_to_restore in &chunked_files {
        let (bytes, index_entry) = restore_chunked_file(
            store,
            &commit_info.reader,
            &commit_info.ancestor_keys,
            workspace,
            &file_to_restore.entry,
            manifest_shards,
        )?;
        stats.files_restored += 1;
        stats.bytes_written += bytes;
        stats.shards_read += file_to_restore.entry.shard_count as usize;
        restored_entries.push(index_entry);

        emit_workspace(
            &options.observer,
            WorkspaceEvent::FileCheckedOut {
                path: file_to_restore.entry.path.clone(),
            },
        );
    }
    if void_dir.exists() {
        let index = if options.paths.is_some() {
            let mut index = plan_index.unwrap_or_else(WorkspaceIndex::empty);
            if !removed_paths.is_empty() {
                let removed: HashSet<String> = removed_paths.into_iter().collect();
                index.entries.retain(|entry| !removed.contains(&entry.path));
            }
            for entry in restored_entries {
                index.upsert_entry(entry);
            }
            index
        } else {
            // Full checkout: build complete index from ALL commit files
            let mut all_entries = restored_entries;
            // Clone paths into owned HashSet to avoid borrow conflict
            let restored_paths: HashSet<String> = all_entries.iter().map(|e| e.path.clone()).collect();

            // Add entries for files that weren't restored (already existed with matching content)
            for (path, content_hash) in &commit_info.all_files {
                if !restored_paths.contains(path.as_str()) {
                    // File was skipped, create index entry from disk metadata
                    let file_path = crate::util::safe_join(workspace, path)?;
                    if file_path.exists() {
                        if let Ok(metadata) = std::fs::metadata(&file_path) {
                            let (mtime_secs, mtime_nanos) = metadata
                                .modified()
                                .ok()
                                .and_then(|mtime| mtime.duration_since(SystemTime::UNIX_EPOCH).ok())
                                .map(|dur| (dur.as_secs(), dur.subsec_nanos()))
                                .unwrap_or((0, 0));

                            let entry = IndexEntry::new(
                                path.clone(),
                                *content_hash,
                                mtime_secs,
                                mtime_nanos,
                                metadata.len(),
                            );
                            all_entries.push(entry);

                            // Write staged blob for skipped files too
                            if !staged::has_staged_blob(&void_dir, content_hash) {
                                let content = fs::read(&file_path)?;
                                staged::write_staged_blob(&void_dir, vault.staged_key()?, content_hash, &content)?;
                            }
                        }
                    }
                }
            }

            // Add non-materialized entries for deferred large files
            all_entries.extend(deferred_entries);

            WorkspaceIndex::new(Some(void_crypto::CommitCid::from_bytes(commit_cid.to_bytes())), all_entries)
        };

        write_workspace_index(&void_dir, vault.index_key()?, &index)?;
    }

    let mut stats = stats;
    stats.files_deferred = files_deferred;
    Ok(stats)
}

/// Checkout specific paths from a commit.
pub fn checkout_paths<S: ObjectStoreExt + Sync>(
    store: &S,
    vault: &KeyVault,
    commit_cid: &crate::VoidCid,
    workspace: &Path,
    paths: &[String],
) -> Result<CheckoutStats> {
    let options = CheckoutOptions {
        paths: Some(paths.to_vec()),
        force: true, // Path-based checkout always overwrites
        observer: None,
        workspace_dir: None,
        include_large: false,
    };

    checkout_tree(store, vault, commit_cid, workspace, &options)
}

/// Restores files to the workspace using parallel processing.
///
/// Uses manifest entries for file offset/length within shards (ShardBody::read_file).
/// When `staged_target` is provided, also writes staged blobs for each file
/// so that subsequent commits can find them via `seal_index()`.
pub fn restore_files<S: ObjectStoreExt + Sync>(
    store: &S,
    reader: &CommitReader,
    staged_key: &SecretKey,
    ancestor_keys: &[ContentKey],
    workspace: &Path,
    files: &[FileToRestore],
    observer: &Option<Arc<dyn VoidObserver>>,
    staged_target: Option<&Path>,
) -> Result<(CheckoutStats, Vec<IndexEntry>)> {
    if files.is_empty() {
        return Ok((CheckoutStats::default(), Vec::new()));
    }

    let total_files = files.len() as u64;

    // Emit initial progress
    emit_workspace(
        observer,
        WorkspaceEvent::Progress {
            stage: "checkout".to_string(),
            current: 0,
            total: total_files,
        },
    );

    // Group files by shard_index to minimize shard reads
    let unique_shards: HashSet<u32> = files.iter().map(|f| f.entry.shard_index).collect();

    // Atomic counters for parallel stats collection
    let files_restored = AtomicUsize::new(0);
    let bytes_written = AtomicU64::new(0);
    let shards_read = AtomicUsize::new(unique_shards.len());

    // Collect index entries for updating
    let index_entries = std::sync::Mutex::new(Vec::new());

    // Process files in parallel, grouped by shard
    let results: Result<Vec<()>> = unique_shards
        .par_iter()
        .map(|shard_index| {
            // Find all files in this shard
            let shard_files: Vec<_> = files.iter().filter(|f| f.entry.shard_index == *shard_index).collect();

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

            // Get shard CID from first file (all files in same shard have same CID)
            let shard_cid = cid::from_bytes(shard_files[0].shard_cid.as_bytes())?;

            // Fetch, decrypt, and decompress shard
            let shard_encrypted: void_crypto::EncryptedShard = store.get_blob(&shard_cid)?;
            let shard_bytes = reader.decrypt_shard(&shard_encrypted, shard_files[0].wrapped_key.as_ref(), ancestor_keys)?;
            let body = shard_bytes.decompress()?;

            // Restore each file from this shard using manifest offsets
            for file_info in shard_files {
                let content = body.read_file(&file_info.entry)?;
                let file_path = crate::util::safe_join(workspace, &file_info.entry.path)?;

                // Create parent directories
                if let Some(parent) = file_path.parent() {
                    fs::create_dir_all(parent)?;
                }

                // Write file
                let mut file = File::create(&file_path)?;
                file.write_all(&content)?;

                let content_len = content.len();

                // Create index entry
                let content_hash = ContentHash::digest(&content);

                // Write staged blob so seal_index() can find it during commit
                if let Some(target) = staged_target {
                    staged::write_staged_blob(target, staged_key, &content_hash, &content)?;
                }

                let metadata = fs::metadata(&file_path)?;
                let (mtime_secs, mtime_nanos) = metadata
                    .modified()
                    .ok()
                    .and_then(|mtime| mtime.duration_since(SystemTime::UNIX_EPOCH).ok())
                    .map(|dur| (dur.as_secs(), dur.subsec_nanos()))
                    .unwrap_or((0, 0));

                let entry = IndexEntry::new(
                    file_info.entry.path.clone(),
                    content_hash,
                    mtime_secs,
                    mtime_nanos,
                    content_len as u64,
                );

                index_entries
                    .lock()
                    .map_err(|_| VoidError::Shard("index lock poisoned".into()))?
                    .push(entry);

                let restored_count = files_restored.fetch_add(1, Ordering::Relaxed) + 1;
                bytes_written.fetch_add(content_len as u64, Ordering::Relaxed);

                // Emit file checked out event
                emit_workspace(
                    observer,
                    WorkspaceEvent::FileCheckedOut {
                        path: file_info.entry.path.clone(),
                    },
                );

                // Emit progress event
                emit_workspace(
                    observer,
                    WorkspaceEvent::Progress {
                        stage: "checkout".to_string(),
                        current: restored_count as u64,
                        total: total_files,
                    },
                );
            }

            Ok(())
        })
        .collect();

    results?;

    let entries = index_entries
        .into_inner()
        .map_err(|_| VoidError::Shard("failed to get index entries".into()))?;

    Ok((
        CheckoutStats {
            files_restored: files_restored.load(Ordering::Relaxed),
            bytes_written: bytes_written.load(Ordering::Relaxed),
            files_skipped: 0,
            shards_read: shards_read.load(Ordering::Relaxed),
            files_deferred: 0,
        },
        entries,
    ))
}

/// Restore a single chunked file by fetching its shards sequentially.
///
/// Each chunk shard is fetched, decrypted, decompressed, and written to disk
/// in order. Only one shard is held in memory at a time. The content hash is
/// verified incrementally after all chunks are written.
fn restore_chunked_file<S: ObjectStoreExt>(
    store: &S,
    reader: &CommitReader,
    ancestor_keys: &[ContentKey],
    workspace: &Path,
    entry: &ManifestEntry,
    shards: &[ShardReference],
) -> Result<(u64, IndexEntry)> {
    use sha2::{Digest, Sha256};

    let file_path = crate::util::safe_join(workspace, &entry.path)?;
    if let Some(parent) = file_path.parent() {
        fs::create_dir_all(parent)?;
    }

    let output = File::create(&file_path)?;
    let mut writer = BufWriter::new(output);
    let mut hasher = Sha256::new();
    let mut total_written = 0u64;

    let start = entry.shard_index as usize;
    let end = start + entry.shard_count as usize;

    for shard_idx in start..end {
        let shard_ref = shards.get(shard_idx).ok_or_else(|| {
            VoidError::Shard(format!(
                "chunk shard index {} out of range for '{}'",
                shard_idx, entry.path
            ))
        })?;

        let shard_cid = cid::from_bytes(shard_ref.cid.as_bytes())?;
        let encrypted: void_crypto::EncryptedShard = store.get_blob(&shard_cid)?;
        let decrypted = reader.decrypt_shard(&encrypted, shard_ref.wrapped_key.as_ref(), ancestor_keys)?;
        let body = decrypted.decompress()?;

        // Each chunk shard's decompressed body IS the chunk content.
        // Use as_bytes() directly instead of read_file() because
        // entry.length is the full file size, not the chunk size.
        let chunk = body.as_bytes();
        writer.write_all(&chunk)?;
        hasher.update(&chunk);
        total_written += chunk.len() as u64;
    }

    writer.flush()?;

    // Verify content hash
    let computed_hash = ContentHash::from_bytes(hasher.finalize().into());
    if computed_hash != entry.content_hash {
        return Err(VoidError::IntegrityError {
            expected: entry.content_hash.to_hex(),
            actual: computed_hash.to_hex(),
        });
    }

    // Build index entry from disk metadata
    let metadata = fs::metadata(&file_path)?;
    let (mtime_secs, mtime_nanos) = metadata
        .modified()
        .ok()
        .and_then(|mtime| mtime.duration_since(SystemTime::UNIX_EPOCH).ok())
        .map(|dur| (dur.as_secs(), dur.subsec_nanos()))
        .unwrap_or((0, 0));

    let index_entry = IndexEntry::new(
        entry.path.clone(),
        entry.content_hash,
        mtime_secs,
        mtime_nanos,
        total_written,
    );

    Ok((total_written, index_entry))
}