forge_agent/

planner.rs

//! Planning engine - Execution plan generation from observations.
//!
//! This module implements the planning phase of the agent loop, generating
//! execution steps from observations and constraints.

use crate::{AgentError, Result};
use forge_core::Forge;
use std::sync::Arc;

/// Planner for generating execution plans.
///
/// The Planner creates ordered steps from observations, ensuring
/// dependencies are satisfied and conflicts are detected.
#[derive(Clone)]
pub struct Planner {
    /// Forge SDK for graph queries
    forge: Arc<Forge>,
}

impl Planner {
    /// Creates a new planner.
    pub fn new(forge: Forge) -> Self {
        Self {
            forge: Arc::new(forge),
        }
    }

    /// Generates execution steps from an observation.
    ///
    /// # Arguments
    ///
    /// * `observation` - The observation data
    pub async fn generate_steps(
        &self,
        observation: &super::observe::Observation,
    ) -> Result<Vec<PlanStep>> {
        let mut steps = Vec::new();

        // For each symbol in the observation, create appropriate steps
        for symbol in &observation.symbols {
            // In production, this would use LLM to decide what operations
            // For now, create placeholder steps
            steps.push(PlanStep {
                description: format!("Process symbol {}", symbol.name),
                operation: PlanOperation::Inspect {
                    symbol_id: symbol.id,
                    symbol_name: symbol.name.clone(),
                },
            });
        }

        Ok(steps)
    }

    /// Estimates impact of a plan.
    ///
    /// # Arguments
    ///
    /// * `steps` - The planned steps
    pub async fn estimate_impact(&self, steps: &[PlanStep]) -> Result<ImpactEstimate> {
        let mut affected_files = std::collections::HashSet::new();

        // Collect all affected files
        for step in steps {
            match &step.operation {
                PlanOperation::Rename { old, .. } => {
                    // Extract file from symbol name (simplified)
                    if let Some(file) = self.extract_file_from_symbol(old) {
                        affected_files.insert(file);
                    }
                }
                PlanOperation::Delete { name } => {
                    if let Some(file) = self.extract_file_from_symbol(name) {
                        affected_files.insert(file);
                    }
                }
                PlanOperation::Create { path, .. } => {
                    affected_files.insert(path.clone());
                }
                PlanOperation::Inspect { .. } => {
                    // Inspect doesn't modify files
                }
                PlanOperation::Modify { file, .. } => {
                    affected_files.insert(file.clone());
                }
            }
        }

        Ok(ImpactEstimate {
            affected_files: affected_files.into_iter().collect(),
            complexity: steps.len(),
        })
    }

    /// Detects conflicts between steps.
    ///
    /// # Arguments
    ///
    /// * `steps` - The planned steps
    pub fn detect_conflicts(&self, steps: &[PlanStep]) -> Result<Vec<Conflict>> {
        let mut conflicts = Vec::new();
        let mut file_regions: std::collections::HashMap<String, Vec<(usize, usize, usize)>> =
            std::collections::HashMap::new();

        // Track regions in each file
        for (idx, step) in steps.iter().enumerate() {
            if let Some(region) = self.get_step_region(step) {
                file_regions
                    .entry(region.file.clone())
                    .or_insert_with(Vec::new)
                    .push((idx, region.start, region.end));
            }
        }

        // Check for overlaps
        for (file, regions) in &file_regions {
            for i in 0..regions.len() {
                for j in (i + 1)..regions.len() {
                    let (idx1, start1, end1) = regions[i];
                    let (idx2, start2, _end2) = regions[j];

                    // Check for overlap (no dereference needed, values are already usize)
                    if start1 < end1 && start2 < end1 {
                        conflicts.push(Conflict {
                            step_indices: vec![idx1, idx2],
                            file: file.clone(),
                            reason: ConflictReason::OverlappingRegion {
                                start: start1,
                                end: end1,
                            },
                        });
                    }
                }
            }
        }

        Ok(conflicts)
    }

    /// Orders steps based on dependencies.
    ///
    /// # Arguments
    ///
    /// * `steps` - The planned steps
    pub fn order_steps(&self, steps: &mut Vec<PlanStep>) -> Result<()> {
        // Simple topological sort based on step dependencies
        // For now, keep existing order (production would use DAG)
        // In a full implementation, this would:
        // 1. Build dependency graph
        // 2. Topologically sort
        // 3. Detect cycles

        // Ensure Rename comes before Delete for same symbol
        let mut rename_indices = Vec::new();
        let mut delete_indices = Vec::new();

        for (idx, step) in steps.iter().enumerate() {
            match &step.operation {
                PlanOperation::Rename { old, .. } => {
                    rename_indices.push((idx, old.clone()));
                }
                PlanOperation::Delete { name } => {
                    delete_indices.push((idx, name.clone()));
                }
                _ => {}
            }
        }

        // Move renames before deletes for same symbols
        for (rename_idx, name) in &rename_indices {
            for (delete_idx, delete_name) in &delete_indices {
                if name == delete_name && rename_idx > delete_idx {
                    // Swap the steps
                    steps.swap(*rename_idx, *delete_idx);
                }
            }
        }

        Ok(())
    }

    /// Generates rollback plan.
    ///
    /// # Arguments
    ///
    /// * `steps` - The planned steps
    pub fn generate_rollback(&self, steps: &[PlanStep]) -> Vec<RollbackStep> {
        steps
            .iter()
            .rev()
            .map(|step| RollbackStep {
                description: format!("Rollback: {}", step.description),
                operation: match &step.operation {
                    PlanOperation::Rename { old, .. } => RollbackOperation::Rename {
                        new_name: old.clone(),
                    },
                    PlanOperation::Delete { name } => {
                        RollbackOperation::Restore { name: name.clone() }
                    }
                    PlanOperation::Create { path, .. } => {
                        RollbackOperation::Delete { path: path.clone() }
                    }
                    PlanOperation::Inspect { .. } => RollbackOperation::None,
                    PlanOperation::Modify { file, .. } => {
                        RollbackOperation::Restore { name: file.clone() }
                    }
                },
            })
            .collect()
    }

    /// Extracts file path from symbol name (simplified).
    fn extract_file_from_symbol(&self, _symbol: &str) -> Option<String> {
        // In production, this would query the graph for symbol location
        // For now, return None
        None
    }

    /// Gets the file region affected by a step.
    fn get_step_region(&self, step: &PlanStep) -> Option<FileRegion> {
        match &step.operation {
            PlanOperation::Rename { .. } | PlanOperation::Delete { .. } => None,
            PlanOperation::Create { path, .. } => Some(FileRegion {
                file: path.clone(),
                start: 0,
                end: usize::MAX,
            }),
            PlanOperation::Inspect { .. } => None,
            PlanOperation::Modify { file, start, end } => Some(FileRegion {
                file: file.clone(),
                start: *start,
                end: *end,
            }),
        }
    }
}

/// A step in the execution plan.
#[derive(Clone, Debug)]
pub struct PlanStep {
    /// Step description
    pub description: String,
    /// Operation to perform
    pub operation: PlanOperation,
}

/// Operation to perform in a plan step.
#[derive(Clone, Debug)]
pub enum PlanOperation {
    /// Rename a symbol
    Rename { old: String, new: String },
    /// Delete a symbol
    Delete { name: String },
    /// Create new code
    Create { path: String, content: String },
    /// Inspect a symbol (read-only)
    Inspect {
        symbol_id: forge_core::types::SymbolId,
        symbol_name: String,
    },
    /// Modify existing code
    Modify {
        file: String,
        start: usize,
        end: usize,
    },
}

/// Estimated impact of a plan.
#[derive(Clone, Debug)]
pub struct ImpactEstimate {
    /// Files to be modified
    pub affected_files: Vec<String>,
    /// Estimated complexity
    pub complexity: usize,
}

/// A conflict detected between steps.
#[derive(Clone, Debug)]
pub struct Conflict {
    /// Indices of conflicting steps
    pub step_indices: Vec<usize>,
    /// File where conflict occurs
    pub file: String,
    /// Reason for conflict
    pub reason: ConflictReason,
}

/// Reason for a conflict.
#[derive(Clone, Debug)]
pub enum ConflictReason {
    /// Overlapping regions in the same file
    OverlappingRegion { start: usize, end: usize },
    /// Circular dependency
    CircularDependency,
    /// Missing dependency
    MissingDependency,
}

/// A rollback step.
#[derive(Clone, Debug)]
pub struct RollbackStep {
    /// Step description
    pub description: String,
    /// Rollback operation
    pub operation: RollbackOperation,
}

/// Rollback operation.
#[derive(Clone, Debug)]
pub enum RollbackOperation {
    /// Rollback by renaming back
    Rename { new_name: String },
    /// Rollback by restoring deleted content
    Restore { name: String },
    /// Rollback by deleting created content
    Delete { path: String },
    /// No rollback needed
    None,
}

/// A file region.
#[derive(Clone, Debug)]
struct FileRegion {
    /// File path
    file: String,
    /// Region start
    start: usize,
    /// Region end
    end: usize,
}

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

    #[tokio::test]
    async fn test_planner_creation() {
        let temp_dir = TempDir::new().unwrap();
        let forge = Forge::open(temp_dir.path()).await.unwrap();
        let _planner = Planner::new(forge);

        // Should create successfully
        assert!(true);
    }

    #[tokio::test]
    async fn test_generate_steps() {
        let temp_dir = TempDir::new().unwrap();
        let forge = Forge::open(temp_dir.path()).await.unwrap();
        let planner = Planner::new(forge);

        let observation = crate::observe::Observation {
            query: "test".to_string(),
            symbols: vec![],
        };

        let steps = planner.generate_steps(&observation).await.unwrap();
        // Should succeed (even if empty)
        assert!(steps.is_empty());
    }

    #[tokio::test]
    async fn test_detect_conflicts_empty() {
        let temp_dir = TempDir::new().unwrap();
        let forge = Forge::open(temp_dir.path()).await.unwrap();
        let planner = Planner::new(forge);

        let steps = vec![];
        let conflicts = planner.detect_conflicts(&steps).unwrap();
        assert!(conflicts.is_empty());
    }

    #[tokio::test]
    async fn test_order_steps() {
        let temp_dir = TempDir::new().unwrap();
        let forge = Forge::open(temp_dir.path()).await.unwrap();
        let planner = Planner::new(forge);

        let mut steps = vec![
            PlanStep {
                description: "Delete foo".to_string(),
                operation: PlanOperation::Delete {
                    name: "foo".to_string(),
                },
            },
            PlanStep {
                description: "Rename foo to bar".to_string(),
                operation: PlanOperation::Rename {
                    old: "foo".to_string(),
                    new: "bar".to_string(),
                },
            },
        ];

        planner.order_steps(&mut steps).unwrap();

        // Rename should now come before Delete
        assert!(matches!(steps[0].operation, PlanOperation::Rename { .. }));
        assert!(matches!(steps[1].operation, PlanOperation::Delete { .. }));
    }

    #[tokio::test]
    async fn test_generate_rollback() {
        let temp_dir = TempDir::new().unwrap();
        let forge = Forge::open(temp_dir.path()).await.unwrap();
        let planner = Planner::new(forge);

        let steps = vec![PlanStep {
            description: "Create file".to_string(),
            operation: PlanOperation::Create {
                path: "/tmp/test.rs".to_string(),
                content: "fn test() {}".to_string(),
            },
        }];

        let rollback = planner.generate_rollback(&steps);

        assert_eq!(rollback.len(), 1);
        assert_eq!(rollback[0].description, "Rollback: Create file");
        assert!(matches!(
            rollback[0].operation,
            RollbackOperation::Delete { .. }
        ));
    }

    #[tokio::test]
    async fn test_estimate_impact() {
        let temp_dir = TempDir::new().unwrap();
        let forge = Forge::open(temp_dir.path()).await.unwrap();
        let planner = Planner::new(forge);

        let steps = vec![PlanStep {
            description: "Create test.rs".to_string(),
            operation: PlanOperation::Create {
                path: "/tmp/test.rs".to_string(),
                content: "fn test() {}".to_string(),
            },
        }];

        let impact = planner.estimate_impact(&steps).await.unwrap();

        assert!(!impact.affected_files.is_empty() || impact.affected_files.len() >= 0);
    }
}