//! Agent Orchestration Language (AOL) AST types.
//!
//! This module defines the Abstract Syntax Tree types for AOL, a declarative
//! DSL for defining multi-agent workflows. AOL supports parallel execution,
//! conditional branching, loops, subworkflow calls, and sophisticated error handling.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;

// ---------------------------------------------------------------------------
// WorkflowDefId
// ---------------------------------------------------------------------------

/// Unique identifier for a workflow definition.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct WorkflowDefId(pub Uuid);

impl WorkflowDefId {
    /// Generate a new random WorkflowDefId.
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }
}

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

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

impl std::str::FromStr for WorkflowDefId {
    type Err = uuid::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(Self(Uuid::parse_str(s)?))
    }
}

// ---------------------------------------------------------------------------
// ParamType
// ---------------------------------------------------------------------------

/// Parameter type enumeration for workflow inputs.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ParamType {
    /// String type.
    String,
    /// Integer type (i64).
    Integer,
    /// Float type (f64).
    Float,
    /// Boolean type.
    Boolean,
    /// Array type (JSON array).
    Array,
    /// Object type (JSON object).
    Object,
}

impl Default for ParamType {
    fn default() -> Self {
        Self::String
    }
}

impl std::fmt::Display for ParamType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ParamType::String => write!(f, "string"),
            ParamType::Integer => write!(f, "integer"),
            ParamType::Float => write!(f, "float"),
            ParamType::Boolean => write!(f, "boolean"),
            ParamType::Array => write!(f, "array"),
            ParamType::Object => write!(f, "object"),
        }
    }
}

// ---------------------------------------------------------------------------
// InputParam
// ---------------------------------------------------------------------------

/// Input parameter specification for workflows.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InputParam {
    /// Parameter name.
    pub name: String,
    /// Parameter type.
    #[serde(default)]
    pub param_type: ParamType,
    /// Whether this parameter is required.
    #[serde(default)]
    pub required: bool,
    /// Default value if not provided.
    #[serde(default)]
    pub default: Option<serde_json::Value>,
    /// Human-readable description.
    #[serde(default)]
    pub description: Option<String>,
    /// Allowed enum values (for validation).
    #[serde(default)]
    pub enum_values: Vec<serde_json::Value>,
}

impl InputParam {
    /// Create a new required input parameter.
    pub fn required(name: impl Into<String>, param_type: ParamType) -> Self {
        Self {
            name: name.into(),
            param_type,
            required: true,
            default: None,
            description: None,
            enum_values: Vec::new(),
        }
    }

    /// Create a new optional input parameter with a default value.
    pub fn optional(
        name: impl Into<String>,
        param_type: ParamType,
        default: impl Into<Option<serde_json::Value>>,
    ) -> Self {
        Self {
            name: name.into(),
            param_type,
            required: false,
            default: default.into(),
            description: None,
            enum_values: Vec::new(),
        }
    }

    /// Add a description to the parameter.
    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.description = Some(description.into());
        self
    }

    /// Set enum values for validation.
    pub fn with_enum_values(mut self, values: Vec<serde_json::Value>) -> Self {
        self.enum_values = values;
        self
    }
}

// ---------------------------------------------------------------------------
// ErrorMode
// ---------------------------------------------------------------------------

/// Error handling mode for workflow steps.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ErrorMode {
    /// Fail the entire workflow immediately on error (default).
    #[default]
    Fail,
    /// Skip this step on error and continue.
    Skip,
    /// Retry the step up to max_retries times.
    Retry,
}

// ---------------------------------------------------------------------------
// AgentRef
// ---------------------------------------------------------------------------

/// Reference to an agent in the system.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum AgentRef {
    /// Reference by unique agent ID.
    ById {
        /// The agent's UUID.
        id: Uuid,
    },
    /// Reference by agent name.
    ByName {
        /// The agent's name.
        name: String,
    },
    /// Reference by role/capability.
    ByRole {
        /// The role to match.
        role: String,
        /// Optional specific capability required.
        capability: Option<String>,
    },
}

impl AgentRef {
    /// Create a reference by ID.
    pub fn by_id(id: Uuid) -> Self {
        Self::ById { id }
    }

    /// Create a reference by name.
    pub fn by_name(name: impl Into<String>) -> Self {
        Self::ByName { name: name.into() }
    }

    /// Create a reference by role.
    pub fn by_role(role: impl Into<String>) -> Self {
        Self::ByRole {
            role: role.into(),
            capability: None,
        }
    }

    /// Create a reference by role with required capability.
    pub fn by_role_with_capability(role: impl Into<String>, capability: impl Into<String>) -> Self {
        Self::ByRole {
            role: role.into(),
            capability: Some(capability.into()),
        }
    }
}

// ---------------------------------------------------------------------------
// CollectStrategy
// ---------------------------------------------------------------------------

/// Strategy for collecting results from parallel or loop iterations.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CollectStrategy {
    /// Merge all results into an array (default).
    #[default]
    Merge,
    /// Take only the first completed result.
    First,
    /// Take only the last completed result.
    Last,
    /// Aggregate results using a custom function.
    Aggregate,
}

// ---------------------------------------------------------------------------
// WorkflowConfig
// ---------------------------------------------------------------------------

/// Configuration options for workflow execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct WorkflowConfig {
    /// Maximum execution time in seconds.
    pub timeout_secs: u64,
    /// Maximum retry attempts for failed steps.
    pub max_retries: u32,
    /// Default error handling mode for steps without explicit mode.
    pub default_error_mode: ErrorMode,
    /// Maximum number of concurrent parallel branches.
    pub max_concurrency: usize,
    /// Whether to persist workflow state for resumption.
    pub persist_state: bool,
    /// Custom metadata.
    pub metadata: HashMap<String, serde_json::Value>,
}

impl Default for WorkflowConfig {
    fn default() -> Self {
        Self {
            timeout_secs: 300, // 5 minutes
            max_retries: 3,
            default_error_mode: ErrorMode::Fail,
            max_concurrency: 10,
            persist_state: true,
            metadata: HashMap::new(),
        }
    }
}

// ---------------------------------------------------------------------------
// AolStep - Step types
// ---------------------------------------------------------------------------

/// A single step in a parallel group.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParallelStep {
    /// Unique step identifier within the workflow.
    pub id: String,
    /// Agent to execute this step.
    pub agent: AgentRef,
    /// Task/prompt to send to the agent.
    pub task: String,
    /// Input bindings from workflow context.
    #[serde(default)]
    pub inputs: HashMap<String, String>,
    /// Output variable name to store result.
    #[serde(default)]
    pub output: Option<String>,
    /// Error handling mode for this step.
    #[serde(default)]
    pub error_mode: Option<ErrorMode>,
    /// Timeout in seconds for this step.
    #[serde(default)]
    pub timeout_secs: Option<u64>,
}

/// A group of steps to execute in parallel.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParallelStepGroup {
    /// Unique group identifier.
    pub id: String,
    /// Steps to execute in parallel.
    pub steps: Vec<ParallelStep>,
    /// How to collect results from parallel branches.
    #[serde(default)]
    pub collect: CollectStrategy,
    /// Output variable name for collected results.
    #[serde(default)]
    pub output: Option<String>,
    /// Maximum number of concurrent executions.
    #[serde(default)]
    pub max_concurrency: Option<usize>,
}

/// A step in a sequential execution chain.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SequentialStep {
    /// Unique step identifier.
    pub id: String,
    /// Agent to execute this step.
    pub agent: AgentRef,
    /// Task/prompt to send to the agent.
    pub task: String,
    /// Input bindings from workflow context.
    #[serde(default)]
    pub inputs: HashMap<String, String>,
    /// Output variable name to store result.
    #[serde(default)]
    pub output: Option<String>,
    /// Error handling mode for this step.
    #[serde(default)]
    pub error_mode: Option<ErrorMode>,
    /// Timeout in seconds for this step.
    #[serde(default)]
    pub timeout_secs: Option<u64>,
    /// Condition to check before executing (JavaScript expression).
    #[serde(default)]
    pub condition: Option<String>,
}

/// A conditional branch with a condition and steps.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConditionalBranch {
    /// Branch identifier.
    pub id: String,
    /// Condition expression (JavaScript).
    pub condition: String,
    /// Steps to execute if condition is true.
    pub steps: Vec<AolStep>,
    /// Output variable name.
    #[serde(default)]
    pub output: Option<String>,
}

/// A conditional step with multiple branches.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConditionalStep {
    /// Unique step identifier.
    pub id: String,
    /// Conditional branches (evaluated in order).
    pub branches: Vec<ConditionalBranch>,
    /// Default branch if no conditions match.
    #[serde(default)]
    pub default: Option<Vec<AolStep>>,
    /// Output variable name.
    #[serde(default)]
    pub output: Option<String>,
}

/// A loop step for iterating over collections.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoopStep {
    /// Unique step identifier.
    pub id: String,
    /// Variable name for the current item.
    pub item_var: String,
    /// Variable name for the current index (optional).
    #[serde(default)]
    pub index_var: Option<String>,
    /// Expression to evaluate for the collection to iterate.
    pub collection: String,
    /// Steps to execute for each item.
    pub steps: Vec<AolStep>,
    /// How to collect results from iterations.
    #[serde(default)]
    pub collect: CollectStrategy,
    /// Output variable name for collected results.
    #[serde(default)]
    pub output: Option<String>,
    /// Maximum number of concurrent iterations (0 = sequential).
    #[serde(default)]
    pub max_concurrency: usize,
}

/// A collect step for aggregating results.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CollectStep {
    /// Unique step identifier.
    pub id: String,
    /// Source variable(s) to collect from.
    pub sources: Vec<String>,
    /// Collection strategy.
    #[serde(default)]
    pub strategy: CollectStrategy,
    /// Aggregation function (for Aggregate strategy).
    #[serde(default)]
    pub aggregate_fn: Option<String>,
    /// Output variable name.
    pub output: String,
}

/// A subworkflow invocation step.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubworkflowStep {
    /// Unique step identifier.
    pub id: String,
    /// Reference to the workflow to invoke.
    pub workflow: WorkflowDefId,
    /// Input bindings for the subworkflow.
    #[serde(default)]
    pub inputs: HashMap<String, String>,
    /// Output variable name for subworkflow result.
    #[serde(default)]
    pub output: Option<String>,
    /// Error handling mode.
    #[serde(default)]
    pub error_mode: Option<ErrorMode>,
    /// Timeout in seconds.
    #[serde(default)]
    pub timeout_secs: Option<u64>,
}

/// A fallback chain for error recovery.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FallbackStep {
    /// Unique step identifier.
    pub id: String,
    /// Primary step to attempt first.
    pub primary: Box<AolStep>,
    /// Fallback steps to try if primary fails.
    pub fallbacks: Vec<AolStep>,
    /// Output variable name.
    #[serde(default)]
    pub output: Option<String>,
}

// ---------------------------------------------------------------------------
// AolStep
// ---------------------------------------------------------------------------

/// A workflow step - the core building block of AOL workflows.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AolStep {
    /// Execute a group of steps in parallel.
    Parallel(ParallelStepGroup),
    /// Execute a single step sequentially.
    Sequential(SequentialStep),
    /// Conditional branching.
    Conditional(ConditionalStep),
    /// Loop over a collection.
    Loop(LoopStep),
    /// Collect/aggregate results.
    Collect(CollectStep),
    /// Invoke a subworkflow.
    Subworkflow(SubworkflowStep),
    /// Fallback chain for error recovery.
    Fallback(FallbackStep),
}

impl AolStep {
    /// Get the step ID.
    pub fn id(&self) -> &str {
        match self {
            AolStep::Parallel(p) => &p.id,
            AolStep::Sequential(s) => &s.id,
            AolStep::Conditional(c) => &c.id,
            AolStep::Loop(l) => &l.id,
            AolStep::Collect(c) => &c.id,
            AolStep::Subworkflow(s) => &s.id,
            AolStep::Fallback(f) => &f.id,
        }
    }

    /// Get the output variable name if any.
    pub fn output(&self) -> Option<&str> {
        match self {
            AolStep::Parallel(p) => p.output.as_deref(),
            AolStep::Sequential(s) => s.output.as_deref(),
            AolStep::Conditional(c) => c.output.as_deref(),
            AolStep::Loop(l) => l.output.as_deref(),
            AolStep::Collect(c) => Some(c.output.as_str()),
            AolStep::Subworkflow(s) => s.output.as_deref(),
            AolStep::Fallback(f) => f.output.as_deref(),
        }
    }
}

// ---------------------------------------------------------------------------
// AolWorkflow
// ---------------------------------------------------------------------------

/// A complete AOL workflow definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AolWorkflow {
    /// Unique workflow identifier.
    pub id: WorkflowDefId,
    /// Human-readable workflow name.
    pub name: String,
    /// Semantic version string.
    #[serde(default = "default_version")]
    pub version: String,
    /// Description of what this workflow does.
    #[serde(default)]
    pub description: String,
    /// Author/creator identifier.
    #[serde(default)]
    pub author: String,
    /// Input parameter specifications.
    #[serde(default)]
    pub inputs: Vec<InputParam>,
    /// Output variable names to expose.
    #[serde(default)]
    pub outputs: Vec<String>,
    /// Workflow configuration.
    #[serde(default)]
    pub config: WorkflowConfig,
    /// The root steps of the workflow.
    pub steps: Vec<AolStep>,
    /// Tags for categorization.
    #[serde(default)]
    pub tags: Vec<String>,
}

fn default_version() -> String {
    "0.1.0".to_string()
}

impl AolWorkflow {
    /// Create a new workflow with the given name.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            id: WorkflowDefId::new(),
            name: name.into(),
            version: default_version(),
            description: String::new(),
            author: String::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
            config: WorkflowConfig::default(),
            steps: Vec::new(),
            tags: Vec::new(),
        }
    }

    /// Add an input parameter.
    pub fn with_input(mut self, param: InputParam) -> Self {
        self.inputs.push(param);
        self
    }

    /// Add an output variable.
    pub fn with_output(mut self, output: impl Into<String>) -> Self {
        self.outputs.push(output.into());
        self
    }

    /// Add a step to the workflow.
    pub fn with_step(mut self, step: AolStep) -> Self {
        self.steps.push(step);
        self
    }

    /// Set the workflow description.
    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.description = description.into();
        self
    }

    /// Set the workflow version.
    pub fn with_version(mut self, version: impl Into<String>) -> Self {
        self.version = version.into();
        self
    }

    /// Set the workflow author.
    pub fn with_author(mut self, author: impl Into<String>) -> Self {
        self.author = author.into();
        self
    }

    /// Add a tag.
    pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
        self.tags.push(tag.into());
        self
    }

    /// Set the configuration.
    pub fn with_config(mut self, config: WorkflowConfig) -> Self {
        self.config = config;
        self
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

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

    #[test]
    fn workflow_def_id_display_roundtrip() {
        let id = WorkflowDefId::new();
        let s = id.to_string();
        let parsed: WorkflowDefId = s.parse().unwrap();
        assert_eq!(id, parsed);
    }

    #[test]
    fn workflow_def_id_default() {
        let a = WorkflowDefId::default();
        let b = WorkflowDefId::default();
        assert_ne!(a, b);
    }

    #[test]
    fn param_type_serde() {
        let pt = ParamType::Integer;
        let json = serde_json::to_string(&pt).unwrap();
        assert_eq!(json, "\"integer\"");
        let back: ParamType = serde_json::from_str(&json).unwrap();
        assert_eq!(back, ParamType::Integer);
    }

    #[test]
    fn param_type_default() {
        assert_eq!(ParamType::default(), ParamType::String);
    }

    #[test]
    fn input_param_required() {
        let param = InputParam::required("query", ParamType::String);
        assert_eq!(param.name, "query");
        assert!(param.required);
        assert_eq!(param.param_type, ParamType::String);
    }

    #[test]
    fn input_param_optional_with_default() {
        let param = InputParam::optional(
            "count",
            ParamType::Integer,
            Some(serde_json::json!(10)),
        );
        assert_eq!(param.name, "count");
        assert!(!param.required);
        assert_eq!(param.default, Some(serde_json::json!(10)));
    }

    #[test]
    fn input_param_with_description() {
        let param = InputParam::required("query", ParamType::String)
            .with_description("Search query string");
        assert_eq!(param.description, Some("Search query string".to_string()));
    }

    #[test]
    fn input_param_with_enum_values() {
        let param = InputParam::required("mode", ParamType::String)
            .with_enum_values(vec![
                serde_json::json!("fast"),
                serde_json::json!("thorough"),
            ]);
        assert_eq!(param.enum_values.len(), 2);
    }

    #[test]
    fn input_param_serde_roundtrip() {
        let param = InputParam {
            name: "test".to_string(),
            param_type: ParamType::Boolean,
            required: true,
            default: Some(serde_json::json!(false)),
            description: Some("A test param".to_string()),
            enum_values: vec![serde_json::json!(true), serde_json::json!(false)],
        };
        let json = serde_json::to_string(&param).unwrap();
        let back: InputParam = serde_json::from_str(&json).unwrap();
        assert_eq!(back.name, param.name);
        assert_eq!(back.param_type, param.param_type);
        assert_eq!(back.required, param.required);
    }

    #[test]
    fn error_mode_default() {
        assert_eq!(ErrorMode::default(), ErrorMode::Fail);
    }

    #[test]
    fn error_mode_serde() {
        let mode = ErrorMode::Retry;
        let json = serde_json::to_string(&mode).unwrap();
        assert_eq!(json, "\"retry\"");
        let back: ErrorMode = serde_json::from_str(&json).unwrap();
        assert_eq!(back, ErrorMode::Retry);
    }

    #[test]
    fn agent_ref_by_id() {
        let id = Uuid::new_v4();
        let r = AgentRef::by_id(id);
        match r {
            AgentRef::ById { id: i } => assert_eq!(i, id),
            _ => panic!("Expected ById"),
        }
    }

    #[test]
    fn agent_ref_by_name() {
        let r = AgentRef::by_name("assistant");
        match r {
            AgentRef::ByName { name } => assert_eq!(name, "assistant"),
            _ => panic!("Expected ByName"),
        }
    }

    #[test]
    fn agent_ref_by_role() {
        let r = AgentRef::by_role("researcher");
        match r {
            AgentRef::ByRole { role, capability } => {
                assert_eq!(role, "researcher");
                assert!(capability.is_none());
            }
            _ => panic!("Expected ByRole"),
        }
    }

    #[test]
    fn agent_ref_by_role_with_capability() {
        let r = AgentRef::by_role_with_capability("worker", "web_search");
        match r {
            AgentRef::ByRole { role, capability } => {
                assert_eq!(role, "worker");
                assert_eq!(capability, Some("web_search".to_string()));
            }
            _ => panic!("Expected ByRole"),
        }
    }

    #[test]
    fn agent_ref_serde_roundtrip() {
        let r = AgentRef::by_name("test-agent");
        let json = serde_json::to_string(&r).unwrap();
        let back: AgentRef = serde_json::from_str(&json).unwrap();
        assert_eq!(r, back);
    }

    #[test]
    fn collect_strategy_default() {
        assert_eq!(CollectStrategy::default(), CollectStrategy::Merge);
    }

    #[test]
    fn collect_strategy_serde() {
        let s = CollectStrategy::First;
        let json = serde_json::to_string(&s).unwrap();
        assert_eq!(json, "\"first\"");
        let back: CollectStrategy = serde_json::from_str(&json).unwrap();
        assert_eq!(back, CollectStrategy::First);
    }

    #[test]
    fn workflow_config_default() {
        let config = WorkflowConfig::default();
        assert_eq!(config.timeout_secs, 300);
        assert_eq!(config.max_retries, 3);
        assert_eq!(config.default_error_mode, ErrorMode::Fail);
        assert_eq!(config.max_concurrency, 10);
        assert!(config.persist_state);
    }

    #[test]
    fn workflow_config_serde_roundtrip() {
        let config = WorkflowConfig {
            timeout_secs: 600,
            max_retries: 5,
            default_error_mode: ErrorMode::Skip,
            max_concurrency: 20,
            persist_state: false,
            metadata: vec![("key".to_string(), serde_json::json!("value"))]
                .into_iter()
                .collect(),
        };
        let json = serde_json::to_string(&config).unwrap();
        let back: WorkflowConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(back.timeout_secs, 600);
        assert_eq!(back.max_retries, 5);
        assert_eq!(back.default_error_mode, ErrorMode::Skip);
    }

    #[test]
    fn parallel_step_creation() {
        let step = ParallelStep {
            id: "step1".to_string(),
            agent: AgentRef::by_name("worker"),
            task: "Process data".to_string(),
            inputs: vec![("input".to_string(), "$.data".to_string())]
                .into_iter()
                .collect(),
            output: Some("result".to_string()),
            error_mode: Some(ErrorMode::Retry),
            timeout_secs: Some(60),
        };
        assert_eq!(step.id, "step1");
        assert!(step.inputs.contains_key("input"));
    }

    #[test]
    fn parallel_step_group_serde() {
        let group = ParallelStepGroup {
            id: "parallel1".to_string(),
            steps: vec![
                ParallelStep {
                    id: "p1".to_string(),
                    agent: AgentRef::by_name("a1"),
                    task: "Task 1".to_string(),
                    inputs: HashMap::new(),
                    output: None,
                    error_mode: None,
                    timeout_secs: None,
                },
            ],
            collect: CollectStrategy::Merge,
            output: Some("combined".to_string()),
            max_concurrency: Some(5),
        };
        let json = serde_json::to_string(&group).unwrap();
        let back: ParallelStepGroup = serde_json::from_str(&json).unwrap();
        assert_eq!(back.id, "parallel1");
        assert_eq!(back.steps.len(), 1);
    }

    #[test]
    fn sequential_step_creation() {
        let step = SequentialStep {
            id: "seq1".to_string(),
            agent: AgentRef::by_name("assistant"),
            task: "Analyze input".to_string(),
            inputs: HashMap::new(),
            output: Some("analysis".to_string()),
            error_mode: None,
            timeout_secs: None,
            condition: Some("$.input != null".to_string()),
        };
        assert_eq!(step.id, "seq1");
        assert!(step.condition.is_some());
    }

    #[test]
    fn conditional_step_creation() {
        let step = ConditionalStep {
            id: "cond1".to_string(),
            branches: vec![
                ConditionalBranch {
                    id: "branch1".to_string(),
                    condition: "$.value > 10".to_string(),
                    steps: vec![],
                    output: None,
                },
            ],
            default: Some(vec![]),
            output: Some("result".to_string()),
        };
        assert_eq!(step.branches.len(), 1);
        assert!(step.default.is_some());
    }

    #[test]
    fn loop_step_creation() {
        let step = LoopStep {
            id: "loop1".to_string(),
            item_var: "item".to_string(),
            index_var: Some("idx".to_string()),
            collection: "$.items".to_string(),
            steps: vec![],
            collect: CollectStrategy::Merge,
            output: Some("results".to_string()),
            max_concurrency: 0,
        };
        assert_eq!(step.item_var, "item");
        assert_eq!(step.index_var, Some("idx".to_string()));
    }

    #[test]
    fn collect_step_creation() {
        let step = CollectStep {
            id: "collect1".to_string(),
            sources: vec!["r1".to_string(), "r2".to_string()],
            strategy: CollectStrategy::Aggregate,
            aggregate_fn: Some("sum".to_string()),
            output: "total".to_string(),
        };
        assert_eq!(step.sources.len(), 2);
        assert_eq!(step.strategy, CollectStrategy::Aggregate);
    }

    #[test]
    fn subworkflow_step_creation() {
        let wf_id = WorkflowDefId::new();
        let step = SubworkflowStep {
            id: "sub1".to_string(),
            workflow: wf_id,
            inputs: vec![("param".to_string(), "$.value".to_string())]
                .into_iter()
                .collect(),
            output: Some("sub_result".to_string()),
            error_mode: Some(ErrorMode::Fail),
            timeout_secs: Some(120),
        };
        assert_eq!(step.workflow, wf_id);
    }

    #[test]
    fn fallback_step_creation() {
        let step = FallbackStep {
            id: "fallback1".to_string(),
            primary: Box::new(AolStep::Sequential(SequentialStep {
                id: "primary".to_string(),
                agent: AgentRef::by_name("primary"),
                task: "Primary task".to_string(),
                inputs: HashMap::new(),
                output: None,
                error_mode: None,
                timeout_secs: None,
                condition: None,
            })),
            fallbacks: vec![],
            output: Some("result".to_string()),
        };
        assert_eq!(step.id, "fallback1");
    }

    #[test]
    fn aol_step_id() {
        let step = AolStep::Sequential(SequentialStep {
            id: "test_id".to_string(),
            agent: AgentRef::by_name("agent"),
            task: "Task".to_string(),
            inputs: HashMap::new(),
            output: None,
            error_mode: None,
            timeout_secs: None,
            condition: None,
        });
        assert_eq!(step.id(), "test_id");
    }

    #[test]
    fn aol_step_output() {
        let step = AolStep::Sequential(SequentialStep {
            id: "test".to_string(),
            agent: AgentRef::by_name("agent"),
            task: "Task".to_string(),
            inputs: HashMap::new(),
            output: Some("my_output".to_string()),
            error_mode: None,
            timeout_secs: None,
            condition: None,
        });
        assert_eq!(step.output(), Some("my_output"));
    }

    #[test]
    fn aol_step_serde_roundtrip() {
        let step = AolStep::Parallel(ParallelStepGroup {
            id: "p1".to_string(),
            steps: vec![],
            collect: CollectStrategy::Merge,
            output: None,
            max_concurrency: None,
        });
        let json = serde_json::to_string(&step).unwrap();
        assert!(json.contains("\"type\":\"parallel\""));
        let back: AolStep = serde_json::from_str(&json).unwrap();
        assert_eq!(back.id(), "p1");
    }

    #[test]
    fn aol_workflow_new() {
        let wf = AolWorkflow::new("test-workflow");
        assert_eq!(wf.name, "test-workflow");
        assert_eq!(wf.version, "0.1.0");
        assert!(wf.inputs.is_empty());
        assert!(wf.steps.is_empty());
    }

    #[test]
    fn aol_workflow_builder() {
        let wf = AolWorkflow::new("test-workflow")
            .with_description("A test workflow")
            .with_version("1.0.0")
            .with_author("test-author")
            .with_tag("test")
            .with_tag("example")
            .with_input(InputParam::required("query", ParamType::String))
            .with_output("result");

        assert_eq!(wf.description, "A test workflow");
        assert_eq!(wf.version, "1.0.0");
        assert_eq!(wf.author, "test-author");
        assert_eq!(wf.tags.len(), 2);
        assert_eq!(wf.inputs.len(), 1);
        assert_eq!(wf.outputs.len(), 1);
    }

    #[test]
    fn aol_workflow_serde_roundtrip() {
        let wf = AolWorkflow::new("test-workflow")
            .with_description("Test description")
            .with_input(InputParam::required("input", ParamType::String))
            .with_step(AolStep::Sequential(SequentialStep {
                id: "step1".to_string(),
                agent: AgentRef::by_name("agent"),
                task: "Do something".to_string(),
                inputs: HashMap::new(),
                output: Some("result".to_string()),
                error_mode: None,
                timeout_secs: None,
                condition: None,
            }));

        let json = serde_json::to_string(&wf).unwrap();
        let back: AolWorkflow = serde_json::from_str(&json).unwrap();

        assert_eq!(back.name, wf.name);
        assert_eq!(back.description, wf.description);
        assert_eq!(back.inputs.len(), 1);
        assert_eq!(back.steps.len(), 1);
    }

    #[test]
    fn aol_workflow_with_config() {
        let config = WorkflowConfig {
            timeout_secs: 600,
            max_retries: 5,
            ..Default::default()
        };
        let wf = AolWorkflow::new("test").with_config(config);
        assert_eq!(wf.config.timeout_secs, 600);
        assert_eq!(wf.config.max_retries, 5);
    }
}
