ralph/

checkpoint.rs

use crate::config::CheckpointsConfig;
use crate::errors::{RalphError, Result};
use crate::providers::Message;
use crate::session::Session;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointMeta {
    pub id: String,
    pub name: String,
    pub turn: u32,
    pub created_at: DateTime<Utc>,
    pub modified_files: Vec<String>,
    /// Git commit SHA, set when git_commit is enabled and file count exceeds threshold.
    #[serde(default)]
    pub git_sha: Option<String>,
}

pub struct Checkpoint {
    pub meta: CheckpointMeta,
    pub messages: Vec<Message>,
    pub dir: PathBuf,
}

impl Checkpoint {
    /// Create a checkpoint at the current session state.
    pub async fn create(session: &Session, name: &str, config: &CheckpointsConfig) -> Result<Self> {
        let seq = next_seq(session);
        let id = format!("{:03}_{}", seq, sanitize_name(name));
        let dir = session.checkpoints_dir().join(&id);
        std::fs::create_dir_all(&dir)?;

        // Snapshot all modified files
        let files_dir = dir.join("files");
        std::fs::create_dir_all(&files_dir)?;

        let mut snapshotted_files = Vec::new();
        for file_path in &session.modified_files {
            if !file_path.exists() {
                continue;
            }
            let rel = file_path
                .strip_prefix(&session.meta.workspace)
                .unwrap_or(file_path);
            let dest = files_dir.join(rel);
            if let Some(parent) = dest.parent() {
                std::fs::create_dir_all(parent)?;
            }
            std::fs::copy(file_path, &dest)?;
            snapshotted_files.push(rel.display().to_string());
        }

        // Save message history at this point
        let messages_json = serde_json::to_string_pretty(&session.messages)?;
        std::fs::write(dir.join("messages.json"), messages_json)?;

        // Optionally commit to git when above the threshold
        let git_sha = if config.git_commit && snapshotted_files.len() > config.git_commit_threshold
        {
            let workspace = PathBuf::from(&session.meta.workspace);
            try_git_commit(
                &workspace,
                &snapshotted_files,
                name,
                session.meta.turn_count,
            )
            .await
        } else {
            None
        };

        let mut meta = CheckpointMeta {
            id: id.clone(),
            name: name.to_string(),
            turn: session.meta.turn_count,
            created_at: Utc::now(),
            modified_files: snapshotted_files,
            git_sha: None,
        };
        meta.git_sha = git_sha;

        let meta_json = serde_json::to_string_pretty(&meta)?;
        std::fs::write(dir.join("checkpoint.json"), meta_json)?;

        if let Some(ref sha) = meta.git_sha {
            println!(
                "[checkpoint] Created '{}' at turn {} ({}) — git commit {}",
                name,
                session.meta.turn_count,
                meta.created_at.format("%H:%M:%S"),
                &sha[..8.min(sha.len())]
            );
        } else {
            println!(
                "[checkpoint] Created '{}' at turn {} ({})",
                name,
                session.meta.turn_count,
                meta.created_at.format("%H:%M:%S")
            );
        }

        let messages = session.messages.clone();
        Ok(Self {
            meta,
            messages,
            dir,
        })
    }

    /// List all checkpoints for a session, sorted by sequence number.
    pub fn list(session: &Session) -> Vec<CheckpointMeta> {
        let cp_dir = session.checkpoints_dir();
        if !cp_dir.exists() {
            return Vec::new();
        }

        let mut checkpoints: Vec<CheckpointMeta> = std::fs::read_dir(&cp_dir)
            .unwrap_or_else(|_| std::fs::read_dir(".").unwrap())
            .flatten()
            .filter(|e| e.path().is_dir())
            .filter_map(|e| {
                let meta_path = e.path().join("checkpoint.json");
                let content = std::fs::read_to_string(&meta_path).ok()?;
                serde_json::from_str(&content).ok()
            })
            .collect();

        checkpoints.sort_by_key(|c| c.id.clone());
        checkpoints
    }

    /// Find a checkpoint by name or sequence number (1-based).
    pub fn find(session: &Session, name_or_num: &str) -> Result<Self> {
        let all = Self::list(session);

        let meta = if let Ok(num) = name_or_num.parse::<usize>() {
            all.into_iter().nth(num - 1)
        } else {
            all.into_iter()
                .find(|c| c.name == name_or_num || c.id == name_or_num)
        }
        .ok_or_else(|| RalphError::CheckpointNotFound(name_or_num.to_string()))?;

        let dir = session.checkpoints_dir().join(&meta.id);
        let messages_json =
            std::fs::read_to_string(dir.join("messages.json")).unwrap_or_else(|_| "[]".to_string());
        let messages: Vec<Message> = serde_json::from_str(&messages_json).unwrap_or_default();

        Ok(Self {
            meta,
            messages,
            dir,
        })
    }

    /// Revert the session to this checkpoint.
    /// - If `files_only` is false, also rolls back the message history.
    /// - Always creates a `before-revert-<timestamp>` safety checkpoint first.
    pub async fn revert_into(
        &self,
        session: &mut Session,
        files_only: bool,
        printer: &crate::output::Printer,
    ) -> Result<()> {
        let workspace = PathBuf::from(&session.meta.workspace);

        // Show what will change
        self.show_revert_preview(session, files_only, printer);

        // Require confirmation — not bypassable
        let confirm = printer.confirm(
            &format!(
                "Reverting to checkpoint '{}' (turn {}). This cannot be skipped.",
                self.meta.name, self.meta.turn
            ),
            true,
            false,
        );
        if confirm != crate::output::ConfirmResult::Yes {
            return Err(RalphError::UserAborted);
        }

        // Always create a safety checkpoint before reverting (never git-commits)
        let safety_name = format!("before-revert-{}", Utc::now().format("%H%M%S"));
        let _ = Checkpoint::create(session, &safety_name, &CheckpointsConfig::default()).await;

        // Build the set of files that existed at checkpoint time
        let cp_files: std::collections::HashSet<String> =
            self.meta.modified_files.iter().cloned().collect();

        // Delete files created after the checkpoint (in session.modified_files but not in checkpoint)
        for file_path in &session.modified_files {
            let rel = file_path
                .strip_prefix(&workspace)
                .unwrap_or(file_path)
                .to_string_lossy()
                .to_string();
            if !cp_files.contains(&rel) && file_path.exists() {
                std::fs::remove_file(file_path)
                    .map_err(|e| RalphError::CheckpointRevertFailed(e.to_string()))?;
            }
        }

        // Restore snapshotted files
        let files_dir = self.dir.join("files");
        if files_dir.exists() {
            for entry in walkdir_flat(&files_dir) {
                let rel = entry.strip_prefix(&files_dir).unwrap_or(&entry);
                let dest = workspace.join(rel);
                if let Some(parent) = dest.parent() {
                    std::fs::create_dir_all(parent)?;
                }
                std::fs::copy(&entry, &dest)
                    .map_err(|e| RalphError::CheckpointRevertFailed(e.to_string()))?;
            }
        }

        if !files_only {
            session.messages = self.messages.clone();
            session.meta.turn_count = self.meta.turn;
            session.flush()?;
        }

        println!(
            "[checkpoint] Reverted to '{}'. Safety checkpoint '{}' was created.",
            self.meta.name, safety_name
        );
        Ok(())
    }

    fn show_revert_preview(
        &self,
        session: &Session,
        files_only: bool,
        printer: &crate::output::Printer,
    ) {
        use colored::Colorize;
        let workspace = PathBuf::from(&session.meta.workspace);

        println!(
            "\n[checkpoint] Reverting to '{}' (turn {}, {})\n",
            self.meta.name,
            self.meta.turn,
            self.meta.created_at.format("%Y-%m-%d %H:%M:%S")
        );

        // Files to restore
        println!("  Files to restore:");
        for f in &self.meta.modified_files {
            let dest = workspace.join(f);
            if dest.exists() {
                println!("    {} {}", "~".yellow(), f);
            } else {
                println!("    {} {} (will be created)", "+".green(), f);
            }
        }

        // Files created after checkpoint that will be removed
        let cp_files: std::collections::HashSet<&str> = self
            .meta
            .modified_files
            .iter()
            .map(|s| s.as_str())
            .collect();
        for file_path in &session.modified_files {
            let rel = file_path
                .strip_prefix(&workspace)
                .unwrap_or(file_path)
                .to_string_lossy()
                .to_string();
            if !cp_files.contains(rel.as_str()) && file_path.exists() {
                println!(
                    "    {} {} (delete — did not exist at checkpoint)",
                    "-".red(),
                    rel
                );
            }
        }

        if !files_only {
            let turns_to_drop = session.meta.turn_count.saturating_sub(self.meta.turn);
            println!(
                "\n  Conversation: {} turn(s) will be discarded",
                turns_to_drop
            );
        } else {
            println!("\n  Conversation history: unchanged (--files-only)");
        }
        println!();
    }
}

/// Stage the given workspace-relative file paths and commit on the current
/// branch. Returns the full commit SHA on success, None on any failure
/// (missing git, nothing to commit, etc.). Never switches branches.
async fn try_git_commit(
    workspace: &Path,
    files: &[String],
    name: &str,
    turn: u32,
) -> Option<String> {
    if !workspace.join(".git").exists() {
        return None;
    }

    // Stage the modified files
    let mut add = tokio::process::Command::new("git");
    add.arg("add").arg("--").current_dir(workspace);
    for f in files {
        add.arg(f);
    }
    let add_ok = add
        .output()
        .await
        .map(|o| o.status.success())
        .unwrap_or(false);
    if !add_ok {
        eprintln!("[checkpoint] Warning: git add failed — checkpoint saved without git commit.");
        return None;
    }

    // Nothing to commit? (`git diff --cached --quiet` exits 0 when clean)
    let clean = tokio::process::Command::new("git")
        .args(["diff", "--cached", "--quiet"])
        .current_dir(workspace)
        .output()
        .await
        .ok()?;
    if clean.status.success() {
        return None; // nothing staged
    }

    // Commit on whichever branch the user is on
    let msg = format!("ralph checkpoint: {} (turn {})", name, turn);
    let commit = tokio::process::Command::new("git")
        .args(["commit", "-m", &msg])
        .current_dir(workspace)
        .output()
        .await
        .ok()?;
    if !commit.status.success() {
        let stderr = String::from_utf8_lossy(&commit.stderr);
        eprintln!(
            "[checkpoint] Warning: git commit failed: {} — checkpoint saved without git commit.",
            stderr.trim()
        );
        return None;
    }

    // Return the SHA
    let sha_out = tokio::process::Command::new("git")
        .args(["rev-parse", "HEAD"])
        .current_dir(workspace)
        .output()
        .await
        .ok()?;
    let sha = String::from_utf8_lossy(&sha_out.stdout).trim().to_string();
    if sha.is_empty() {
        None
    } else {
        Some(sha)
    }
}

fn next_seq(session: &Session) -> usize {
    Checkpoint::list(session).len() + 1
}

fn sanitize_name(name: &str) -> String {
    name.chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '-' {
                c
            } else {
                '_'
            }
        })
        .collect()
}

fn walkdir_flat(dir: &Path) -> Vec<PathBuf> {
    let mut files = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                files.extend(walkdir_flat(&path));
            } else {
                files.push(path);
            }
        }
    }
    files
}