// vapora-backend: Workflow dependency scheduler
// Phase 3: Topological sort for dependency resolution and parallel execution

use std::collections::{HashMap, VecDeque};

use thiserror::Error;

use crate::workflow::state::WorkflowStep;

#[derive(Debug, Error)]
pub enum SchedulerError {
    #[error("Circular dependency detected in workflow")]
    CircularDependency,

    #[error("Invalid step reference: {0}")]
    InvalidStepReference(String),
}

pub struct Scheduler;

impl Scheduler {
    /// Resolve dependencies using topological sort (Kahn's algorithm)
    /// Returns levels of steps that can be executed in parallel
    pub fn resolve_dependencies(
        steps: &[WorkflowStep],
    ) -> Result<Vec<Vec<String>>, SchedulerError> {
        if steps.is_empty() {
            return Ok(Vec::new());
        }

        // Build dependency graph
        let mut graph: HashMap<String, Vec<String>> = HashMap::new();
        let mut in_degree: HashMap<String, usize> = HashMap::new();

        // Initialize graph with all steps
        for step in steps {
            in_degree.insert(step.id.clone(), step.depends_on.len());
            graph.insert(step.id.clone(), Vec::new());
        }

        // Build adjacency list (reverse dependencies)
        for step in steps {
            for dep in &step.depends_on {
                // Validate dependency exists
                if !in_degree.contains_key(dep) {
                    return Err(SchedulerError::InvalidStepReference(dep.clone()));
                }

                // Add edge from dependency to dependent
                if let Some(dependents) = graph.get_mut(dep) {
                    dependents.push(step.id.clone());
                }
            }
        }

        // Kahn's algorithm for topological sort
        let mut queue: VecDeque<String> = VecDeque::new();

        // Find all steps with no dependencies
        for (step_id, degree) in &in_degree {
            if *degree == 0 {
                queue.push_back(step_id.clone());
            }
        }

        let mut sorted_levels: Vec<Vec<String>> = Vec::new();
        let mut processed_count = 0;

        while !queue.is_empty() {
            let mut current_level = Vec::new();
            let level_size = queue.len();

            // Process all steps at current level
            for _ in 0..level_size {
                if let Some(step_id) = queue.pop_front() {
                    current_level.push(step_id.clone());
                    processed_count += 1;

                    // Reduce in-degree for all dependents
                    if let Some(dependents) = graph.get(&step_id) {
                        for dependent in dependents {
                            if let Some(degree) = in_degree.get_mut(dependent) {
                                *degree -= 1;

                                // If in-degree becomes 0, add to queue
                                if *degree == 0 {
                                    queue.push_back(dependent.clone());
                                }
                            }
                        }
                    }
                }
            }

            if !current_level.is_empty() {
                sorted_levels.push(current_level);
            }
        }

        // Check for circular dependencies
        if processed_count != steps.len() {
            return Err(SchedulerError::CircularDependency);
        }

        Ok(sorted_levels)
    }

    /// Get steps that can be executed in parallel at each level
    pub fn get_parallel_groups(steps: &[WorkflowStep]) -> Result<Vec<Vec<String>>, SchedulerError> {
        let sorted_levels = Self::resolve_dependencies(steps)?;

        // Filter to only include parallelizable steps
        let parallel_groups = sorted_levels
            .into_iter()
            .map(|level| {
                level
                    .into_iter()
                    .filter(|step_id| {
                        steps
                            .iter()
                            .find(|s| &s.id == step_id)
                            .map(|s| s.can_parallelize)
                            .unwrap_or(false)
                    })
                    .collect()
            })
            .filter(|level: &Vec<String>| !level.is_empty())
            .collect();

        Ok(parallel_groups)
    }

    /// Get execution order for sequential execution
    pub fn get_sequential_order(steps: &[WorkflowStep]) -> Result<Vec<String>, SchedulerError> {
        let levels = Self::resolve_dependencies(steps)?;
        Ok(levels.into_iter().flatten().collect())
    }

    /// Validate workflow has no circular dependencies
    pub fn validate_workflow(steps: &[WorkflowStep]) -> Result<(), SchedulerError> {
        Self::resolve_dependencies(steps)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::workflow::state::StepStatus;

    fn create_step(id: &str, deps: Vec<&str>, parallel: bool) -> WorkflowStep {
        WorkflowStep {
            id: id.to_string(),
            name: format!("Step {}", id),
            agent_role: "developer".to_string(),
            status: StepStatus::Pending,
            depends_on: deps.iter().map(|s| s.to_string()).collect(),
            can_parallelize: parallel,
            started_at: None,
            completed_at: None,
            result: None,
            error: None,
        }
    }

    #[test]
    fn test_simple_dependency_chain() {
        let steps = vec![
            create_step("a", vec![], true),
            create_step("b", vec!["a"], true),
            create_step("c", vec!["b"], true),
        ];

        let result = Scheduler::resolve_dependencies(&steps);
        assert!(result.is_ok());

        let sorted = result.unwrap();
        assert_eq!(sorted.len(), 3);
        assert_eq!(sorted[0], vec!["a"]);
        assert_eq!(sorted[1], vec!["b"]);
        assert_eq!(sorted[2], vec!["c"]);
    }

    #[test]
    fn test_parallel_execution() {
        let steps = vec![
            create_step("a", vec![], true),
            create_step("b", vec!["a"], true),
            create_step("c", vec!["a"], true),
            create_step("d", vec!["b", "c"], true),
        ];

        let result = Scheduler::resolve_dependencies(&steps);
        assert!(result.is_ok());

        let sorted = result.unwrap();
        assert_eq!(sorted.len(), 3);
        assert_eq!(sorted[0], vec!["a"]);
        assert_eq!(sorted[1].len(), 2); // b and c can run in parallel
        assert!(sorted[1].contains(&"b".to_string()));
        assert!(sorted[1].contains(&"c".to_string()));
        assert_eq!(sorted[2], vec!["d"]);
    }

    #[test]
    fn test_circular_dependency() {
        let steps = vec![
            create_step("a", vec!["b"], true),
            create_step("b", vec!["c"], true),
            create_step("c", vec!["a"], true),
        ];

        let result = Scheduler::resolve_dependencies(&steps);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            SchedulerError::CircularDependency
        ));
    }

    #[test]
    fn test_invalid_dependency_reference() {
        let steps = vec![
            create_step("a", vec![], true),
            create_step("b", vec!["nonexistent"], true),
        ];

        let result = Scheduler::resolve_dependencies(&steps);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            SchedulerError::InvalidStepReference(_)
        ));
    }

    #[test]
    fn test_complex_dag() {
        //     a
        //    / \
        //   b   c
        //   |   |
        //   d   e
        //    \ /
        //     f
        let steps = vec![
            create_step("a", vec![], true),
            create_step("b", vec!["a"], true),
            create_step("c", vec!["a"], true),
            create_step("d", vec!["b"], true),
            create_step("e", vec!["c"], true),
            create_step("f", vec!["d", "e"], true),
        ];

        let result = Scheduler::resolve_dependencies(&steps);
        assert!(result.is_ok());

        let sorted = result.unwrap();
        assert_eq!(sorted.len(), 4);
        assert_eq!(sorted[0], vec!["a"]);
        assert_eq!(sorted[1].len(), 2); // b and c parallel
        assert_eq!(sorted[2].len(), 2); // d and e parallel
        assert_eq!(sorted[3], vec!["f"]);
    }

    #[test]
    fn test_parallel_groups_filter() {
        let steps = vec![
            create_step("a", vec![], true),
            create_step("b", vec!["a"], false), // Not parallelizable
            create_step("c", vec!["a"], true),
        ];

        let result = Scheduler::get_parallel_groups(&steps);
        assert!(result.is_ok());

        let groups = result.unwrap();
        // Only "a" and "c" should be in parallel groups
        assert!(groups[0].contains(&"a".to_string()));
        assert!(groups[1].contains(&"c".to_string()));
        assert!(!groups[1].contains(&"b".to_string()));
    }

    #[test]
    fn test_sequential_order() {
        let steps = vec![
            create_step("c", vec!["a"], true),
            create_step("a", vec![], true),
            create_step("b", vec!["a"], true),
        ];

        let result = Scheduler::get_sequential_order(&steps);
        assert!(result.is_ok());

        let order = result.unwrap();
        assert_eq!(order.len(), 3);
        assert_eq!(order[0], "a");
        // b and c can be in any order after a
        assert!(order[1..].contains(&"b".to_string()));
        assert!(order[1..].contains(&"c".to_string()));
    }

    #[test]
    fn test_empty_steps() {
        let steps: Vec<WorkflowStep> = vec![];
        let result = Scheduler::resolve_dependencies(&steps);
        assert!(result.is_ok());
        assert!(result.unwrap().is_empty());
    }
}
