path: root/makima/src/daemon/task/context_recovery.rs

//! Context recovery for task resumption and restart.
//!
//! This module provides functionality to rebuild context when tasks resume or restart,
//! ensuring Claude can quickly orient itself with the current state of the worktree.

use std::path::Path;
use std::process::Command;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// Summary of git status information.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GitStatusSummary {
    /// Number of staged files.
    pub staged: u32,
    /// Number of modified (unstaged) files.
    pub modified: u32,
    /// Number of untracked files.
    pub untracked: u32,
    /// Number of deleted files.
    pub deleted: u32,
    /// Brief description (e.g., "3 modified, 1 staged").
    pub summary: String,
}

/// Information about a checkpoint (git commit).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointInfo {
    /// Commit SHA (short form).
    pub sha: String,
    /// Commit message (first line).
    pub message: String,
    /// Commit timestamp.
    pub timestamp: DateTime<Utc>,
}

/// A single progress entry from the progress log.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProgressEntry {
    /// Timestamp of the entry.
    pub timestamp: DateTime<Utc>,
    /// Entry type (e.g., "ITERATION", "TASK_START", "CHECKPOINT").
    pub entry_type: String,
    /// Entry message or summary.
    pub message: String,
}

/// Context recovery data for a task.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextRecovery {
    /// Current git branch name.
    pub current_branch: String,
    /// Git status summary.
    pub git_status: GitStatusSummary,
    /// Last checkpoint information (most recent commit).
    pub last_checkpoint: Option<CheckpointInfo>,
    /// Recent progress entries from progress.log.
    pub recent_progress: Vec<ProgressEntry>,
    /// Current phase of the task.
    pub current_phase: String,
}

impl ContextRecovery {
    /// Format the context recovery as a markdown header suitable for injection into prompts.
    pub fn to_markdown(&self) -> String {
        let mut output = String::new();
        output.push_str("## Context Recovery\n");
        output.push_str(&format!("- Current branch: {}\n", self.current_branch));
        output.push_str(&format!("- Git status: {}\n", self.git_status.summary));

        if let Some(ref checkpoint) = self.last_checkpoint {
            output.push_str(&format!(
                "- Last checkpoint: {} - {} ({})\n",
                checkpoint.sha,
                checkpoint.message,
                checkpoint.timestamp.format("%Y-%m-%d %H:%M:%S UTC")
            ));
        } else {
            output.push_str("- Last checkpoint: none\n");
        }

        if self.recent_progress.is_empty() {
            output.push_str("- Progress log (recent): none\n");
        } else {
            output.push_str("- Progress log (recent):\n");
            for entry in &self.recent_progress {
                output.push_str(&format!(
                    "  - [{}] {}: {}\n",
                    entry.timestamp.format("%H:%M:%S"),
                    entry.entry_type,
                    entry.message
                ));
            }
        }

        output.push_str(&format!("- Current phase: {}\n", self.current_phase));
        output.push_str("\n---\n\n");

        output
    }
}

/// Build context recovery information for a task worktree.
///
/// # Arguments
/// * `worktree_path` - Path to the task's worktree directory
/// * `phase` - Current task phase (e.g., "execute", "review")
/// * `max_progress_entries` - Maximum number of recent progress entries to include
///
/// # Returns
/// A formatted markdown string containing context recovery information.
pub fn build_context_recovery(
    worktree_path: &Path,
    phase: &str,
    max_progress_entries: usize,
) -> String {
    let context = build_context_recovery_data(worktree_path, phase, max_progress_entries);
    context.to_markdown()
}

/// Build context recovery data structure for a task worktree.
///
/// # Arguments
/// * `worktree_path` - Path to the task's worktree directory
/// * `phase` - Current task phase (e.g., "execute", "review")
/// * `max_progress_entries` - Maximum number of recent progress entries to include
///
/// # Returns
/// A ContextRecovery struct containing all gathered information.
pub fn build_context_recovery_data(
    worktree_path: &Path,
    phase: &str,
    max_progress_entries: usize,
) -> ContextRecovery {
    let current_branch = get_current_branch(worktree_path).unwrap_or_else(|| "unknown".to_string());
    let git_status = get_git_status_summary(worktree_path);
    let last_checkpoint = get_last_checkpoint(worktree_path);
    let recent_progress = get_recent_progress(worktree_path, max_progress_entries);

    ContextRecovery {
        current_branch,
        git_status,
        last_checkpoint,
        recent_progress,
        current_phase: phase.to_string(),
    }
}

/// Get the current git branch name.
fn get_current_branch(worktree_path: &Path) -> Option<String> {
    let output = Command::new("git")
        .args(["rev-parse", "--abbrev-ref", "HEAD"])
        .current_dir(worktree_path)
        .output()
        .ok()?;

    if output.status.success() {
        let branch = String::from_utf8_lossy(&output.stdout).trim().to_string();
        if branch.is_empty() || branch == "HEAD" {
            // Detached HEAD state - try to get the commit SHA
            let sha_output = Command::new("git")
                .args(["rev-parse", "--short", "HEAD"])
                .current_dir(worktree_path)
                .output()
                .ok()?;
            if sha_output.status.success() {
                Some(format!("(detached at {})", String::from_utf8_lossy(&sha_output.stdout).trim()))
            } else {
                Some("(detached)".to_string())
            }
        } else {
            Some(branch)
        }
    } else {
        None
    }
}

/// Get a summary of the git status.
fn get_git_status_summary(worktree_path: &Path) -> GitStatusSummary {
    let output = Command::new("git")
        .args(["status", "--porcelain"])
        .current_dir(worktree_path)
        .output();

    let output = match output {
        Ok(o) if o.status.success() => o,
        _ => return GitStatusSummary {
            summary: "unable to get git status".to_string(),
            ..Default::default()
        },
    };

    let status_output = String::from_utf8_lossy(&output.stdout);
    let mut staged = 0u32;
    let mut modified = 0u32;
    let mut untracked = 0u32;
    let mut deleted = 0u32;

    for line in status_output.lines() {
        if line.len() < 2 {
            continue;
        }
        let index_status = line.chars().next().unwrap_or(' ');
        let worktree_status = line.chars().nth(1).unwrap_or(' ');

        // Count staged changes (index column)
        match index_status {
            'A' | 'M' | 'R' | 'C' => staged += 1,
            'D' => {
                staged += 1;
                deleted += 1;
            }
            _ => {}
        }

        // Count unstaged changes (worktree column)
        match worktree_status {
            'M' => modified += 1,
            'D' => {
                modified += 1;
                deleted += 1;
            }
            '?' => untracked += 1,
            _ => {}
        }
    }

    // Build summary string
    let mut parts = Vec::new();
    if staged > 0 {
        parts.push(format!("{} staged", staged));
    }
    if modified > 0 {
        parts.push(format!("{} modified", modified));
    }
    if untracked > 0 {
        parts.push(format!("{} untracked", untracked));
    }

    let summary = if parts.is_empty() {
        "clean".to_string()
    } else {
        parts.join(", ")
    };

    GitStatusSummary {
        staged,
        modified,
        untracked,
        deleted,
        summary,
    }
}

/// Get information about the last checkpoint (most recent commit).
fn get_last_checkpoint(worktree_path: &Path) -> Option<CheckpointInfo> {
    // Get the most recent commit info
    let output = Command::new("git")
        .args([
            "log",
            "-1",
            "--format=%h|%s|%aI", // short sha | subject | ISO 8601 date
        ])
        .current_dir(worktree_path)
        .output()
        .ok()?;

    if !output.status.success() {
        return None;
    }

    let log_output = String::from_utf8_lossy(&output.stdout).trim().to_string();
    let parts: Vec<&str> = log_output.splitn(3, '|').collect();

    if parts.len() < 3 {
        return None;
    }

    let sha = parts[0].to_string();
    let message = parts[1].to_string();
    let timestamp = DateTime::parse_from_rfc3339(parts[2])
        .ok()?
        .with_timezone(&Utc);

    Some(CheckpointInfo {
        sha,
        message,
        timestamp,
    })
}

/// Get recent progress entries from the progress.log file.
fn get_recent_progress(worktree_path: &Path, max_entries: usize) -> Vec<ProgressEntry> {
    let progress_log_path = worktree_path.join("progress.log");

    let content = match std::fs::read_to_string(&progress_log_path) {
        Ok(c) => c,
        Err(_) => return Vec::new(),
    };

    let mut entries = Vec::new();

    // Parse the progress.log file
    // Expected format:
    // [2024-01-15T10:30:00Z] ENTRY_TYPE
    // key: value
    // key: value
    // (blank line)
    let mut lines = content.lines().peekable();

    while let Some(line) = lines.next() {
        // Look for timestamp lines
        if line.starts_with('[') && line.contains(']') {
            if let Some(entry) = parse_progress_entry(line, &mut lines) {
                entries.push(entry);
            }
        }
    }

    // Return the most recent entries (up to max_entries)
    let start = if entries.len() > max_entries {
        entries.len() - max_entries
    } else {
        0
    };

    entries[start..].to_vec()
}

/// Parse a single progress entry from the log.
fn parse_progress_entry<'a, I>(header_line: &str, lines: &mut std::iter::Peekable<I>) -> Option<ProgressEntry>
where
    I: Iterator<Item = &'a str>,
{
    // Parse header: [2024-01-15T10:30:00Z] ENTRY_TYPE
    let close_bracket = header_line.find(']')?;
    let timestamp_str = &header_line[1..close_bracket];
    let entry_type = header_line[close_bracket + 1..].trim().to_string();

    let timestamp = DateTime::parse_from_rfc3339(timestamp_str)
        .ok()?
        .with_timezone(&Utc);

    // Collect message from following lines until blank line or next entry
    let mut message_parts = Vec::new();

    while let Some(&next_line) = lines.peek() {
        if next_line.is_empty() || next_line.starts_with('#') {
            lines.next(); // consume blank/comment line
            break;
        }
        if next_line.starts_with('[') {
            // Next entry starts, don't consume
            break;
        }

        let line = lines.next().unwrap();
        // Extract key-value or just the line
        if let Some(colon_pos) = line.find(':') {
            let key = line[..colon_pos].trim();
            let value = line[colon_pos + 1..].trim();
            if key == "progress" || key == "message" || key == "reason" {
                message_parts.push(value.to_string());
            }
        }
    }

    let message = if message_parts.is_empty() {
        entry_type.clone()
    } else {
        message_parts.join("; ")
    };

    Some(ProgressEntry {
        timestamp,
        entry_type,
        message,
    })
}

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

    fn setup_git_repo() -> TempDir {
        let temp_dir = TempDir::new().unwrap();
        let path = temp_dir.path();

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

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

        Command::new("git")
            .args(["config", "user.name", "Test User"])
            .current_dir(path)
            .output()
            .unwrap();

        // Create initial commit
        fs::write(path.join("README.md"), "# Test").unwrap();
        Command::new("git")
            .args(["add", "README.md"])
            .current_dir(path)
            .output()
            .unwrap();
        Command::new("git")
            .args(["commit", "-m", "Initial commit"])
            .current_dir(path)
            .output()
            .unwrap();

        temp_dir
    }

    #[test]
    fn test_get_current_branch() {
        let temp_dir = setup_git_repo();
        let branch = get_current_branch(temp_dir.path());
        // Default branch could be main or master depending on git config
        assert!(branch.is_some());
        let branch_name = branch.unwrap();
        assert!(branch_name == "main" || branch_name == "master");
    }

    #[test]
    fn test_git_status_clean() {
        let temp_dir = setup_git_repo();
        let status = get_git_status_summary(temp_dir.path());
        assert_eq!(status.staged, 0);
        assert_eq!(status.modified, 0);
        assert_eq!(status.untracked, 0);
        assert_eq!(status.summary, "clean");
    }

    #[test]
    fn test_git_status_with_changes() {
        let temp_dir = setup_git_repo();
        let path = temp_dir.path();

        // Create untracked file
        fs::write(path.join("new_file.txt"), "new content").unwrap();

        // Modify tracked file
        fs::write(path.join("README.md"), "# Modified").unwrap();

        let status = get_git_status_summary(path);
        assert_eq!(status.modified, 1);
        assert_eq!(status.untracked, 1);
        assert!(status.summary.contains("modified"));
        assert!(status.summary.contains("untracked"));
    }

    #[test]
    fn test_get_last_checkpoint() {
        let temp_dir = setup_git_repo();
        let checkpoint = get_last_checkpoint(temp_dir.path());
        assert!(checkpoint.is_some());

        let checkpoint = checkpoint.unwrap();
        assert_eq!(checkpoint.message, "Initial commit");
        assert!(!checkpoint.sha.is_empty());
    }

    #[test]
    fn test_build_context_recovery() {
        let temp_dir = setup_git_repo();
        let markdown = build_context_recovery(temp_dir.path(), "execute", 5);

        assert!(markdown.contains("## Context Recovery"));
        assert!(markdown.contains("Current branch:"));
        assert!(markdown.contains("Git status:"));
        assert!(markdown.contains("Last checkpoint:"));
        assert!(markdown.contains("Current phase: execute"));
    }

    #[test]
    fn test_progress_log_parsing() {
        let temp_dir = setup_git_repo();
        let path = temp_dir.path();

        // Create a mock progress.log
        let progress_content = r#"# progress.log
# Auto-generated by Makima - DO NOT EDIT MANUALLY

[2024-01-15T10:30:00Z] TASK_START
task_id: test-123
task_name: Test Task

[2024-01-15T10:35:00Z] ITERATION 1
progress: Analyzed codebase structure
files_modified: 2

[2024-01-15T10:40:00Z] ITERATION 2
progress: Implemented first feature
files_modified: 3
"#;
        fs::write(path.join("progress.log"), progress_content).unwrap();

        let entries = get_recent_progress(path, 5);
        assert_eq!(entries.len(), 3);
        assert_eq!(entries[0].entry_type, "TASK_START");
        assert_eq!(entries[1].entry_type, "ITERATION 1");
        assert_eq!(entries[2].entry_type, "ITERATION 2");
    }

    #[test]
    fn test_context_recovery_to_markdown() {
        let context = ContextRecovery {
            current_branch: "feature/test".to_string(),
            git_status: GitStatusSummary {
                staged: 1,
                modified: 2,
                untracked: 0,
                deleted: 0,
                summary: "1 staged, 2 modified".to_string(),
            },
            last_checkpoint: Some(CheckpointInfo {
                sha: "abc1234".to_string(),
                message: "Added feature X".to_string(),
                timestamp: Utc::now(),
            }),
            recent_progress: vec![
                ProgressEntry {
                    timestamp: Utc::now(),
                    entry_type: "ITERATION 1".to_string(),
                    message: "Started implementation".to_string(),
                },
            ],
            current_phase: "execute".to_string(),
        };

        let markdown = context.to_markdown();
        assert!(markdown.contains("## Context Recovery"));
        assert!(markdown.contains("feature/test"));
        assert!(markdown.contains("1 staged, 2 modified"));
        assert!(markdown.contains("abc1234"));
        assert!(markdown.contains("Added feature X"));
        assert!(markdown.contains("ITERATION 1"));
        assert!(markdown.contains("execute"));
    }
}