enact_core/flow/

repeat.rs

//! Loop Flow - Iteration with exit condition
//!
//! Repeatedly execute a callable until an exit condition is met.

use crate::callable::Callable;
use std::sync::Arc;

/// Loop exit condition
pub enum LoopCondition {
    /// Fixed number of iterations
    MaxIterations(usize),
    /// Output matches predicate
    OutputMatches(Box<dyn Fn(&str) -> bool + Send + Sync>),
    /// Output contains string
    OutputContains(String),
    /// Combined: max iterations OR output matches
    Either {
        max_iterations: usize,
        predicate: Box<dyn Fn(&str) -> bool + Send + Sync>,
    },
}

impl LoopCondition {
    /// Check if loop should exit
    pub fn should_exit(&self, iteration: usize, output: &str) -> bool {
        match self {
            LoopCondition::MaxIterations(max) => iteration >= *max,
            LoopCondition::OutputMatches(pred) => pred(output),
            LoopCondition::OutputContains(needle) => output.contains(needle),
            LoopCondition::Either {
                max_iterations,
                predicate,
            } => iteration >= *max_iterations || predicate(output),
        }
    }

    /// Create a max iterations condition
    pub fn max(n: usize) -> Self {
        LoopCondition::MaxIterations(n)
    }

    /// Create an output contains condition
    pub fn until_contains(s: impl Into<String>) -> Self {
        LoopCondition::OutputContains(s.into())
    }

    /// Create a predicate condition
    pub fn until(pred: impl Fn(&str) -> bool + Send + Sync + 'static) -> Self {
        LoopCondition::OutputMatches(Box::new(pred))
    }

    /// Create a combined condition
    pub fn max_or_until(max: usize, pred: impl Fn(&str) -> bool + Send + Sync + 'static) -> Self {
        LoopCondition::Either {
            max_iterations: max,
            predicate: Box::new(pred),
        }
    }
}

/// Loop execution flow
pub struct LoopFlow<C: Callable> {
    /// The callable to execute repeatedly
    callable: Arc<C>,
    /// Exit condition
    condition: LoopCondition,
    /// Flow name
    name: String,
    /// Whether to pass previous output as input (feedback loop)
    feedback: bool,
}

impl<C: Callable> LoopFlow<C> {
    /// Create a new loop flow
    pub fn new(name: impl Into<String>, callable: Arc<C>, condition: LoopCondition) -> Self {
        Self {
            callable,
            condition,
            name: name.into(),
            feedback: true, // Default: use output as next input
        }
    }

    /// Create a loop that runs N times
    pub fn times(name: impl Into<String>, n: usize, callable: Arc<C>) -> Self {
        Self::new(name, callable, LoopCondition::MaxIterations(n))
    }

    /// Create a loop that runs until output contains string
    pub fn until_contains(name: impl Into<String>, s: impl Into<String>, callable: Arc<C>) -> Self {
        Self::new(name, callable, LoopCondition::OutputContains(s.into()))
    }

    /// Set whether to use output as next input
    pub fn with_feedback(mut self, feedback: bool) -> Self {
        self.feedback = feedback;
        self
    }

    /// Execute the loop
    pub async fn execute(&self, input: &str) -> anyhow::Result<String> {
        let mut current_input = input.to_string();
        let mut iteration = 0;

        loop {
            let output = self.callable.run(&current_input).await?;

            if self.condition.should_exit(iteration, &output) {
                return Ok(output);
            }

            // Prepare next iteration
            if self.feedback {
                current_input = output;
            }
            iteration += 1;
        }
    }

    /// Execute with iteration tracking
    pub async fn execute_with_history(&self, input: &str) -> anyhow::Result<LoopHistory> {
        let mut current_input = input.to_string();
        let mut iteration = 0;
        let mut outputs = Vec::new();

        loop {
            let output = self.callable.run(&current_input).await?;
            outputs.push(output.clone());

            if self.condition.should_exit(iteration, &output) {
                return Ok(LoopHistory {
                    iterations: iteration + 1,
                    outputs,
                    final_output: output,
                });
            }

            if self.feedback {
                current_input = output;
            }
            iteration += 1;
        }
    }

    /// Get the flow name
    pub fn name(&self) -> &str {
        &self.name
    }
}

/// History of a loop execution
#[derive(Debug)]
pub struct LoopHistory {
    /// Number of iterations executed
    pub iterations: usize,
    /// Output from each iteration
    pub outputs: Vec<String>,
    /// Final output
    pub final_output: String,
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// Mock callable that tracks calls and transforms input
    #[allow(clippy::type_complexity)]
    struct MockCallable {
        name: String,
        call_count: Arc<AtomicUsize>,
        transform: Box<dyn Fn(&str, usize) -> String + Send + Sync>,
    }

    impl MockCallable {
        fn new(
            name: &str,
            transform: impl Fn(&str, usize) -> String + Send + Sync + 'static,
        ) -> Self {
            Self {
                name: name.to_string(),
                call_count: Arc::new(AtomicUsize::new(0)),
                transform: Box::new(transform),
            }
        }

        /// Simple incrementing callable - appends iteration count
        fn incrementing(name: &str) -> Self {
            Self::new(name, |input, n| format!("{}:{}", input, n))
        }

        /// Callable that emits "DONE" on the Nth call
        fn done_on_call(name: &str, n: usize) -> Self {
            Self::new(name, move |input, call| {
                if call >= n - 1 {
                    "DONE".to_string()
                } else {
                    format!("{}:{}", input, call)
                }
            })
        }

        fn get_call_count(&self) -> usize {
            self.call_count.load(Ordering::SeqCst)
        }
    }

    #[async_trait]
    impl Callable for MockCallable {
        fn name(&self) -> &str {
            &self.name
        }

        async fn run(&self, input: &str) -> anyhow::Result<String> {
            let n = self.call_count.fetch_add(1, Ordering::SeqCst);
            Ok((self.transform)(input, n))
        }
    }

    // ============ LoopCondition Tests ============

    #[test]
    fn test_condition_max_iterations() {
        let cond = LoopCondition::MaxIterations(3);
        assert!(!cond.should_exit(0, "any"));
        assert!(!cond.should_exit(1, "any"));
        assert!(!cond.should_exit(2, "any"));
        assert!(cond.should_exit(3, "any")); // Exit at iteration 3
        assert!(cond.should_exit(5, "any")); // Also exit if past
    }

    #[test]
    fn test_condition_output_matches() {
        let cond = LoopCondition::OutputMatches(Box::new(|s| s.len() > 5));
        assert!(!cond.should_exit(0, "hi"));
        assert!(!cond.should_exit(10, "short"));
        assert!(cond.should_exit(0, "longer"));
        assert!(cond.should_exit(0, "this is long enough"));
    }

    #[test]
    fn test_condition_output_contains() {
        let cond = LoopCondition::OutputContains("DONE".to_string());
        assert!(!cond.should_exit(0, "not yet"));
        assert!(!cond.should_exit(5, "still working"));
        assert!(cond.should_exit(0, "DONE"));
        assert!(cond.should_exit(0, "task DONE here"));
    }

    #[test]
    fn test_condition_either() {
        let cond = LoopCondition::Either {
            max_iterations: 5,
            predicate: Box::new(|s| s.contains("STOP")),
        };

        // Not exit before max or predicate
        assert!(!cond.should_exit(0, "working"));
        assert!(!cond.should_exit(4, "still going"));

        // Exit at max iterations
        assert!(cond.should_exit(5, "working"));

        // Exit when predicate matches, even before max
        assert!(cond.should_exit(2, "STOP now"));
    }

    #[test]
    fn test_condition_helpers() {
        // max()
        let cond = LoopCondition::max(2);
        assert!(!cond.should_exit(1, "x"));
        assert!(cond.should_exit(2, "x"));

        // until_contains()
        let cond = LoopCondition::until_contains("END");
        assert!(!cond.should_exit(0, "not"));
        assert!(cond.should_exit(0, "END"));

        // until()
        let cond = LoopCondition::until(|s| s == "target");
        assert!(!cond.should_exit(0, "other"));
        assert!(cond.should_exit(0, "target"));

        // max_or_until()
        let cond = LoopCondition::max_or_until(3, |s| s.starts_with("!"));
        assert!(!cond.should_exit(0, "a"));
        assert!(cond.should_exit(3, "a")); // max hit
        assert!(cond.should_exit(0, "!bang")); // predicate hit
    }

    // ============ LoopFlow Construction Tests ============

    #[tokio::test]
    async fn test_loop_flow_new() {
        let callable = Arc::new(MockCallable::incrementing("inc"));
        let flow = LoopFlow::new("test_loop", callable, LoopCondition::MaxIterations(2));
        assert_eq!(flow.name(), "test_loop");
    }

    #[tokio::test]
    async fn test_loop_flow_times() {
        let callable = Arc::new(MockCallable::incrementing("inc"));
        let flow = LoopFlow::times("timer", 3, callable);
        assert_eq!(flow.name(), "timer");
    }

    #[tokio::test]
    async fn test_loop_flow_until_contains() {
        let callable = Arc::new(MockCallable::done_on_call("done", 2));
        let flow = LoopFlow::until_contains("stopper", "DONE", callable);
        assert_eq!(flow.name(), "stopper");
    }

    // ============ LoopFlow Execute Tests ============

    #[tokio::test]
    async fn test_loop_execute_max_iterations() {
        let callable = Arc::new(MockCallable::incrementing("inc"));
        let flow = LoopFlow::times("loop", 3, callable.clone());

        let result = flow.execute("start").await.unwrap();
        // With MaxIterations(3), exit condition is iteration >= 3
        // iteration 0: run, check 0>=3=false, inc to 1
        // iteration 1: run, check 1>=3=false, inc to 2
        // iteration 2: run, check 2>=3=false, inc to 3
        // iteration 3: run, check 3>=3=true, exit
        // So 4 total calls
        assert_eq!(callable.get_call_count(), 4);
        assert!(result.contains("start"));
    }

    #[tokio::test]
    async fn test_loop_execute_until_contains() {
        let callable = Arc::new(MockCallable::done_on_call("done", 3));
        let flow = LoopFlow::until_contains("wait_done", "DONE", callable.clone());

        let result = flow.execute("input").await.unwrap();
        assert_eq!(result, "DONE");
        assert_eq!(callable.get_call_count(), 3);
    }

    #[tokio::test]
    async fn test_loop_execute_with_predicate() {
        let callable = Arc::new(MockCallable::new("counter", |_, n| format!("count:{}", n)));
        let flow = LoopFlow::new(
            "until_five",
            callable.clone(),
            LoopCondition::until(|s| s == "count:5"),
        );

        let result = flow.execute("x").await.unwrap();
        assert_eq!(result, "count:5");
        assert_eq!(callable.get_call_count(), 6); // 0,1,2,3,4,5
    }

    #[tokio::test]
    async fn test_loop_execute_either_max_first() {
        let callable = Arc::new(MockCallable::new("counter", |_, n| format!("v{}", n)));
        let flow = LoopFlow::new(
            "either",
            callable.clone(),
            LoopCondition::max_or_until(3, |s| s == "never"),
        );

        let result = flow.execute("x").await.unwrap();
        // Exit at iteration >= 3, so 4 calls total
        assert_eq!(callable.get_call_count(), 4);
        assert_eq!(result, "v3");
    }

    #[tokio::test]
    async fn test_loop_execute_either_predicate_first() {
        let callable = Arc::new(MockCallable::done_on_call("done", 2));
        let flow = LoopFlow::new(
            "either",
            callable.clone(),
            LoopCondition::max_or_until(10, |s| s == "DONE"),
        );

        let result = flow.execute("x").await.unwrap();
        assert_eq!(result, "DONE");
        assert_eq!(callable.get_call_count(), 2); // Exit before max
    }

    // ============ Feedback Mode Tests ============

    #[tokio::test]
    async fn test_loop_with_feedback_enabled() {
        // Track inputs to verify feedback
        let inputs: Arc<std::sync::Mutex<Vec<String>>> =
            Arc::new(std::sync::Mutex::new(Vec::new()));
        let inputs_clone = inputs.clone();

        let callable = Arc::new(MockCallable::new("fb", move |input, n| {
            inputs_clone.lock().unwrap().push(input.to_string());
            format!("out{}", n)
        }));

        let flow = LoopFlow::times("feedback_on", 3, callable).with_feedback(true);
        flow.execute("start").await.unwrap();

        let recorded = inputs.lock().unwrap().clone();
        // MaxIterations(3) runs 4 times (exits when iteration >= 3)
        assert_eq!(recorded, vec!["start", "out0", "out1", "out2"]);
    }

    #[tokio::test]
    async fn test_loop_with_feedback_disabled() {
        let inputs: Arc<std::sync::Mutex<Vec<String>>> =
            Arc::new(std::sync::Mutex::new(Vec::new()));
        let inputs_clone = inputs.clone();

        let callable = Arc::new(MockCallable::new("no_fb", move |input, n| {
            inputs_clone.lock().unwrap().push(input.to_string());
            format!("out{}", n)
        }));

        let flow = LoopFlow::times("feedback_off", 3, callable).with_feedback(false);
        flow.execute("same").await.unwrap();

        let recorded = inputs.lock().unwrap().clone();
        // Without feedback, input stays the same; MaxIterations(3) runs 4 times
        assert_eq!(recorded, vec!["same", "same", "same", "same"]);
    }

    // ============ Execute with History Tests ============

    #[tokio::test]
    async fn test_loop_execute_with_history() {
        let callable = Arc::new(MockCallable::new("hist", |_, n| format!("iter{}", n)));
        let flow = LoopFlow::times("history_test", 4, callable);

        let history = flow.execute_with_history("start").await.unwrap();

        // MaxIterations(4) exits when iteration >= 4, so runs 5 times
        assert_eq!(history.iterations, 5);
        assert_eq!(history.outputs.len(), 5);
        assert_eq!(
            history.outputs,
            vec!["iter0", "iter1", "iter2", "iter3", "iter4"]
        );
        assert_eq!(history.final_output, "iter4");
    }

    #[tokio::test]
    async fn test_loop_execute_with_history_early_exit() {
        let callable = Arc::new(MockCallable::done_on_call("early", 2));
        let flow = LoopFlow::until_contains("early_exit", "DONE", callable);

        let history = flow.execute_with_history("x").await.unwrap();

        assert_eq!(history.iterations, 2);
        assert_eq!(history.outputs.len(), 2);
        assert_eq!(history.final_output, "DONE");
    }

    // ============ Error Handling Tests ============

    #[tokio::test]
    async fn test_loop_error_propagation() {
        struct FailingCallable {
            fail_on: usize,
            call_count: Arc<AtomicUsize>,
        }

        #[async_trait]
        impl Callable for FailingCallable {
            fn name(&self) -> &str {
                "failing"
            }

            async fn run(&self, _input: &str) -> anyhow::Result<String> {
                let n = self.call_count.fetch_add(1, Ordering::SeqCst);
                if n >= self.fail_on {
                    anyhow::bail!("Intentional failure at iteration {}", n)
                }
                Ok(format!("ok{}", n))
            }
        }

        let callable = Arc::new(FailingCallable {
            fail_on: 2,
            call_count: Arc::new(AtomicUsize::new(0)),
        });

        let flow = LoopFlow::times("fail_loop", 5, callable);
        let result = flow.execute("start").await;

        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Intentional failure"));
    }

    // ============ Edge Cases ============

    #[tokio::test]
    async fn test_loop_zero_iterations() {
        let callable = Arc::new(MockCallable::incrementing("zero"));
        let flow = LoopFlow::times("zero_loop", 0, callable.clone());

        // With max 0, should_exit returns true on iteration 0
        // So we run once and immediately exit
        let result = flow.execute("input").await.unwrap();
        assert_eq!(callable.get_call_count(), 1);
        assert!(result.contains("input"));
    }

    #[tokio::test]
    async fn test_loop_immediate_exit_predicate() {
        let callable = Arc::new(MockCallable::new("imm", |_, _| "STOP".to_string()));
        let flow = LoopFlow::new(
            "immediate",
            callable.clone(),
            LoopCondition::until_contains("STOP"),
        );

        let result = flow.execute("x").await.unwrap();
        assert_eq!(result, "STOP");
        assert_eq!(callable.get_call_count(), 1);
    }

    #[tokio::test]
    async fn test_loop_single_iteration() {
        let callable = Arc::new(MockCallable::incrementing("single"));
        let flow = LoopFlow::times("one", 1, callable.clone());

        let history = flow.execute_with_history("x").await.unwrap();
        // MaxIterations(1) exits when iteration >= 1, so runs 2 times
        assert_eq!(history.iterations, 2);
        assert_eq!(callable.get_call_count(), 2);
    }
}