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

//! Learning injection for enhanced task prompts.
//!
//! This module implements Ralph-inspired learning injection into Claude prompts.
//! It prepends context recovery information and recent learnings to base prompts,
//! enabling cross-task learning and context preservation.
//!
//! The injection includes:
//! - Context recovery info (branch, git status, last checkpoint, recent progress)
//! - Previous learnings from progress log (patterns, gotchas, tips)
//!
//! Output format:
//! ```markdown
//! ## Context Recovery
//! [from context_recovery.rs]
//!
//! ## Previous Learnings
//! These patterns were discovered in previous iterations:
//! - Pattern: Always run typecheck before commit
//! - Gotcha: Remember to update imports when moving files
//! - Tip: Use the existing helper in utils.rs
//!
//! ---
//! [Original task prompt follows]
//! ```

use std::path::Path;

use regex::Regex;

use super::context_recovery::{self, ContextRecoveryInfo};
use super::progress_log::{self, Learning};
use super::ManagedTask;

/// Default maximum number of learning entries to inject into prompts.
const DEFAULT_MAX_ENTRIES_INJECTED: usize = 20;

/// Error type for learning injection operations.
#[derive(Debug, thiserror::Error)]
pub enum LearningInjectionError {
    #[error("Context recovery error: {0}")]
    ContextRecovery(#[from] context_recovery::ContextRecoveryError),

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

/// Configuration for learning injection.
#[derive(Debug, Clone)]
pub struct LearningInjector {
    /// Maximum number of learning entries to inject into prompts.
    pub max_entries_injected: usize,
    /// Whether to include context recovery information in the prompt.
    pub include_context_recovery: bool,
}

impl Default for LearningInjector {
    fn default() -> Self {
        Self {
            max_entries_injected: DEFAULT_MAX_ENTRIES_INJECTED,
            include_context_recovery: true,
        }
    }
}

impl LearningInjector {
    /// Create a new LearningInjector with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the maximum number of entries to inject.
    pub fn with_max_entries(mut self, max_entries: usize) -> Self {
        self.max_entries_injected = max_entries;
        self
    }

    /// Set whether to include context recovery.
    pub fn with_context_recovery(mut self, include: bool) -> Self {
        self.include_context_recovery = include;
        self
    }

    /// Inject learnings and context recovery into a base prompt.
    ///
    /// This function:
    /// 1. Prepends context recovery info (if enabled)
    /// 2. Prepends recent learnings from progress log
    /// 3. Returns the enhanced prompt
    ///
    /// # Arguments
    /// * `worktree_path` - Path to the task's worktree directory
    /// * `task` - The managed task struct
    /// * `base_prompt` - The original task prompt to enhance
    ///
    /// # Returns
    /// Enhanced prompt with context and learnings prepended.
    pub async fn inject_learnings(
        &self,
        worktree_path: &Path,
        task: &ManagedTask,
        base_prompt: &str,
    ) -> Result<String, LearningInjectionError> {
        let mut enhanced_prompt = String::new();

        // Add context recovery if enabled
        if self.include_context_recovery {
            let context_info = context_recovery::build_context_recovery(worktree_path, task).await?;
            enhanced_prompt.push_str(&context_info.to_markdown());
            enhanced_prompt.push('\n');
        }

        // Gather learnings from recent progress log entries
        let learnings = self.gather_learnings(worktree_path)?;

        // Add learnings section if there are any
        if !learnings.is_empty() {
            enhanced_prompt.push_str(&self.format_learnings_section(&learnings));
            enhanced_prompt.push('\n');
        }

        // Add separator before original prompt (if we added any content)
        if !enhanced_prompt.is_empty() {
            enhanced_prompt.push_str("---\n\n");
        }

        // Append the original prompt
        enhanced_prompt.push_str(base_prompt);

        Ok(enhanced_prompt)
    }

    /// Gather learnings from recent progress log entries.
    ///
    /// Reads the progress log and extracts unique learnings up to the
    /// configured maximum number of entries.
    fn gather_learnings(&self, worktree_path: &Path) -> Result<Vec<Learning>, LearningInjectionError> {
        // Read enough entries to potentially have max_entries_injected learnings
        // Since each entry may have multiple learnings, we read more entries
        let entries = progress_log::read_recent_entries(worktree_path, self.max_entries_injected * 2)?;

        let mut learnings = Vec::new();
        let mut seen_descriptions = std::collections::HashSet::new();

        // Iterate in reverse to get most recent learnings first
        for entry in entries.iter().rev() {
            for learning in &entry.learnings {
                // Deduplicate by description
                if !seen_descriptions.contains(&learning.description) {
                    seen_descriptions.insert(learning.description.clone());
                    learnings.push(learning.clone());

                    if learnings.len() >= self.max_entries_injected {
                        break;
                    }
                }
            }
            if learnings.len() >= self.max_entries_injected {
                break;
            }
        }

        Ok(learnings)
    }

    /// Format the learnings section for prompt injection.
    fn format_learnings_section(&self, learnings: &[Learning]) -> String {
        let mut section = String::new();

        section.push_str("## Previous Learnings\n");
        section.push_str("These patterns were discovered in previous iterations:\n");

        for learning in learnings {
            section.push_str(&format!("- {}: {}\n", learning.kind, learning.description));
        }

        section
    }
}

/// Extract learnings from Claude output text.
///
/// Parses Claude output looking for learning patterns marked with:
/// - LEARNING: Description of what was learned
/// - PATTERN: A reusable pattern discovered
/// - GOTCHA: Something to watch out for
/// - TIP: A helpful suggestion
///
/// # Arguments
/// * `output` - The Claude output text to parse
///
/// # Returns
/// Vector of learnings extracted from the output.
pub fn extract_learnings_from_output(output: &str) -> Vec<Learning> {
    let mut learnings = Vec::new();

    // Patterns to look for in Claude output
    // Format: MARKER: description text (until end of line or next marker)
    let patterns = [
        ("LEARNING", "Learning"),
        ("PATTERN", "Pattern"),
        ("GOTCHA", "Gotcha"),
        ("TIP", "Tip"),
    ];

    for line in output.lines() {
        let trimmed = line.trim();

        for (marker, kind) in &patterns {
            // Check for "MARKER:" at start of line (case-insensitive)
            let pattern_prefix = format!("{}:", marker);

            if trimmed.to_uppercase().starts_with(&pattern_prefix) {
                // Extract the description after the marker
                let description = trimmed[pattern_prefix.len()..].trim();

                if !description.is_empty() {
                    learnings.push(Learning {
                        kind: kind.to_string(),
                        description: description.to_string(),
                    });
                }
            }

            // Also check for "**MARKER:**" markdown bold format
            let bold_prefix = format!("**{}:**", marker);
            let bold_prefix_upper = format!("**{}:**", marker.to_uppercase());

            if trimmed.starts_with(&bold_prefix) || trimmed.starts_with(&bold_prefix_upper) {
                let prefix_len = if trimmed.starts_with(&bold_prefix) {
                    bold_prefix.len()
                } else {
                    bold_prefix_upper.len()
                };
                let description = trimmed[prefix_len..].trim().trim_end_matches("**").trim();

                if !description.is_empty() {
                    learnings.push(Learning {
                        kind: kind.to_string(),
                        description: description.to_string(),
                    });
                }
            }
        }
    }

    // Also look for learnings in code blocks or structured output
    // Example: `LEARNING: description` or within markdown bullets
    let bullet_regex = Regex::new(r"^[-*]\s*(LEARNING|PATTERN|GOTCHA|TIP):\s*(.+)$").unwrap();

    for line in output.lines() {
        let trimmed = line.trim();

        if let Some(captures) = bullet_regex.captures(trimmed) {
            let marker = captures.get(1).map(|m| m.as_str()).unwrap_or("");
            let description = captures.get(2).map(|m| m.as_str()).unwrap_or("").trim();

            if !description.is_empty() {
                let kind = match marker.to_uppercase().as_str() {
                    "LEARNING" => "Learning",
                    "PATTERN" => "Pattern",
                    "GOTCHA" => "Gotcha",
                    "TIP" => "Tip",
                    _ => continue,
                };

                // Check for duplicates before adding
                let already_exists = learnings
                    .iter()
                    .any(|l| l.kind == kind && l.description == description);

                if !already_exists {
                    learnings.push(Learning {
                        kind: kind.to_string(),
                        description: description.to_string(),
                    });
                }
            }
        }
    }

    learnings
}

/// Convenience function to inject learnings with default settings.
///
/// This is a simpler interface for common use cases.
pub async fn inject_learnings(
    worktree_path: &Path,
    task: &ManagedTask,
    base_prompt: &str,
) -> Result<String, LearningInjectionError> {
    LearningInjector::default()
        .inject_learnings(worktree_path, task, base_prompt)
        .await
}

/// Convenience function to inject learnings with context recovery info provided.
///
/// Use this when you already have context recovery info computed to avoid
/// recomputing it.
pub fn inject_learnings_with_context(
    worktree_path: &Path,
    context_info: Option<&ContextRecoveryInfo>,
    base_prompt: &str,
    max_entries: usize,
) -> Result<String, LearningInjectionError> {
    let mut enhanced_prompt = String::new();

    // Add context recovery if provided
    if let Some(info) = context_info {
        enhanced_prompt.push_str(&info.to_markdown());
        enhanced_prompt.push('\n');
    }

    // Gather learnings from recent progress log entries
    let entries = progress_log::read_recent_entries(worktree_path, max_entries * 2)?;

    let mut learnings = Vec::new();
    let mut seen_descriptions = std::collections::HashSet::new();

    for entry in entries.iter().rev() {
        for learning in &entry.learnings {
            if !seen_descriptions.contains(&learning.description) {
                seen_descriptions.insert(learning.description.clone());
                learnings.push(learning.clone());

                if learnings.len() >= max_entries {
                    break;
                }
            }
        }
        if learnings.len() >= max_entries {
            break;
        }
    }

    // Add learnings section if there are any
    if !learnings.is_empty() {
        enhanced_prompt.push_str("## Previous Learnings\n");
        enhanced_prompt.push_str("These patterns were discovered in previous iterations:\n");

        for learning in &learnings {
            enhanced_prompt.push_str(&format!("- {}: {}\n", learning.kind, learning.description));
        }

        enhanced_prompt.push('\n');
    }

    // Add separator before original prompt (if we added any content)
    if !enhanced_prompt.is_empty() {
        enhanced_prompt.push_str("---\n\n");
    }

    // Append the original prompt
    enhanced_prompt.push_str(base_prompt);

    Ok(enhanced_prompt)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_learning_injector_default() {
        let injector = LearningInjector::default();
        assert_eq!(injector.max_entries_injected, DEFAULT_MAX_ENTRIES_INJECTED);
        assert!(injector.include_context_recovery);
    }

    #[test]
    fn test_learning_injector_builder() {
        let injector = LearningInjector::new()
            .with_max_entries(10)
            .with_context_recovery(false);

        assert_eq!(injector.max_entries_injected, 10);
        assert!(!injector.include_context_recovery);
    }

    #[test]
    fn test_extract_learnings_basic() {
        let output = r#"
I discovered some useful patterns during this work.

LEARNING: Always check if the file exists before reading
PATTERN: Use async/await for IO operations
GOTCHA: Don't forget to handle error cases
TIP: Use the existing validation helper in utils.rs
"#;

        let learnings = extract_learnings_from_output(output);

        assert_eq!(learnings.len(), 4);
        assert_eq!(learnings[0].kind, "Learning");
        assert_eq!(
            learnings[0].description,
            "Always check if the file exists before reading"
        );
        assert_eq!(learnings[1].kind, "Pattern");
        assert_eq!(learnings[1].description, "Use async/await for IO operations");
        assert_eq!(learnings[2].kind, "Gotcha");
        assert_eq!(
            learnings[2].description,
            "Don't forget to handle error cases"
        );
        assert_eq!(learnings[3].kind, "Tip");
        assert_eq!(
            learnings[3].description,
            "Use the existing validation helper in utils.rs"
        );
    }

    #[test]
    fn test_extract_learnings_case_insensitive() {
        let output = r#"
learning: This should still be captured
Pattern: Mixed case works
GOTCHA: All caps works too
tip: Lower case as well
"#;

        let learnings = extract_learnings_from_output(output);

        assert_eq!(learnings.len(), 4);
    }

    #[test]
    fn test_extract_learnings_markdown_bold() {
        let output = r#"
**LEARNING:** This is a bold learning
**PATTERN:** Bold pattern text
"#;

        let learnings = extract_learnings_from_output(output);

        assert_eq!(learnings.len(), 2);
        assert_eq!(learnings[0].kind, "Learning");
        assert_eq!(learnings[0].description, "This is a bold learning");
    }

    #[test]
    fn test_extract_learnings_bullet_points() {
        let output = r#"
Here are some learnings:
- PATTERN: First pattern from bullets
- TIP: A helpful tip
* GOTCHA: Watch out for this
"#;

        let learnings = extract_learnings_from_output(output);

        assert_eq!(learnings.len(), 3);
    }

    #[test]
    fn test_extract_learnings_empty_description() {
        let output = r#"
LEARNING:
PATTERN:
TIP: This one is valid
"#;

        let learnings = extract_learnings_from_output(output);

        // Only the valid one should be captured
        assert_eq!(learnings.len(), 1);
        assert_eq!(learnings[0].kind, "Tip");
    }

    #[test]
    fn test_extract_learnings_no_duplicates() {
        let output = r#"
PATTERN: Same pattern twice
- PATTERN: Same pattern twice
"#;

        let learnings = extract_learnings_from_output(output);

        // Should deduplicate
        assert_eq!(learnings.len(), 1);
    }

    #[test]
    fn test_format_learnings_section() {
        let injector = LearningInjector::default();
        let learnings = vec![
            Learning::pattern("Use XYZ pattern"),
            Learning::gotcha("Watch out for edge cases"),
            Learning::tip("Check the documentation"),
        ];

        let section = injector.format_learnings_section(&learnings);

        assert!(section.contains("## Previous Learnings"));
        assert!(section.contains("These patterns were discovered in previous iterations:"));
        assert!(section.contains("- Pattern: Use XYZ pattern"));
        assert!(section.contains("- Gotcha: Watch out for edge cases"));
        assert!(section.contains("- Tip: Check the documentation"));
    }

    #[test]
    fn test_inject_learnings_with_context_no_entries() {
        let temp_dir = std::env::temp_dir().join(format!(
            "learning_injection_test_{}",
            std::process::id()
        ));
        std::fs::create_dir_all(&temp_dir).unwrap();

        let result = inject_learnings_with_context(
            &temp_dir,
            None,
            "Original prompt content",
            10,
        )
        .unwrap();

        // With no context and no learnings, should just return the base prompt
        assert_eq!(result, "Original prompt content");

        std::fs::remove_dir_all(&temp_dir).ok();
    }

    #[test]
    fn test_inject_learnings_with_context_info() {
        let temp_dir = std::env::temp_dir().join(format!(
            "learning_injection_context_test_{}",
            std::process::id()
        ));
        std::fs::create_dir_all(&temp_dir).unwrap();

        let context_info = ContextRecoveryInfo {
            current_branch: "feature/test".to_string(),
            uncommitted_changes_count: 2,
            last_checkpoint: None,
            current_phase: Some("execute".to_string()),
            recent_progress: vec![],
        };

        let result = inject_learnings_with_context(
            &temp_dir,
            Some(&context_info),
            "Original prompt content",
            10,
        )
        .unwrap();

        assert!(result.contains("## Context Recovery"));
        assert!(result.contains("feature/test"));
        assert!(result.contains("---"));
        assert!(result.contains("Original prompt content"));

        std::fs::remove_dir_all(&temp_dir).ok();
    }

    #[test]
    fn test_gather_learnings_deduplication() {
        let temp_dir = std::env::temp_dir().join(format!(
            "learning_dedup_test_{}",
            std::process::id()
        ));
        std::fs::create_dir_all(&temp_dir).unwrap();

        // Create progress log entries with duplicate learnings
        use super::super::progress_log::{append_entry, ProgressLogEntry, TaskCompletionStatus};

        let entry1 = ProgressLogEntry::new("task-1", "First task", TaskCompletionStatus::Done)
            .with_learnings(vec![
                Learning::pattern("Common pattern"),
                Learning::tip("First tip"),
            ]);
        append_entry(&temp_dir, &entry1).unwrap();

        let entry2 = ProgressLogEntry::new("task-2", "Second task", TaskCompletionStatus::Done)
            .with_learnings(vec![
                Learning::pattern("Common pattern"), // Duplicate
                Learning::gotcha("New gotcha"),
            ]);
        append_entry(&temp_dir, &entry2).unwrap();

        let injector = LearningInjector::new().with_max_entries(10);
        let learnings = injector.gather_learnings(&temp_dir).unwrap();

        // Should have 3 unique learnings (not 4 with duplicate)
        assert_eq!(learnings.len(), 3);

        // Check that "Common pattern" appears only once
        let pattern_count = learnings
            .iter()
            .filter(|l| l.description == "Common pattern")
            .count();
        assert_eq!(pattern_count, 1);

        std::fs::remove_dir_all(&temp_dir).ok();
    }

    #[test]
    fn test_gather_learnings_max_entries() {
        let temp_dir = std::env::temp_dir().join(format!(
            "learning_max_test_{}",
            std::process::id()
        ));
        std::fs::create_dir_all(&temp_dir).unwrap();

        use super::super::progress_log::{append_entry, ProgressLogEntry, TaskCompletionStatus};

        // Create an entry with many learnings
        let learnings: Vec<Learning> = (0..10)
            .map(|i| Learning::pattern(format!("Pattern {}", i)))
            .collect();

        let entry = ProgressLogEntry::new("task-many", "Many learnings", TaskCompletionStatus::Done)
            .with_learnings(learnings);
        append_entry(&temp_dir, &entry).unwrap();

        // Set max to 5
        let injector = LearningInjector::new().with_max_entries(5);
        let gathered = injector.gather_learnings(&temp_dir).unwrap();

        assert_eq!(gathered.len(), 5);

        std::fs::remove_dir_all(&temp_dir).ok();
    }
}