axon-lang 1.21.1

//! Step dependency analysis — variable-based dependency graph between steps.
//!
//! Analyzes `$variable` / `${variable}` references in step prompts to build
//! a dependency graph. This enables:
//!   - Detection of which steps can potentially run in parallel
//!   - Validation that referenced variables are actually produced
//!   - Execution plan visualization with dependency chains
//!
//! Built-in variables ($result, $step_name, $flow_name, etc.) are excluded
//! from dependency analysis as they are runtime-injected, not step-produced.

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

// ── Built-in variables (not produced by steps) ─────────────────────────────

const BUILTIN_VARS: &[&str] = &[
    "result",
    "step_name",
    "step_type",
    "flow_name",
    "persona_name",
    "unit_index",
    "step_index",
];

fn is_builtin(var: &str) -> bool {
    BUILTIN_VARS.contains(&var)
}

// ── Variable extraction ────────────────────────────────────────────────────

/// Extract all variable references from a string.
/// Returns the set of variable names referenced via $name or ${name}.
pub fn extract_refs(text: &str) -> HashSet<String> {
    let mut refs = HashSet::new();
    let bytes = text.as_bytes();
    let mut i = 0;

    while i < bytes.len() {
        if bytes[i] == b'$' && i + 1 < bytes.len() {
            if bytes[i + 1] == b'{' {
                // ${name} form
                if let Some(close) = text[i + 2..].find('}') {
                    let var_name = &text[i + 2..i + 2 + close];
                    if !var_name.is_empty() {
                        refs.insert(var_name.to_string());
                    }
                    i += 3 + close;
                    continue;
                }
            } else if bytes[i + 1].is_ascii_alphabetic() || bytes[i + 1] == b'_' {
                // $name form
                let start = i + 1;
                let mut end = start;
                while end < bytes.len()
                    && (bytes[end].is_ascii_alphanumeric() || bytes[end] == b'_')
                {
                    end += 1;
                }
                let var_name = &text[start..end];
                refs.insert(var_name.to_string());
                i = end;
                continue;
            }
        }
        i += 1;
    }

    refs
}

// ── Step info for analysis ─────────────────────────────────────────────────

/// Minimal step representation for dependency analysis.
#[derive(Debug, Clone)]
pub struct StepInfo {
    pub name: String,
    pub step_type: String,
    pub user_prompt: String,
    /// For tool/memory steps: the argument expression.
    pub argument: String,
}

// ── Dependency analysis result ─────────────────────────────────────────────

/// Analysis result for a single step.
#[derive(Debug, Clone)]
pub struct StepDependency {
    /// Step name.
    pub name: String,
    /// Step type.
    pub step_type: String,
    /// Steps this step depends on (via variable references).
    pub depends_on: Vec<String>,
    /// All variable references found (including builtins).
    pub all_refs: Vec<String>,
    /// Variable references that are step-produced (non-builtin).
    pub step_refs: Vec<String>,
    /// Whether this step has no step dependencies (can run first).
    pub is_root: bool,
}

/// Full dependency graph for a unit's steps.
#[derive(Debug)]
pub struct DependencyGraph {
    pub steps: Vec<StepDependency>,
    /// Steps that can potentially run in parallel (no mutual dependencies).
    pub parallel_groups: Vec<Vec<String>>,
    /// Steps that reference undefined variables (not produced by any prior step).
    pub unresolved_refs: Vec<(String, String)>,
    /// Maximum depth of the dependency chain.
    pub max_depth: usize,
}

// ── Analysis ───────────────────────────────────────────────────────────────

/// Analyze dependencies between steps in a unit.
pub fn analyze(steps: &[StepInfo]) -> DependencyGraph {
    // 1. Build the set of step names (these are the "producers")
    let step_names: HashSet<&str> = steps.iter().map(|s| s.name.as_str()).collect();

    // 2. For each step, extract variable refs and resolve dependencies
    let mut deps: Vec<StepDependency> = Vec::new();
    let mut unresolved: Vec<(String, String)> = Vec::new();

    for step in steps {
        // Scan both user_prompt and argument for references
        let mut all_refs: HashSet<String> = extract_refs(&step.user_prompt);
        if !step.argument.is_empty() {
            all_refs.extend(extract_refs(&step.argument));
        }

        let mut step_refs: Vec<String> = Vec::new();
        let mut depends_on: Vec<String> = Vec::new();

        for r in &all_refs {
            if is_builtin(r) {
                continue;
            }
            if step_names.contains(r.as_str()) {
                step_refs.push(r.clone());
                depends_on.push(r.clone());
            } else {
                unresolved.push((step.name.clone(), r.clone()));
            }
        }

        depends_on.sort();
        depends_on.dedup();
        step_refs.sort();

        let mut all_refs_sorted: Vec<String> = all_refs.into_iter().collect();
        all_refs_sorted.sort();

        deps.push(StepDependency {
            name: step.name.clone(),
            step_type: step.step_type.clone(),
            is_root: depends_on.is_empty(),
            depends_on,
            all_refs: all_refs_sorted,
            step_refs,
        });
    }

    // 3. Detect parallel groups (steps with no mutual dependencies)
    let parallel_groups = find_parallel_groups(&deps);

    // 4. Calculate max depth
    let max_depth = calculate_max_depth(&deps);

    DependencyGraph {
        steps: deps,
        parallel_groups,
        unresolved_refs: unresolved,
        max_depth,
    }
}

/// Find groups of steps that can potentially execute in parallel.
/// Steps at the same depth level with no mutual dependencies form a group.
fn find_parallel_groups(deps: &[StepDependency]) -> Vec<Vec<String>> {
    // Build transitive dependency sets via depth calculation
    let dep_map: HashMap<&str, &StepDependency> =
        deps.iter().map(|d| (d.name.as_str(), d)).collect();

    // Calculate depth for each step
    let mut depth_cache: HashMap<String, usize> = HashMap::new();
    fn step_depth(
        name: &str,
        dep_map: &HashMap<&str, &StepDependency>,
        cache: &mut HashMap<String, usize>,
    ) -> usize {
        if let Some(&cached) = cache.get(name) {
            return cached;
        }
        let d = match dep_map.get(name) {
            Some(d) => d,
            None => return 0,
        };
        if d.depends_on.is_empty() {
            cache.insert(name.to_string(), 0);
            return 0;
        }
        let max_child = d
            .depends_on
            .iter()
            .map(|dep| step_depth(dep, dep_map, cache))
            .max()
            .unwrap_or(0);
        let result = max_child + 1;
        cache.insert(name.to_string(), result);
        result
    }

    for d in deps {
        step_depth(&d.name, &dep_map, &mut depth_cache);
    }

    // Group steps by depth level
    let mut by_depth: HashMap<usize, Vec<String>> = HashMap::new();
    for d in deps {
        let depth = depth_cache.get(&d.name).copied().unwrap_or(0);
        by_depth.entry(depth).or_default().push(d.name.clone());
    }

    // Return only groups with more than one step (actual parallelism)
    let mut groups: Vec<Vec<String>> = by_depth
        .into_values()
        .filter(|g| g.len() > 1)
        .collect();
    groups.sort_by_key(|g| g[0].clone());
    groups
}

/// Calculate the maximum dependency chain depth.
fn calculate_max_depth(deps: &[StepDependency]) -> usize {
    let dep_map: HashMap<&str, &StepDependency> =
        deps.iter().map(|d| (d.name.as_str(), d)).collect();

    fn depth(
        name: &str,
        dep_map: &HashMap<&str, &StepDependency>,
        cache: &mut HashMap<String, usize>,
    ) -> usize {
        if let Some(&cached) = cache.get(name) {
            return cached;
        }
        let d = match dep_map.get(name) {
            Some(d) => d,
            None => return 0,
        };
        if d.depends_on.is_empty() {
            cache.insert(name.to_string(), 0);
            return 0;
        }
        let max_child = d
            .depends_on
            .iter()
            .map(|dep| depth(dep, dep_map, cache))
            .max()
            .unwrap_or(0);
        let result = max_child + 1;
        cache.insert(name.to_string(), result);
        result
    }

    let mut cache = HashMap::new();
    deps.iter()
        .map(|d| depth(&d.name, &dep_map, &mut cache))
        .max()
        .unwrap_or(0)
}

// ── Tests ──────────────────────────────────────────────────────────────────

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

    #[test]
    fn extract_refs_dollar_name() {
        let refs = extract_refs("Use $result from $Analyze");
        assert!(refs.contains("result"));
        assert!(refs.contains("Analyze"));
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn extract_refs_braced() {
        let refs = extract_refs("Given ${Extract} and ${Validate}");
        assert!(refs.contains("Extract"));
        assert!(refs.contains("Validate"));
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn extract_refs_mixed() {
        let refs = extract_refs("$result is ${Analyze} plus $flow_name");
        assert!(refs.contains("result"));
        assert!(refs.contains("Analyze"));
        assert!(refs.contains("flow_name"));
        assert_eq!(refs.len(), 3);
    }

    #[test]
    fn extract_refs_no_vars() {
        let refs = extract_refs("plain text with no variables");
        assert!(refs.is_empty());
    }

    #[test]
    fn extract_refs_dollar_at_end() {
        let refs = extract_refs("trailing $");
        assert!(refs.is_empty());
    }

    #[test]
    fn analyze_independent_steps() {
        let steps = vec![
            StepInfo {
                name: "A".into(),
                step_type: "step".into(),
                user_prompt: "Do task A".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "B".into(),
                step_type: "step".into(),
                user_prompt: "Do task B".into(),
                argument: String::new(),
            },
        ];

        let graph = analyze(&steps);
        assert_eq!(graph.steps.len(), 2);
        assert!(graph.steps[0].is_root);
        assert!(graph.steps[1].is_root);
        assert_eq!(graph.max_depth, 0);
        // Both independent → one parallel group
        assert_eq!(graph.parallel_groups.len(), 1);
        assert_eq!(graph.parallel_groups[0].len(), 2);
    }

    #[test]
    fn analyze_linear_chain() {
        let steps = vec![
            StepInfo {
                name: "Extract".into(),
                step_type: "step".into(),
                user_prompt: "Extract entities".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "Analyze".into(),
                step_type: "step".into(),
                user_prompt: "Analyze ${Extract}".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "Report".into(),
                step_type: "step".into(),
                user_prompt: "Report on ${Analyze}".into(),
                argument: String::new(),
            },
        ];

        let graph = analyze(&steps);

        // Extract is root
        assert!(graph.steps[0].is_root);
        assert!(graph.steps[0].depends_on.is_empty());

        // Analyze depends on Extract
        assert!(!graph.steps[1].is_root);
        assert_eq!(graph.steps[1].depends_on, vec!["Extract"]);

        // Report depends on Analyze
        assert!(!graph.steps[2].is_root);
        assert_eq!(graph.steps[2].depends_on, vec!["Analyze"]);

        // Max depth is 2 (Extract→Analyze→Report)
        assert_eq!(graph.max_depth, 2);

        // No parallel groups (all sequential)
        assert!(graph.parallel_groups.is_empty());
    }

    #[test]
    fn analyze_diamond_pattern() {
        // A → B, A → C, B+C → D
        let steps = vec![
            StepInfo {
                name: "A".into(),
                step_type: "step".into(),
                user_prompt: "Start".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "B".into(),
                step_type: "step".into(),
                user_prompt: "Process ${A} path B".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "C".into(),
                step_type: "step".into(),
                user_prompt: "Process ${A} path C".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "D".into(),
                step_type: "step".into(),
                user_prompt: "Merge ${B} and ${C}".into(),
                argument: String::new(),
            },
        ];

        let graph = analyze(&steps);

        assert!(graph.steps[0].is_root); // A
        assert_eq!(graph.steps[1].depends_on, vec!["A"]); // B→A
        assert_eq!(graph.steps[2].depends_on, vec!["A"]); // C→A
        assert_eq!(graph.steps[3].depends_on, vec!["B", "C"]); // D→B,C

        // B and C can be parallel
        assert!(!graph.parallel_groups.is_empty());
        let has_bc_group = graph.parallel_groups.iter().any(|g| {
            g.len() == 2 && g.contains(&"B".to_string()) && g.contains(&"C".to_string())
        });
        assert!(has_bc_group);

        // Max depth: A→B→D or A→C→D = 2
        assert_eq!(graph.max_depth, 2);
    }

    #[test]
    fn analyze_builtin_vars_excluded() {
        let steps = vec![
            StepInfo {
                name: "S1".into(),
                step_type: "step".into(),
                user_prompt: "Current step is $step_name in $flow_name".into(),
                argument: String::new(),
            },
        ];

        let graph = analyze(&steps);
        assert!(graph.steps[0].is_root);
        assert!(graph.steps[0].depends_on.is_empty());
        // All refs include builtins
        assert!(graph.steps[0].all_refs.contains(&"step_name".to_string()));
        assert!(graph.steps[0].all_refs.contains(&"flow_name".to_string()));
        // But step_refs is empty (no step-produced refs)
        assert!(graph.steps[0].step_refs.is_empty());
    }

    #[test]
    fn analyze_unresolved_refs() {
        let steps = vec![
            StepInfo {
                name: "S1".into(),
                step_type: "step".into(),
                user_prompt: "Use ${NonExistent} data".into(),
                argument: String::new(),
            },
        ];

        let graph = analyze(&steps);
        assert_eq!(graph.unresolved_refs.len(), 1);
        assert_eq!(graph.unresolved_refs[0], ("S1".to_string(), "NonExistent".to_string()));
    }

    #[test]
    fn analyze_argument_refs() {
        let steps = vec![
            StepInfo {
                name: "Gather".into(),
                step_type: "step".into(),
                user_prompt: "Gather data".into(),
                argument: String::new(),
            },
            StepInfo {
                name: "Calc".into(),
                step_type: "use_tool".into(),
                user_prompt: "Calculate".into(),
                argument: "${Gather}".into(),
            },
        ];

        let graph = analyze(&steps);
        assert_eq!(graph.steps[1].depends_on, vec!["Gather"]);
    }

    #[test]
    fn analyze_empty_steps() {
        let graph = analyze(&[]);
        assert!(graph.steps.is_empty());
        assert!(graph.parallel_groups.is_empty());
        assert_eq!(graph.max_depth, 0);
    }

    #[test]
    fn max_depth_flat() {
        let steps = vec![
            StepInfo { name: "A".into(), step_type: "step".into(), user_prompt: "a".into(), argument: String::new() },
            StepInfo { name: "B".into(), step_type: "step".into(), user_prompt: "b".into(), argument: String::new() },
            StepInfo { name: "C".into(), step_type: "step".into(), user_prompt: "c".into(), argument: String::new() },
        ];
        assert_eq!(analyze(&steps).max_depth, 0);
    }
}