do-memory-core 0.1.34

//! Checkpoint and handoff operations.
//!
//! Core functions for creating checkpoints, generating handoff packs, and resuming work.

use crate::episode::ExecutionStep;
use crate::error::{Error, Result};
use crate::memory::SelfLearningMemory;
use crate::memory::pattern_search::PatternSearchResult;
use crate::patterns::Heuristic;
use chrono::Utc;
use tracing::{info, instrument, warn};
use uuid::Uuid;

use super::{CheckpointMeta, HandoffPack};
use crate::episode::Episode;
use crate::types::TaskOutcome;

/// Create a checkpoint for an in-progress episode.
///
/// Saves the current state of an episode for handoffs or recovery.
/// Returns an error if the episode doesn't exist, is already completed, or storage fails.
#[instrument(skip(memory), fields(episode_id = %episode_id))]
pub async fn checkpoint_episode(
    memory: &SelfLearningMemory,
    episode_id: Uuid,
    reason: String,
) -> Result<CheckpointMeta> {
    checkpoint_episode_with_note(memory, episode_id, reason, None).await
}

/// Create a checkpoint with an optional note.
///
/// Same as [`checkpoint_episode`] but allows specifying a note.
#[instrument(skip(memory), fields(episode_id = %episode_id))]
pub async fn checkpoint_episode_with_note(
    memory: &SelfLearningMemory,
    episode_id: Uuid,
    reason: String,
    note: Option<String>,
) -> Result<CheckpointMeta> {
    info!("Creating checkpoint for episode: {}", episode_id);

    let mut episode = memory.get_episode(episode_id).await?;

    if episode.is_complete() {
        warn!(
            "Cannot create checkpoint for completed episode: {}",
            episode_id
        );
        return Err(Error::InvalidState(
            "Cannot create checkpoint for completed episode".to_string(),
        ));
    }

    let step_number = episode.steps.len();
    let checkpoint = CheckpointMeta::new(
        episode.episode_id,
        reason,
        step_number,
        note,
        episode.salient_features.clone(),
    );
    episode.checkpoints.push(checkpoint.clone());
    memory.update_episode_full(&episode).await?;

    info!(
        checkpoint_id = %checkpoint.checkpoint_id,
        step_number = checkpoint.step_number,
        "Created checkpoint"
    );

    Ok(checkpoint)
}

/// Call this function wherever an episode is finalized (outcome set + end_time set).
pub fn maybe_create_abstention_checkpoint(episode: &mut Episode) {
    if let Some(TaskOutcome::Abstained {
        ref reason,
        stopped_at_step,
        ..
    }) = episode.outcome
    {
        // Only auto-checkpoint if there are actual steps to preserve
        if stopped_at_step > 0 {
            let checkpoint = CheckpointMeta {
                checkpoint_id: Uuid::new_v4(),
                episode_id: episode.episode_id,
                step_number: stopped_at_step,
                timestamp: Utc::now(),
                label: format!("[ABSTAIN] {reason}"),
                salient_features_snapshot: episode.salient_features.clone(),
                is_abstention_checkpoint: true,
                note: None,
            };
            episode.checkpoints.push(checkpoint);
        }
    }
}

/// Returns all checkpoints from a given episode that were created on abstention.
/// Used to find resumable states for handoff to another agent.
pub fn list_abstention_checkpoints(episode: &Episode) -> Vec<&CheckpointMeta> {
    episode
        .checkpoints
        .iter()
        .filter(|c| c.is_abstention_checkpoint)
        .collect()
}

/// Generate a handoff pack from a checkpoint.
///
/// Creates a comprehensive context package for transferring work between agents,
/// including current goal, lessons learned, relevant patterns, and suggested next steps.
/// Returns an error if the checkpoint or episode doesn't exist.
#[instrument(skip(memory), fields(checkpoint_id = %checkpoint_id))]
pub async fn get_handoff_pack(
    memory: &SelfLearningMemory,
    checkpoint_id: Uuid,
) -> Result<HandoffPack> {
    info!("Generating handoff pack for checkpoint: {}", checkpoint_id);

    let (episode, checkpoint) = find_checkpoint(memory, checkpoint_id).await?;

    let steps_completed: Vec<ExecutionStep> = episode
        .steps
        .iter()
        .take(checkpoint.step_number)
        .cloned()
        .collect();

    let (what_worked, what_failed, salient_facts) =
        extract_lessons(memory, &episode, checkpoint.step_number);

    let suggested_next_steps = generate_suggested_next_steps(memory, &episode).await;
    let relevant_patterns = get_relevant_patterns(memory, &episode).await;
    let relevant_heuristics = get_relevant_heuristics(memory, &episode).await;

    let handoff = HandoffPack {
        checkpoint_id: checkpoint.checkpoint_id,
        episode_id: episode.episode_id,
        timestamp: Utc::now(),
        current_goal: episode.task_description.clone(),
        steps_completed,
        what_worked,
        what_failed,
        salient_facts,
        suggested_next_steps,
        relevant_patterns,
        relevant_heuristics,
    };

    info!(
        step_count = handoff.step_count(),
        pattern_count = handoff.relevant_patterns.len(),
        heuristic_count = handoff.relevant_heuristics.len(),
        "Generated handoff pack"
    );

    Ok(handoff)
}

/// Resume work from a handoff pack.
///
/// Creates a new episode initialized with context from a handoff pack,
/// including the same goal, lessons learned, and suggested next steps.
#[instrument(skip(memory, handoff), fields(checkpoint_id = %handoff.checkpoint_id))]
pub async fn resume_from_handoff(
    memory: &SelfLearningMemory,
    handoff: HandoffPack,
) -> Result<Uuid> {
    info!(
        "Resuming from handoff pack: checkpoint_id={}, steps={}",
        handoff.checkpoint_id,
        handoff.step_count()
    );

    let context = crate::types::TaskContext {
        domain: "resumed".to_string(),
        language: None,
        framework: None,
        complexity: crate::types::ComplexityLevel::Moderate,
        tags: vec![
            "resumed".to_string(),
            format!("from-{}", handoff.episode_id),
        ],
    };

    let new_episode_id = memory
        .start_episode(
            handoff.current_goal.clone(),
            context,
            crate::types::TaskType::Other,
        )
        .await;

    // Store handoff context in episode metadata using normal update path
    let mut episode = memory.get_episode(new_episode_id).await?;
    episode.metadata.insert(
        "resumed_from_checkpoint".to_string(),
        handoff.checkpoint_id.to_string(),
    );
    episode.metadata.insert(
        "resumed_from_episode".to_string(),
        handoff.episode_id.to_string(),
    );
    episode.metadata.insert(
        "what_worked".to_string(),
        serde_json::to_string(&handoff.what_worked).unwrap_or_default(),
    );
    episode.metadata.insert(
        "what_failed".to_string(),
        serde_json::to_string(&handoff.what_failed).unwrap_or_default(),
    );
    episode.metadata.insert(
        "salient_facts".to_string(),
        serde_json::to_string(&handoff.salient_facts).unwrap_or_default(),
    );
    episode.metadata.insert(
        "suggested_next_steps".to_string(),
        serde_json::to_string(&handoff.suggested_next_steps).unwrap_or_default(),
    );
    memory.update_episode_full(&episode).await?;

    info!(new_episode_id = %new_episode_id, "Created new episode for resumption");

    Ok(new_episode_id)
}

/// Find an episode and checkpoint by checkpoint ID.
async fn find_checkpoint(
    memory: &SelfLearningMemory,
    checkpoint_id: Uuid,
) -> Result<(crate::episode::Episode, CheckpointMeta)> {
    let episodes = memory.get_all_episodes().await?;

    for episode in episodes {
        if let Some(checkpoint) = episode
            .checkpoints
            .iter()
            .find(|c| c.checkpoint_id == checkpoint_id)
        {
            let checkpoint = checkpoint.clone();
            return Ok((episode, checkpoint));
        }
    }

    Err(Error::NotFound(checkpoint_id))
}

/// Extract lessons learned from an episode up to a step.
fn extract_lessons(
    memory: &SelfLearningMemory,
    episode: &crate::episode::Episode,
    step_number: usize,
) -> (Vec<String>, Vec<String>, Vec<String>) {
    let mut what_worked = Vec::new();
    let mut what_failed = Vec::new();
    let mut salient_facts = Vec::new();

    let mut partial_episode = episode.clone();
    partial_episode.steps.truncate(step_number);

    let features = memory.salient_extractor.extract(&partial_episode);

    for step in &partial_episode.steps {
        if step.is_success() {
            what_worked.push(format!("{}: {}", step.tool, step.action));
        } else {
            what_failed.push(format!("{}: {}", step.tool, step.action));
        }
    }

    for decision in &features.critical_decisions {
        salient_facts.push(format!("Decision: {decision}"));
    }
    for insight in &features.key_insights {
        salient_facts.push(format!("Insight: {insight}"));
    }
    for recovery in &features.error_recovery_patterns {
        salient_facts.push(format!("Recovery: {recovery}"));
    }

    (what_worked, what_failed, salient_facts)
}

/// Generate suggested next steps based on episode context.
async fn generate_suggested_next_steps(
    memory: &SelfLearningMemory,
    episode: &crate::episode::Episode,
) -> Vec<String> {
    let mut suggestions = Vec::new();

    let patterns = memory.retrieve_relevant_patterns(&episode.context, 3).await;

    for pattern in patterns {
        match &pattern {
            crate::patterns::Pattern::ToolSequence { tools, .. } => {
                let remaining_tools: Vec<_> = tools
                    .iter()
                    .filter(|t| !episode.steps.iter().any(|s| &s.tool == *t))
                    .collect();
                if !remaining_tools.is_empty() {
                    suggestions.push(format!(
                        "Consider using tools: {}",
                        remaining_tools
                            .iter()
                            .map(|t| t.as_str())
                            .collect::<Vec<_>>()
                            .join(", ")
                    ));
                }
            }
            crate::patterns::Pattern::DecisionPoint {
                condition, action, ..
            } => {
                suggestions.push(format!("Consider: {condition} (then {action})"));
            }
            _ => {}
        }
    }

    if episode.steps.len() > 5 && episode.successful_steps_count() < episode.steps.len() / 2 {
        suggestions.push("Consider breaking down the task into smaller steps".to_string());
        suggestions.push("Review the what_failed list to avoid repeating mistakes".to_string());
    }

    suggestions.truncate(5);
    suggestions
}

/// Get relevant patterns for an episode.
async fn get_relevant_patterns(
    memory: &SelfLearningMemory,
    episode: &crate::episode::Episode,
) -> Vec<PatternSearchResult> {
    let config = crate::memory::pattern_search::SearchConfig::default();
    memory
        .search_patterns(&episode.task_description, &episode.context, config)
        .await
        .unwrap_or_default()
        .into_iter()
        .take(5)
        .collect()
}

/// Get relevant heuristics for an episode.
async fn get_relevant_heuristics(
    memory: &SelfLearningMemory,
    episode: &crate::episode::Episode,
) -> Vec<Heuristic> {
    let all_heuristics = memory.get_all_heuristics().await.unwrap_or_default();

    all_heuristics
        .into_iter()
        .filter(|h| {
            let condition_lower = h.condition.to_lowercase();
            let task_lower = episode.task_description.to_lowercase();
            condition_lower
                .split_whitespace()
                .any(|word| task_lower.contains(word))
        })
        .take(5)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{TaskContext, TaskType};

    #[tokio::test]
    async fn test_checkpoint_episode_not_found() {
        let memory = SelfLearningMemory::new();
        let result = checkpoint_episode(&memory, Uuid::new_v4(), "test".to_string()).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_get_handoff_pack_not_found() {
        let memory = SelfLearningMemory::new();
        let result = get_handoff_pack(&memory, Uuid::new_v4()).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_checkpoint_completed_episode() {
        use crate::episode::ExecutionStep;
        use crate::memory::MemoryConfig;
        use crate::types::ExecutionResult;

        let test_config = MemoryConfig {
            quality_threshold: 0.3,
            ..Default::default()
        };
        let memory = SelfLearningMemory::with_config(test_config);

        let episode_id = memory
            .start_episode(
                "Test task".to_string(),
                TaskContext::default(),
                TaskType::Testing,
            )
            .await;

        let mut step = ExecutionStep::new(1, "test_tool".to_string(), "test action".to_string());
        step.result = Some(ExecutionResult::Success {
            output: "test output".to_string(),
        });
        memory.log_step(episode_id, step).await;

        memory
            .complete_episode(
                episode_id,
                crate::types::TaskOutcome::Success {
                    verdict: "Done".to_string(),
                    artifacts: vec![],
                },
            )
            .await
            .unwrap();

        let result = checkpoint_episode(&memory, episode_id, "test".to_string()).await;
        assert!(result.is_err());
    }
}

#[cfg(test)]
mod abstention_tests {
    use super::*;
    use crate::episode::Episode;
    use crate::types::{TaskContext, TaskOutcome, TaskType};

    fn create_test_episode_with_steps(step_count: usize) -> Episode {
        let mut episode = Episode::new(
            "Test task".to_string(),
            TaskContext::default(),
            TaskType::Testing,
        );
        for i in 1..=step_count {
            episode.add_step(crate::episode::ExecutionStep::new(
                i,
                "tool".to_string(),
                "action".to_string(),
            ));
        }
        episode
    }

    #[test]
    fn test_abstention_auto_checkpoint_created() {
        let mut episode = create_test_episode_with_steps(5);
        episode.outcome = Some(TaskOutcome::Abstained {
            reason: "Tool not found".to_string(),
            stopped_at_step: 3,
            infeasibility_signals: vec!["404_tool".to_string()],
        });
        maybe_create_abstention_checkpoint(&mut episode);
        assert_eq!(episode.checkpoints.len(), 1);
        assert!(episode.checkpoints[0].is_abstention_checkpoint);
        assert_eq!(episode.checkpoints[0].step_number, 3);
        assert!(episode.checkpoints[0].label.starts_with("[ABSTAIN]"));
    }

    #[test]
    fn test_no_checkpoint_for_zero_step_abstention() {
        let mut episode = create_test_episode_with_steps(0);
        episode.outcome = Some(TaskOutcome::Abstained {
            reason: "Immediate rejection".to_string(),
            stopped_at_step: 0,
            infeasibility_signals: vec![],
        });
        maybe_create_abstention_checkpoint(&mut episode);
        // No steps = no useful state to checkpoint
        assert_eq!(episode.checkpoints.len(), 0);
    }

    #[test]
    fn test_non_abstention_outcome_no_auto_checkpoint() {
        let mut episode = create_test_episode_with_steps(5);
        episode.outcome = Some(TaskOutcome::Failure {
            reason: "Error".to_string(),
            error_details: None,
        });
        maybe_create_abstention_checkpoint(&mut episode);
        assert_eq!(episode.checkpoints.len(), 0);
    }

    #[test]
    fn test_abstention_checkpoint_backward_compat_deserialization() {
        // Old JSON without is_abstention_checkpoint and episode_id fields should deserialize with defaults
        // and aliases should handle timestamp/label
        let json = r#"{
            "checkpoint_id": "00000000-0000-0000-0000-000000000000",
            "step_number": 3,
            "created_at": "2026-01-01T00:00:00Z",
            "reason": "old checkpoint",
            "salient_features_snapshot": null,
            "note": "old note"
        }"#;
        let meta: CheckpointMeta = serde_json::from_str(json).unwrap();
        assert!(!meta.is_abstention_checkpoint);
        assert_eq!(meta.label, "old checkpoint");
        assert_eq!(meta.step_number, 3);
    }
}