openfang/crates/openfang-kernel/src/aol/executor.rs

//! AOL Workflow Execution Engine.
//!
//! Executes compiled AOL workflows, handling parallel execution,
//! conditional branching, loops, and error handling.

use crate::aol::template::{expand_template, TemplateContext};
use crate::aol::validator::validate_workflow;
use crate::aol::{AolError, AolResult, CompiledWorkflow};
use futures::future::BoxFuture;
use openfang_types::aol::{
    AolStep, AgentRef, CollectStrategy, ErrorMode, ParallelStepGroup, WorkflowDefId,
};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
use uuid::Uuid;

/// Unique identifier for a workflow execution instance.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ExecutionId(pub Uuid);

impl ExecutionId {
    /// Generate a new execution ID.
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }
}

impl Default for ExecutionId {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Display for ExecutionId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Status of a workflow execution.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExecutionStatus {
    /// Execution is pending.
    Pending,
    /// Execution is running.
    Running,
    /// Execution completed successfully.
    Completed,
    /// Execution failed.
    Failed,
    /// Execution was cancelled.
    Cancelled,
}

/// Result of a single step execution.
#[derive(Debug, Clone)]
pub struct StepExecutionResult {
    /// Step ID.
    pub step_id: String,
    /// Agent that executed the step.
    pub agent_id: Option<String>,
    /// Output from the step.
    pub output: Value,
    /// Whether the step succeeded.
    pub success: bool,
    /// Error message if failed.
    pub error: Option<String>,
    /// Duration of execution.
    pub duration_ms: u64,
    /// Number of retries attempted.
    pub retries: u32,
}

/// Result of a workflow execution.
#[derive(Debug, Clone)]
pub struct ExecutionResult {
    /// Execution ID.
    pub id: ExecutionId,
    /// Workflow definition ID.
    pub workflow_id: WorkflowDefId,
    /// Execution status.
    pub status: ExecutionStatus,
    /// Step results.
    pub step_results: Vec<StepExecutionResult>,
    /// Final output variables.
    pub outputs: HashMap<String, Value>,
    /// Error message if failed.
    pub error: Option<String>,
    /// Total execution time.
    pub duration_ms: u64,
    /// Started at.
    pub started_at: Instant,
    /// Completed at.
    pub completed_at: Option<Instant>,
}

impl ExecutionResult {
    /// Create a new pending execution result.
    pub fn new(id: ExecutionId, workflow_id: WorkflowDefId) -> Self {
        Self {
            id,
            workflow_id,
            status: ExecutionStatus::Pending,
            step_results: Vec::new(),
            outputs: HashMap::new(),
            error: None,
            duration_ms: 0,
            started_at: Instant::now(),
            completed_at: None,
        }
    }
}

/// Trait for executing agent tasks.
#[async_trait::async_trait]
pub trait AgentExecutor: Send + Sync {
    /// Execute a task on an agent.
    async fn execute(
        &self,
        agent_ref: &AgentRef,
        task: &str,
        inputs: &HashMap<String, String>,
        timeout_secs: u64,
    ) -> AolResult<Value>;
}

/// Default agent executor that returns mock results.
pub struct MockAgentExecutor;

#[async_trait::async_trait]
impl AgentExecutor for MockAgentExecutor {
    async fn execute(
        &self,
        agent_ref: &AgentRef,
        task: &str,
        _inputs: &HashMap<String, String>,
        _timeout_secs: u64,
    ) -> AolResult<Value> {
        // Mock implementation - in real usage, this would call the kernel
        let agent_name = match agent_ref {
            AgentRef::ById { id } => id.to_string(),
            AgentRef::ByName { name } => name.clone(),
            AgentRef::ByRole { role, .. } => format!("role:{}", role),
        };

        Ok(Value::String(format!(
            "Mock response from {} for task: {}",
            agent_name,
            task.chars().take(50).collect::<String>()
        )))
    }
}

/// The AOL workflow executor.
pub struct AolExecutor {
    /// Agent executor implementation.
    agent_executor: Arc<dyn AgentExecutor>,
    /// Active executions.
    executions: Arc<RwLock<HashMap<ExecutionId, ExecutionResult>>>,
    /// Default timeout in seconds.
    default_timeout_secs: u64,
    /// Maximum retry attempts.
    max_retries: u32,
}

impl AolExecutor {
    /// Create a new executor with a custom agent executor.
    pub fn new(agent_executor: Arc<dyn AgentExecutor>) -> Self {
        Self {
            agent_executor,
            executions: Arc::new(RwLock::new(HashMap::new())),
            default_timeout_secs: 300,
            max_retries: 3,
        }
    }

    /// Create a new executor with mock agent executor.
    pub fn with_mock() -> Self {
        Self::new(Arc::new(MockAgentExecutor))
    }

    /// Set default timeout.
    pub fn with_timeout(mut self, timeout_secs: u64) -> Self {
        self.default_timeout_secs = timeout_secs;
        self
    }

    /// Set max retries.
    pub fn with_max_retries(mut self, max_retries: u32) -> Self {
        self.max_retries = max_retries;
        self
    }

    /// Execute a compiled workflow.
    pub async fn execute(
        &self,
        workflow: &CompiledWorkflow,
        inputs: HashMap<String, Value>,
    ) -> AolResult<ExecutionResult> {
        // Validate if not already validated
        if !workflow.validated {
            validate_workflow(&workflow.workflow)?;
        }

        let exec_id = ExecutionId::new();
        let mut result = ExecutionResult::new(exec_id, workflow.id);
        result.status = ExecutionStatus::Running;

        // Store execution
        self.executions.write().await.insert(exec_id, result.clone());

        // Create template context
        let mut ctx = TemplateContext::new();
        for (k, v) in &inputs {
            ctx.add_input(k, v.clone());
        }

        // Execute steps
        let start_time = Instant::now();
        match self.execute_steps(&workflow.workflow.steps, &mut ctx, &mut result).await {
            Ok(()) => {
                result.status = ExecutionStatus::Completed;
                result.outputs = ctx.outputs.clone();
            }
            Err(e) => {
                result.status = ExecutionStatus::Failed;
                result.error = Some(e.to_string());
            }
        }
        result.duration_ms = start_time.elapsed().as_millis() as u64;
        result.completed_at = Some(Instant::now());

        // Update stored execution
        self.executions.write().await.insert(exec_id, result.clone());

        info!(
            execution_id = %exec_id,
            workflow_id = %workflow.id,
            status = ?result.status,
            duration_ms = result.duration_ms,
            "Workflow execution completed"
        );

        Ok(result)
    }

    /// Execute a list of steps (boxed for recursion).
    fn execute_steps<'a>(
        &'a self,
        steps: &'a [AolStep],
        ctx: &'a mut TemplateContext,
        result: &'a mut ExecutionResult,
    ) -> BoxFuture<'a, AolResult<()>> {
        Box::pin(async move {
            for step in steps {
                self.execute_step(step, ctx, result).await?;
            }
            Ok(())
        })
    }

    /// Execute a single step.
    async fn execute_step(
        &self,
        step: &AolStep,
        ctx: &mut TemplateContext,
        result: &mut ExecutionResult,
    ) -> AolResult<()> {
        match step {
            AolStep::Parallel(pg) => self.execute_parallel(pg, ctx, result).await,
            AolStep::Sequential(ss) => {
                let step_result = self.execute_agent_step(
                    &ss.id,
                    &ss.agent,
                    &ss.task,
                    &ss.inputs,
                    ss.error_mode,
                    ss.timeout_secs,
                    ctx,
                ).await?;

                if let Some(output) = &ss.output {
                    ctx.set_output(output, step_result.output.clone());
                }
                result.step_results.push(step_result);
                Ok(())
            }
            AolStep::Conditional(cs) => {
                // Evaluate branches in order
                for branch in &cs.branches {
                    if self.evaluate_condition(&branch.condition, ctx)? {
                        debug!(branch_id = %branch.id, "Condition matched");
                        self.execute_steps(&branch.steps, ctx, result).await?;

                        if let Some(output) = &branch.output {
                            if let Some(last_result) = result.step_results.last() {
                                ctx.set_output(output, last_result.output.clone());
                            }
                        }
                        return Ok(());
                    }
                }

                // No branch matched, execute default if present
                if let Some(default_steps) = &cs.default {
                    debug!("Executing default branch");
                    self.execute_steps(default_steps, ctx, result).await?;

                    if let Some(output) = &cs.output {
                        if let Some(last_result) = result.step_results.last() {
                            ctx.set_output(output, last_result.output.clone());
                        }
                    }
                }

                Ok(())
            }
            AolStep::Loop(ls) => {
                let collection = self.evaluate_collection(&ls.collection, ctx)?;
                let mut all_results = Vec::new();

                for (index, item) in collection.iter().enumerate() {
                    // Set loop variables
                    ctx.add_loop_var(&ls.item_var, item.clone());
                    if let Some(index_var) = &ls.index_var {
                        ctx.add_loop_var(index_var, Value::Number(serde_json::Number::from(index as i64)));
                    }

                    // Execute loop body
                    self.execute_steps(&ls.steps, ctx, result).await?;

                    // Collect result
                    if let Some(last) = result.step_results.last() {
                        all_results.push(last.output.clone());
                    }
                }

                // Apply collect strategy
                let collected = apply_collect_strategy(&ls.collect, all_results);
                if let Some(output) = &ls.output {
                    ctx.set_output(output, collected);
                }

                Ok(())
            }
            AolStep::Collect(cs) => {
                let mut values = Vec::new();
                for source in &cs.sources {
                    if let Some(value) = ctx.get(source) {
                        values.push(value.clone());
                    } else {
                        return Err(AolError::Execution(format!(
                            "Collect source '{}' not found",
                            source
                        )));
                    }
                }

                let collected = apply_collect_strategy(&cs.strategy, values);
                ctx.set_output(&cs.output, collected);
                Ok(())
            }
            AolStep::Subworkflow(ss) => {
                // In a real implementation, this would recursively execute another workflow
                warn!(step_id = %ss.id, "Subworkflow execution not fully implemented");
                ctx.set_output(
                    ss.output.as_deref().unwrap_or(&ss.id),
                    Value::String(format!("Subworkflow {} result", ss.workflow)),
                );
                Ok(())
            }
            AolStep::Fallback(fs) => {
                // Try primary step first, then fallbacks in order
                // For each step, we execute it directly without further fallback recursion
                let mut steps_to_try: Vec<&AolStep> = vec![&fs.primary];
                steps_to_try.extend(fs.fallbacks.iter());

                let mut last_error: Option<AolError> = None;
                for fallback_step in steps_to_try {
                    // Execute the step directly without fallback handling
                    let exec_result = match fallback_step {
                        AolStep::Parallel(pg) => self.execute_parallel(pg, ctx, result).await,
                        AolStep::Sequential(ss) => {
                            let step_result = self.execute_agent_step(
                                &ss.id,
                                &ss.agent,
                                &ss.task,
                                &ss.inputs,
                                ss.error_mode,
                                ss.timeout_secs,
                                ctx,
                            ).await;
                            match step_result {
                                Ok(sr) => {
                                    if let Some(output) = &ss.output {
                                        ctx.set_output(output, sr.output.clone());
                                    }
                                    result.step_results.push(sr);
                                    Ok(())
                                }
                                Err(e) => Err(e),
                            }
                        }
                        AolStep::Conditional(cs) => {
                            let mut matched = false;
                            for branch in &cs.branches {
                                if self.evaluate_condition(&branch.condition, ctx)? {
                                    self.execute_steps(&branch.steps, ctx, result).await?;
                                    if let Some(output) = &branch.output {
                                        if let Some(last_result) = result.step_results.last() {
                                            ctx.set_output(output, last_result.output.clone());
                                        }
                                    }
                                    matched = true;
                                    break;
                                }
                            }
                            if !matched {
                                if let Some(default_steps) = &cs.default {
                                    self.execute_steps(default_steps, ctx, result).await?;
                                }
                            }
                            Ok(())
                        }
                        AolStep::Loop(ls) => {
                            let collection = self.evaluate_collection(&ls.collection, ctx)?;
                            for (index, item) in collection.iter().enumerate() {
                                ctx.add_loop_var(&ls.item_var, item.clone());
                                if let Some(index_var) = &ls.index_var {
                                    ctx.add_loop_var(index_var, Value::Number(serde_json::Number::from(index as i64)));
                                }
                                self.execute_steps(&ls.steps, ctx, result).await?;
                            }
                            Ok(())
                        }
                        AolStep::Collect(cs) => {
                            let mut values = Vec::new();
                            for source in &cs.sources {
                                if let Some(value) = ctx.get(source) {
                                    values.push(value.clone());
                                }
                            }
                            let collected = apply_collect_strategy(&cs.strategy, values);
                            ctx.set_output(&cs.output, collected);
                            Ok(())
                        }
                        AolStep::Subworkflow(ss) => {
                            ctx.set_output(
                                ss.output.as_deref().unwrap_or(&ss.id),
                                Value::String(format!("Subworkflow {} result", ss.workflow)),
                            );
                            Ok(())
                        }
                        AolStep::Fallback(nested_fs) => {
                            // For nested fallback, just try the primary of the nested fallback
                            // This prevents infinite recursion
                            self.execute_steps(std::slice::from_ref(&nested_fs.primary), ctx, result).await
                        }
                    };

                    match exec_result {
                        Ok(()) => {
                            if let Some(output) = &fs.output {
                                if let Some(last_result) = result.step_results.last() {
                                    ctx.set_output(output, last_result.output.clone());
                                }
                            }
                            return Ok(());
                        }
                        Err(e) => {
                            warn!(error = %e, "Fallback step failed, trying next");
                            last_error = Some(e);
                        }
                    }
                }

                // All fallbacks failed
                Err(last_error.unwrap_or_else(|| AolError::Execution(format!(
                    "All fallbacks failed for step {}",
                    fs.id
                ))))
            }
        }
    }

    /// Execute a parallel step group.
    async fn execute_parallel(
        &self,
        pg: &ParallelStepGroup,
        ctx: &mut TemplateContext,
        result: &mut ExecutionResult,
    ) -> AolResult<()> {
        let concurrency = pg.max_concurrency.unwrap_or(10);
        let steps = &pg.steps;

        // For simplicity, execute in batches based on concurrency
        let mut all_results = Vec::new();
        let mut batch_results = Vec::new();

        for chunk in steps.chunks(concurrency) {
            let mut tasks = Vec::new();

            for step in chunk {
                let task = self.execute_agent_step(
                    &step.id,
                    &step.agent,
                    &step.task,
                    &step.inputs,
                    step.error_mode,
                    step.timeout_secs,
                    ctx,
                );
                tasks.push(task);
            }

            // Execute batch concurrently
            let batch = futures::future::join_all(tasks).await;

            for task_result in batch {
                let step_result = task_result?;
                batch_results.push(step_result.output.clone());
                result.step_results.push(step_result);
            }
        }

        // Apply collect strategy
        all_results.extend(batch_results);
        let collected = apply_collect_strategy(&pg.collect, all_results);

        if let Some(output) = &pg.output {
            ctx.set_output(output, collected);
        }

        Ok(())
    }

    /// Execute a single agent step with retry support.
    async fn execute_agent_step(
        &self,
        step_id: &str,
        agent: &AgentRef,
        task_template: &str,
        inputs: &HashMap<String, String>,
        error_mode: Option<ErrorMode>,
        timeout_secs: Option<u64>,
        ctx: &TemplateContext,
    ) -> AolResult<StepExecutionResult> {
        // Expand task template
        let task = expand_template(task_template, ctx)?;
        let expanded_inputs = crate::aol::template::expand_templates_in_map(inputs, ctx)?;

        let timeout = timeout_secs.unwrap_or(self.default_timeout_secs);
        let error_mode = error_mode.unwrap_or(ErrorMode::Fail);
        let max_retries = if matches!(error_mode, ErrorMode::Retry) {
            self.max_retries
        } else {
            0
        };

        let start_time = Instant::now();
        let mut retries = 0;
        #[allow(unused_assignments)]
        let mut last_error = None;

        loop {
            match self.agent_executor.execute(agent, &task, &expanded_inputs, timeout).await {
                Ok(output) => {
                    return Ok(StepExecutionResult {
                        step_id: step_id.to_string(),
                        agent_id: None,
                        output,
                        success: true,
                        error: None,
                        duration_ms: start_time.elapsed().as_millis() as u64,
                        retries,
                    });
                }
                Err(e) => {
                    last_error = Some(e.to_string());

                    if retries < max_retries {
                        retries += 1;
                        debug!(step_id = %step_id, retry = retries, "Retrying step");
                        tokio::time::sleep(Duration::from_millis(100 * retries as u64)).await;
                    } else if matches!(error_mode, ErrorMode::Skip) {
                        warn!(step_id = %step_id, error = ?last_error, "Step failed, skipping");
                        return Ok(StepExecutionResult {
                            step_id: step_id.to_string(),
                            agent_id: None,
                            output: Value::Null,
                            success: false,
                            error: last_error,
                            duration_ms: start_time.elapsed().as_millis() as u64,
                            retries,
                        });
                    } else {
                        return Err(AolError::Execution(format!(
                            "Step {} failed: {}",
                            step_id,
                            last_error.unwrap_or_default()
                        )));
                    }
                }
            }
        }
    }

    /// Evaluate a condition expression.
    fn evaluate_condition(&self, condition: &str, ctx: &TemplateContext) -> AolResult<bool> {
        // First expand any template variables
        let expanded = expand_template(condition, ctx)?;

        // Simple condition evaluation
        // Supports: ==, !=, >, <, >=, <=, contains, starts_with, ends_with
        let expanded = expanded.trim();

        // Check for comparison operators
        if let Some(eq_pos) = expanded.find("==") {
            let left = expanded[..eq_pos].trim();
            let right = expanded[eq_pos + 2..].trim();
            return Ok(left == right);
        }

        if let Some(ne_pos) = expanded.find("!=") {
            let left = expanded[..ne_pos].trim();
            let right = expanded[ne_pos + 2..].trim();
            return Ok(left != right);
        }

        if let Some(gt_pos) = expanded.find('>') {
            let left = expanded[..gt_pos].trim();
            let right = expanded[gt_pos + 1..].trim();
            if let (Ok(l), Ok(r)) = (left.parse::<f64>(), right.parse::<f64>()) {
                return Ok(l > r);
            }
        }

        if let Some(lt_pos) = expanded.find('<') {
            let left = expanded[..lt_pos].trim();
            let right = expanded[lt_pos + 1..].trim();
            if let (Ok(l), Ok(r)) = (left.parse::<f64>(), right.parse::<f64>()) {
                return Ok(l < r);
            }
        }

        if let Some(ge_pos) = expanded.find(">=") {
            let left = expanded[..ge_pos].trim();
            let right = expanded[ge_pos + 2..].trim();
            if let (Ok(l), Ok(r)) = (left.parse::<f64>(), right.parse::<f64>()) {
                return Ok(l >= r);
            }
        }

        if let Some(le_pos) = expanded.find("<=") {
            let left = expanded[..le_pos].trim();
            let right = expanded[le_pos + 2..].trim();
            if let (Ok(l), Ok(r)) = (left.parse::<f64>(), right.parse::<f64>()) {
                return Ok(l <= r);
            }
        }

        // Check for boolean literals
        match expanded.to_lowercase().as_str() {
            "true" | "yes" | "1" => return Ok(true),
            "false" | "no" | "0" => return Ok(false),
            _ => {}
        }

        // Default: non-empty string is truthy
        Ok(!expanded.is_empty())
    }

    /// Evaluate a collection expression.
    fn evaluate_collection(&self, collection: &str, ctx: &TemplateContext) -> AolResult<Vec<Value>> {
        let expanded = expand_template(collection, ctx)?;

        // Try to parse as JSON array
        if let Ok(Value::Array(arr)) = serde_json::from_str::<Value>(&expanded) {
            return Ok(arr);
        }

        // Try to get from context
        if let Some(value) = ctx.get(&expanded) {
            match value {
                Value::Array(arr) => return Ok(arr.clone()),
                Value::String(s) => {
                    // Try parsing string as JSON array
                    if let Ok(Value::Array(arr)) = serde_json::from_str::<Value>(s) {
                        return Ok(arr);
                    }
                    // Split by comma
                    return Ok(s.split(',').map(|s| Value::String(s.trim().to_string())).collect());
                }
                _ => {}
            }
        }

        // Return single-element array
        Ok(vec![Value::String(expanded)])
    }

    /// Get an execution by ID.
    pub async fn get_execution(&self, id: ExecutionId) -> Option<ExecutionResult> {
        self.executions.read().await.get(&id).cloned()
    }

    /// List all executions.
    pub async fn list_executions(&self) -> Vec<ExecutionResult> {
        self.executions.read().await.values().cloned().collect()
    }
}

/// Apply a collect strategy to a list of values.
fn apply_collect_strategy(strategy: &CollectStrategy, values: Vec<Value>) -> Value {
    match strategy {
        CollectStrategy::Merge => Value::Array(values),
        CollectStrategy::First => values.into_iter().next().unwrap_or(Value::Null),
        CollectStrategy::Last => values.into_iter().last().unwrap_or(Value::Null),
        CollectStrategy::Aggregate => {
            // Simple aggregation: concatenate strings, sum numbers
            let mut total = 0.0;
            let mut all_strings = true;

            for v in &values {
                match v {
                    Value::Number(n) => {
                        all_strings = false;
                        if let Some(f) = n.as_f64() {
                            total += f;
                        }
                    }
                    Value::String(_) => {}
                    _ => all_strings = false,
                }
            }

            if all_strings {
                let strings: Vec<&str> = values
                    .iter()
                    .filter_map(|v| v.as_str())
                    .collect();
                Value::String(strings.join("\n"))
            } else {
                Value::Number(serde_json::Number::from_f64(total).unwrap_or_else(|| 0.into()))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use openfang_types::aol::{
        AolWorkflow, ConditionalBranch, ConditionalStep, InputParam, LoopStep,
        ParallelStep, ParallelStepGroup, ParamType, SequentialStep, WorkflowConfig,
    };
    use std::collections::HashMap;

    fn make_simple_workflow() -> AolWorkflow {
        AolWorkflow {
            id: WorkflowDefId::new(),
            name: "test".to_string(),
            version: "1.0.0".to_string(),
            description: String::new(),
            author: String::new(),
            inputs: vec![InputParam::required("input", ParamType::String)],
            outputs: vec!["result".to_string()],
            config: WorkflowConfig::default(),
            steps: vec![AolStep::Sequential(SequentialStep {
                id: "step1".to_string(),
                agent: AgentRef::by_name("test-agent"),
                task: "Process: {{input.input}}".to_string(),
                inputs: HashMap::new(),
                output: Some("result".to_string()),
                error_mode: None,
                timeout_secs: None,
                condition: None,
            })],
            tags: vec![],
        }
    }

    #[test]
    fn test_execution_id_generation() {
        let id1 = ExecutionId::new();
        let id2 = ExecutionId::new();
        assert_ne!(id1, id2);
    }

    #[tokio::test]
    async fn test_execute_simple_workflow() {
        let executor = AolExecutor::with_mock();
        let workflow = make_simple_workflow();
        let compiled = CompiledWorkflow::new(workflow);

        let mut inputs = HashMap::new();
        inputs.insert("input".to_string(), Value::String("test data".to_string()));

        let result = executor.execute(&compiled, inputs).await.unwrap();

        assert_eq!(result.status, ExecutionStatus::Completed);
        assert!(!result.step_results.is_empty());
        assert!(result.outputs.contains_key("result"));
    }

    #[tokio::test]
    async fn test_execute_parallel_workflow() {
        let executor = AolExecutor::with_mock();

        let workflow = AolWorkflow {
            id: WorkflowDefId::new(),
            name: "parallel-test".to_string(),
            version: "1.0.0".to_string(),
            description: String::new(),
            author: String::new(),
            inputs: vec![],
            outputs: vec!["combined".to_string()],
            config: WorkflowConfig::default(),
            steps: vec![AolStep::Parallel(ParallelStepGroup {
                id: "parallel1".to_string(),
                steps: vec![
                    ParallelStep {
                        id: "p1".to_string(),
                        agent: AgentRef::by_name("agent1"),
                        task: "Task 1".to_string(),
                        inputs: HashMap::new(),
                        output: Some("r1".to_string()),
                        error_mode: None,
                        timeout_secs: None,
                    },
                    ParallelStep {
                        id: "p2".to_string(),
                        agent: AgentRef::by_name("agent2"),
                        task: "Task 2".to_string(),
                        inputs: HashMap::new(),
                        output: Some("r2".to_string()),
                        error_mode: None,
                        timeout_secs: None,
                    },
                ],
                collect: CollectStrategy::Merge,
                output: Some("combined".to_string()),
                max_concurrency: Some(2),
            })],
            tags: vec![],
        };

        let compiled = CompiledWorkflow::new(workflow);
        let result = executor.execute(&compiled, HashMap::new()).await.unwrap();

        assert_eq!(result.status, ExecutionStatus::Completed);
        assert_eq!(result.step_results.len(), 2);
    }

    #[tokio::test]
    async fn test_execute_conditional_workflow() {
        let executor = AolExecutor::with_mock();

        let workflow = AolWorkflow {
            id: WorkflowDefId::new(),
            name: "conditional-test".to_string(),
            version: "1.0.0".to_string(),
            description: String::new(),
            author: String::new(),
            inputs: vec![InputParam::required("value", ParamType::Integer)],
            outputs: vec![],
            config: WorkflowConfig::default(),
            steps: vec![AolStep::Conditional(ConditionalStep {
                id: "cond1".to_string(),
                branches: vec![
                    ConditionalBranch {
                        id: "high".to_string(),
                        condition: "{{input.value}} > 10".to_string(),
                        steps: vec![AolStep::Sequential(SequentialStep {
                            id: "high-step".to_string(),
                            agent: AgentRef::by_name("high-agent"),
                            task: "High value".to_string(),
                            inputs: HashMap::new(),
                            output: None,
                            error_mode: None,
                            timeout_secs: None,
                            condition: None,
                        })],
                        output: None,
                    },
                    ConditionalBranch {
                        id: "low".to_string(),
                        condition: "{{input.value}} <= 10".to_string(),
                        steps: vec![AolStep::Sequential(SequentialStep {
                            id: "low-step".to_string(),
                            agent: AgentRef::by_name("low-agent"),
                            task: "Low value".to_string(),
                            inputs: HashMap::new(),
                            output: None,
                            error_mode: None,
                            timeout_secs: None,
                            condition: None,
                        })],
                        output: None,
                    },
                ],
                default: None,
                output: None,
            })],
            tags: vec![],
        };

        let compiled = CompiledWorkflow::new(workflow);

        // Test high value
        let mut inputs = HashMap::new();
        inputs.insert("value".to_string(), Value::Number(15.into()));
        let result = executor.execute(&compiled, inputs).await.unwrap();
        assert_eq!(result.status, ExecutionStatus::Completed);

        // Test low value
        let mut inputs = HashMap::new();
        inputs.insert("value".to_string(), Value::Number(5.into()));
        let result = executor.execute(&compiled, inputs).await.unwrap();
        assert_eq!(result.status, ExecutionStatus::Completed);
    }

    #[tokio::test]
    async fn test_execute_loop_workflow() {
        let executor = AolExecutor::with_mock();

        let workflow = AolWorkflow {
            id: WorkflowDefId::new(),
            name: "loop-test".to_string(),
            version: "1.0.0".to_string(),
            description: String::new(),
            author: String::new(),
            inputs: vec![],
            outputs: vec!["results".to_string()],
            config: WorkflowConfig::default(),
            steps: vec![AolStep::Loop(LoopStep {
                id: "loop1".to_string(),
                item_var: "item".to_string(),
                index_var: Some("idx".to_string()),
                collection: "[\"a\", \"b\", \"c\"]".to_string(),
                steps: vec![AolStep::Sequential(SequentialStep {
                    id: "process".to_string(),
                    agent: AgentRef::by_name("worker"),
                    task: "Process {{loop.item}}".to_string(),
                    inputs: HashMap::new(),
                    output: None,
                    error_mode: None,
                    timeout_secs: None,
                    condition: None,
                })],
                collect: CollectStrategy::Merge,
                output: Some("results".to_string()),
                max_concurrency: 0,
            })],
            tags: vec![],
        };

        let compiled = CompiledWorkflow::new(workflow);
        let result = executor.execute(&compiled, HashMap::new()).await.unwrap();

        assert_eq!(result.status, ExecutionStatus::Completed);
        assert_eq!(result.step_results.len(), 3); // 3 items in loop
    }

    #[test]
    fn test_evaluate_condition_equality() {
        let executor = AolExecutor::with_mock();
        let ctx = TemplateContext::new();

        assert!(executor.evaluate_condition("hello == hello", &ctx).unwrap());
        assert!(!executor.evaluate_condition("hello == world", &ctx).unwrap());
    }

    #[test]
    fn test_evaluate_condition_numeric() {
        let executor = AolExecutor::with_mock();
        let ctx = TemplateContext::new();

        assert!(executor.evaluate_condition("15 > 10", &ctx).unwrap());
        assert!(!executor.evaluate_condition("5 > 10", &ctx).unwrap());
        assert!(executor.evaluate_condition("10 >= 10", &ctx).unwrap());
        assert!(executor.evaluate_condition("5 < 10", &ctx).unwrap());
    }

    #[test]
    fn test_apply_collect_strategy() {
        let values = vec![
            Value::String("a".to_string()),
            Value::String("b".to_string()),
            Value::String("c".to_string()),
        ];

        // Merge
        let result = apply_collect_strategy(&CollectStrategy::Merge, values.clone());
        assert_eq!(result, Value::Array(values.clone()));

        // First
        let result = apply_collect_strategy(&CollectStrategy::First, values.clone());
        assert_eq!(result, Value::String("a".to_string()));

        // Last
        let result = apply_collect_strategy(&CollectStrategy::Last, values.clone());
        assert_eq!(result, Value::String("c".to_string()));
    }
}