path: root/makima/src/daemon/storage/patch.rs

//! Git patch creation and application for checkpoint recovery.

use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use std::io::{Read, Write};
use std::path::Path;
use thiserror::Error;
use tokio::process::Command;

/// Errors that can occur during patch operations.
#[derive(Error, Debug)]
pub enum PatchError {
    #[error("Git command failed: {0}")]
    GitCommand(String),

    #[error("Compression error: {0}")]
    Compression(#[from] std::io::Error),

    #[error("Patch too large: {size} bytes (max: {max} bytes)")]
    TooLarge { size: usize, max: usize },

    #[error("Empty patch (no changes)")]
    EmptyPatch,

    #[error("Failed to apply patch: {0}")]
    ApplyFailed(String),
}

/// Create a compressed git diff from worktree changes.
///
/// Generates a diff between `base_sha` and HEAD, then compresses it with gzip.
/// Returns the compressed patch bytes and the number of files changed.
pub async fn create_patch(
    worktree_path: &Path,
    base_sha: &str,
) -> Result<(Vec<u8>, usize), PatchError> {
    // Get the diff between base commit and HEAD
    let output = Command::new("git")
        .current_dir(worktree_path)
        .args(["diff", base_sha, "HEAD", "--binary"])
        .output()
        .await
        .map_err(|e| PatchError::GitCommand(format!("Failed to run git diff: {}", e)))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(PatchError::GitCommand(format!("git diff failed: {}", stderr)));
    }

    let diff_data = output.stdout;
    if diff_data.is_empty() {
        return Err(PatchError::EmptyPatch);
    }

    // Count files changed
    let files_output = Command::new("git")
        .current_dir(worktree_path)
        .args(["diff", base_sha, "HEAD", "--name-only"])
        .output()
        .await
        .map_err(|e| PatchError::GitCommand(format!("Failed to count files: {}", e)))?;

    let files_count = if files_output.status.success() {
        String::from_utf8_lossy(&files_output.stdout)
            .lines()
            .filter(|l| !l.is_empty())
            .count()
    } else {
        0
    };

    // Compress with gzip
    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
    encoder.write_all(&diff_data)?;
    let compressed = encoder.finish()?;

    Ok((compressed, files_count))
}

/// Apply a compressed patch to restore worktree state.
///
/// The worktree should already be checked out at `base_sha` before calling this.
pub async fn apply_patch(worktree_path: &Path, patch_data: &[u8]) -> Result<(), PatchError> {
    // Decompress gzip
    let mut decoder = GzDecoder::new(patch_data);
    let mut decompressed = Vec::new();
    decoder.read_to_end(&mut decompressed)?;

    if decompressed.is_empty() {
        return Err(PatchError::EmptyPatch);
    }

    // Apply the patch using git apply
    let mut child = Command::new("git")
        .current_dir(worktree_path)
        .args(["apply", "--binary", "-"])
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .map_err(|e| PatchError::GitCommand(format!("Failed to spawn git apply: {}", e)))?;

    // Write patch to stdin
    if let Some(mut stdin) = child.stdin.take() {
        use tokio::io::AsyncWriteExt;
        stdin.write_all(&decompressed).await?;
        drop(stdin); // Close stdin to signal EOF
    }

    let output = child
        .wait_with_output()
        .await
        .map_err(|e| PatchError::GitCommand(format!("Failed to wait for git apply: {}", e)))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(PatchError::ApplyFailed(stderr.to_string()));
    }

    Ok(())
}

/// Get the parent commit SHA (HEAD~1) from a worktree.
pub async fn get_parent_sha(worktree_path: &Path) -> Result<String, PatchError> {
    let output = Command::new("git")
        .current_dir(worktree_path)
        .args(["rev-parse", "HEAD~1"])
        .output()
        .await
        .map_err(|e| PatchError::GitCommand(format!("Failed to get parent SHA: {}", e)))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(PatchError::GitCommand(format!(
            "git rev-parse HEAD~1 failed: {}",
            stderr
        )));
    }

    Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
}

/// Checkout a specific commit in the worktree.
pub async fn checkout_commit(worktree_path: &Path, sha: &str) -> Result<(), PatchError> {
    let output = Command::new("git")
        .current_dir(worktree_path)
        .args(["checkout", sha])
        .output()
        .await
        .map_err(|e| PatchError::GitCommand(format!("Failed to checkout: {}", e)))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(PatchError::GitCommand(format!(
            "git checkout {} failed: {}",
            sha, stderr
        )));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    async fn setup_test_repo() -> TempDir {
        let dir = TempDir::new().unwrap();
        let path = dir.path();

        // Initialize git repo
        Command::new("git")
            .current_dir(path)
            .args(["init"])
            .output()
            .await
            .unwrap();

        // Configure git user
        Command::new("git")
            .current_dir(path)
            .args(["config", "user.email", "test@test.com"])
            .output()
            .await
            .unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["config", "user.name", "Test"])
            .output()
            .await
            .unwrap();

        // Create initial commit
        fs::write(path.join("file.txt"), "initial").unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["add", "."])
            .output()
            .await
            .unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["commit", "-m", "initial"])
            .output()
            .await
            .unwrap();

        dir
    }

    #[tokio::test]
    async fn test_create_and_apply_patch() {
        let dir = setup_test_repo().await;
        let path = dir.path();

        // Get base SHA
        let base_sha = get_parent_sha(path).await;
        // This will fail since there's only one commit
        assert!(base_sha.is_err());

        // Make another commit first
        fs::write(path.join("file.txt"), "modified").unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["add", "."])
            .output()
            .await
            .unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["commit", "-m", "modified"])
            .output()
            .await
            .unwrap();

        // Now get the base SHA
        let base_sha = get_parent_sha(path).await.unwrap();

        // Create patch
        let (patch_data, files_count) = create_patch(path, &base_sha).await.unwrap();
        assert!(!patch_data.is_empty());
        assert_eq!(files_count, 1);

        // Reset to base and apply patch
        checkout_commit(path, &base_sha).await.unwrap();
        assert_eq!(fs::read_to_string(path.join("file.txt")).unwrap(), "initial");

        apply_patch(path, &patch_data).await.unwrap();
        assert_eq!(
            fs::read_to_string(path.join("file.txt")).unwrap(),
            "modified"
        );
    }

    #[tokio::test]
    async fn test_empty_patch() {
        let dir = setup_test_repo().await;
        let path = dir.path();

        // Make another commit
        fs::write(path.join("file.txt"), "modified").unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["add", "."])
            .output()
            .await
            .unwrap();
        Command::new("git")
            .current_dir(path)
            .args(["commit", "-m", "modified"])
            .output()
            .await
            .unwrap();

        // Get current HEAD
        let head_output = Command::new("git")
            .current_dir(path)
            .args(["rev-parse", "HEAD"])
            .output()
            .await
            .unwrap();
        let head_sha = String::from_utf8_lossy(&head_output.stdout)
            .trim()
            .to_string();

        // Try to create patch from HEAD to HEAD (no changes)
        let result = create_patch(path, &head_sha).await;
        assert!(matches!(result, Err(PatchError::EmptyPatch)));
    }
}