bonsai_mdsl/

tree.rs

use std::collections::HashMap;
use std::time::Instant;

use crate::context::TreeContext;
use crate::error::{BonsaiError, Result};
use crate::nodes::{Guard, Node, NodeResult, NodeState, NodeType};
use crate::parser::parse_mdsl;

/// A behavior tree that can be executed with a context
pub struct BehaviorTree {
    root: Node,
    _named_trees: HashMap<String, Node>,
    start_time: Option<Instant>,
}

impl BehaviorTree {
    /// Create a new behavior tree from an MDSL string
    pub fn from_mdsl(mdsl: &str) -> Result<Self> {
        let root = parse_mdsl(mdsl)?;
        Ok(Self {
            root,
            _named_trees: HashMap::new(),
            start_time: None,
        })
    }

    /// Create a behavior tree from a pre-built node
    pub fn from_node(root: Node) -> Self {
        Self {
            root,
            _named_trees: HashMap::new(),
            start_time: None,
        }
    }

    /// Execute one tick of the behavior tree
    pub fn tick(&mut self, context: &TreeContext) -> Result<NodeResult> {
        self.tick_with_delta(context, 0.0)
    }

    /// Execute one tick of the behavior tree with delta time
    pub fn tick_with_delta(
        &mut self,
        context: &TreeContext,
        delta_time: f64,
    ) -> Result<NodeResult> {
        if self.start_time.is_none() {
            self.start_time = Some(Instant::now());
        }

        let result = execute_node(&mut self.root, context, delta_time)?;

        // Don't auto-reset - let the user decide when to reset
        // if matches!(result, NodeResult::Success | NodeResult::Failure) {
        //     self.reset();
        // }

        Ok(result)
    }

    /// Reset the tree to its initial state
    pub fn reset(&mut self) {
        self.root.reset();
        self.start_time = None;
    }

    /// Get the current state of the tree
    pub fn get_state(&self) -> NodeState {
        self.root.state
    }

    /// Check if the tree is currently running
    pub fn is_running(&self) -> bool {
        matches!(self.root.state, NodeState::Running)
    }
}

/// Execute a single node - made into a free function to avoid borrowing issues
fn execute_node(node: &mut Node, context: &TreeContext, delta_time: f64) -> Result<NodeResult> {
    // Check guards first
    if let Some(guard) = &node.guard {
        if !evaluate_guard(guard, context)? {
            return Ok(NodeResult::Failure);
        }
    }

    // Execute entry callback if transitioning to running
    if node.state == NodeState::Ready {
        if let Some(callbacks) = &node.callbacks {
            if let Some((callback_name, args)) = &callbacks.entry {
                context.execute_callback(callback_name, args)?;
            }
        }
    }

    // Execute step callback
    if let Some(callbacks) = &node.callbacks {
        if let Some((callback_name, args)) = &callbacks.step {
            context.execute_callback(callback_name, args)?;
        }
    }

    // Execute the node based on its type
    let result = match &mut node.node_type {
        NodeType::Root { child } => execute_node(child, context, delta_time)?,
        NodeType::Sequence { children } => execute_sequence(children, context, delta_time)?,
        NodeType::Selector { children } => execute_selector(children, context, delta_time)?,
        NodeType::Parallel { children } => execute_parallel(children, context, delta_time)?,
        NodeType::Race { children } => execute_race(children, context, delta_time)?,
        NodeType::All { children } => execute_all(children, context, delta_time)?,
        NodeType::Lotto { children, weights } => {
            execute_lotto(children, weights.as_ref(), context, delta_time)?
        }
        NodeType::While {
            condition,
            args,
            children,
        } => execute_while(condition, args, children, context, delta_time)?,
        NodeType::Until {
            condition,
            args,
            children,
        } => execute_until(condition, args, children, context, delta_time)?,
        NodeType::WhileAll {
            condition,
            args,
            children,
        } => execute_while_all(condition, args, children, context, delta_time)?,
        NodeType::Repeat { child, iterations } => {
            execute_repeat(child, *iterations, context, delta_time)?
        }
        NodeType::Retry { child, attempts } => {
            execute_retry(child, *attempts, context, delta_time)?
        }
        NodeType::Flip { child } => execute_flip(child, context, delta_time)?,
        NodeType::Succeed { child } => execute_succeed(child, context, delta_time)?,
        NodeType::Fail { child } => execute_fail(child, context, delta_time)?,
        NodeType::Action { name, args } => context.execute_action(name, args)?,
        NodeType::Condition { name, args } => {
            if context.evaluate_condition(name, args)? {
                NodeResult::Success
            } else {
                NodeResult::Failure
            }
        }
        NodeType::Wait { duration } => {
            if let Some(duration_ms) = duration {
                // Add delta time to elapsed time
                node.elapsed_time += delta_time * 1000.0; // Convert seconds to milliseconds

                if node.elapsed_time >= *duration_ms as f64 {
                    // Wait duration completed, reset for next time
                    node.elapsed_time = 0.0;
                    NodeResult::Success
                } else {
                    // Still waiting
                    NodeResult::Running
                }
            } else {
                // Wait indefinitely
                NodeResult::Running
            }
        }
        NodeType::Branch {
            reference: _reference,
        } => {
            // Branch nodes reference other named behavior trees
            // This feature requires implementing a tree registry and would be a significant architectural change
            return Err(BonsaiError::NodeExecutionError(
                "Branch nodes are not yet implemented - this is a planned feature for referencing named trees".to_string(),
            ));
        }
    };

    // Update node state
    node.state = result.into();

    // Execute exit callback if transitioning to a terminal state
    if matches!(result, NodeResult::Success | NodeResult::Failure) {
        if let Some(callbacks) = &node.callbacks {
            if let Some((callback_name, args)) = &callbacks.exit {
                context.execute_callback(callback_name, args)?;
            }
        }
    }

    Ok(result)
}

fn execute_sequence(
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    let mut processed_wait_in_this_tick = false;

    for child in children {
        // Skip children that have already succeeded
        if child.state == NodeState::Success {
            continue;
        }

        let is_wait_node = matches!(child.node_type, NodeType::Wait { .. });

        // If we already processed a wait node and this is another wait node, stop
        if processed_wait_in_this_tick && is_wait_node {
            return Ok(NodeResult::Running);
        }

        // Execute this child
        let result = execute_node(child, context, delta_time)?;

        // Track if we processed a wait node
        if is_wait_node && matches!(result, NodeResult::Success | NodeResult::Running) {
            processed_wait_in_this_tick = true;
        }

        match result {
            NodeResult::Success => {
                // For instant actions (non-wait), continue in same tick
                // For wait nodes, we'll check in the next iteration
                if !is_wait_node {
                    continue;
                }
                // Wait completed, continue to next child if it's instant
                continue;
            }
            NodeResult::Running => return Ok(NodeResult::Running),
            NodeResult::Failure => return Ok(NodeResult::Failure),
            NodeResult::Ready => return Ok(NodeResult::Running),
        }
    }
    // All children have succeeded
    Ok(NodeResult::Success)
}

fn execute_selector(
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    for child in children {
        let result = execute_node(child, context, delta_time)?;
        match result {
            NodeResult::Success => return Ok(NodeResult::Success),
            NodeResult::Running => return Ok(NodeResult::Running),
            NodeResult::Failure => continue,
            NodeResult::Ready => continue,
        }
    }
    Ok(NodeResult::Failure)
}

fn execute_parallel(
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    let mut has_running = false;
    let mut has_failure = false;

    for child in children {
        let result = execute_node(child, context, delta_time)?;
        match result {
            NodeResult::Running => has_running = true,
            NodeResult::Failure => has_failure = true,
            _ => {}
        }
    }

    if has_running {
        Ok(NodeResult::Running)
    } else if has_failure {
        Ok(NodeResult::Failure)
    } else {
        Ok(NodeResult::Success)
    }
}

fn execute_race(
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    for child in children.iter_mut() {
        let result = execute_node(child, context, delta_time)?;
        match result {
            NodeResult::Success => return Ok(NodeResult::Success),
            NodeResult::Running => continue,
            NodeResult::Failure => continue,
            NodeResult::Ready => continue,
        }
    }

    // If no child succeeded, check if any are still running
    for child in children {
        if matches!(child.state, NodeState::Running) {
            return Ok(NodeResult::Running);
        }
    }

    Ok(NodeResult::Failure)
}

fn execute_all(
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    let mut has_running = false;
    let mut has_success = false;

    for child in children {
        let result = execute_node(child, context, delta_time)?;
        match result {
            NodeResult::Running => has_running = true,
            NodeResult::Success => has_success = true,
            _ => {}
        }
    }

    if has_running {
        Ok(NodeResult::Running)
    } else if has_success {
        Ok(NodeResult::Success)
    } else {
        Ok(NodeResult::Failure)
    }
}

fn execute_lotto(
    children: &mut [Node],
    weights: Option<&Vec<u32>>,
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    if children.is_empty() {
        return Ok(NodeResult::Failure);
    }

    let index = if let Some(weights) = weights {
        // Weighted random selection
        if weights.len() != children.len() {
            // Fallback to uniform random if weights don't match children count
            (std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
                % children.len() as u128) as usize
        } else {
            // Calculate total weight
            let total_weight: u32 = weights.iter().sum();
            if total_weight == 0 {
                // Fallback to uniform random if all weights are zero
                (std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap()
                    .as_nanos()
                    % children.len() as u128) as usize
            } else {
                // Generate random number and find weighted index
                let mut random_value = (std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap()
                    .as_nanos()
                    % total_weight as u128) as u32;

                let mut selected_index = 0;
                for (i, &weight) in weights.iter().enumerate() {
                    if random_value < weight {
                        selected_index = i;
                        break;
                    }
                    random_value -= weight;
                }
                selected_index
            }
        }
    } else {
        // Uniform random selection
        (std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
            % children.len() as u128) as usize
    };

    execute_node(&mut children[index], context, delta_time)
}

fn execute_repeat(
    child: &mut Node,
    iterations: Option<u32>,
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    let max_iterations = iterations.unwrap_or(1);

    // Track iterations in elapsed_time field
    let current_iteration = child.elapsed_time as u32;

    if current_iteration >= max_iterations {
        return Ok(NodeResult::Success);
    }

    // For instant actions, execute all iterations in one tick
    loop {
        let result = execute_node(child, context, delta_time)?;

        match result {
            NodeResult::Success => {
                // Increment iteration count
                child.elapsed_time = (child.elapsed_time as u32 + 1) as f64;

                // Check if we've completed all iterations
                if child.elapsed_time as u32 >= max_iterations {
                    return Ok(NodeResult::Success);
                } else {
                    // Reset child for next iteration and continue in same tick
                    child.state = NodeState::Ready;
                    continue; // Execute next iteration immediately
                }
            }
            NodeResult::Running => {
                // If child is still running, we need to wait for next tick
                return Ok(NodeResult::Running);
            }
            NodeResult::Failure => {
                // Repeat stops on failure
                return Ok(NodeResult::Failure);
            }
            NodeResult::Ready => {
                // This shouldn't happen, but treat as Running
                return Ok(NodeResult::Running);
            }
        }
    }
}

fn execute_retry(
    child: &mut Node,
    attempts: Option<u32>,
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    let max_attempts = attempts.unwrap_or(1);

    // Track attempts in elapsed_time field (reusing the field for different purpose)
    let current_attempt = child.elapsed_time as u32;

    if current_attempt >= max_attempts {
        return Ok(NodeResult::Failure); // Exhausted all attempts
    }

    // For instant actions, execute all attempts in one tick
    loop {
        let result = execute_node(child, context, delta_time)?;

        match result {
            NodeResult::Success => {
                // Success on any attempt
                return Ok(NodeResult::Success);
            }
            NodeResult::Running => {
                // If child is still running, we need to wait for next tick
                return Ok(NodeResult::Running);
            }
            NodeResult::Failure => {
                // Increment attempt count
                child.elapsed_time = (child.elapsed_time as u32 + 1) as f64;

                // Check if we have more attempts
                if child.elapsed_time as u32 >= max_attempts {
                    return Ok(NodeResult::Failure); // Exhausted all attempts
                } else {
                    // Reset child for next attempt and continue in same tick
                    child.state = NodeState::Ready;
                    continue; // Try next attempt immediately
                }
            }
            NodeResult::Ready => {
                // This shouldn't happen, but treat as Running
                return Ok(NodeResult::Running);
            }
        }
    }
}

fn execute_flip(child: &mut Node, context: &TreeContext, delta_time: f64) -> Result<NodeResult> {
    let result = execute_node(child, context, delta_time)?;
    Ok(match result {
        NodeResult::Success => NodeResult::Failure,
        NodeResult::Failure => NodeResult::Success,
        other => other,
    })
}

fn execute_succeed(child: &mut Node, context: &TreeContext, delta_time: f64) -> Result<NodeResult> {
    let result = execute_node(child, context, delta_time)?;
    Ok(match result {
        NodeResult::Running => NodeResult::Running,
        _ => NodeResult::Success,
    })
}

fn execute_fail(child: &mut Node, context: &TreeContext, delta_time: f64) -> Result<NodeResult> {
    let result = execute_node(child, context, delta_time)?;
    Ok(match result {
        NodeResult::Running => NodeResult::Running,
        _ => NodeResult::Failure,
    })
}

fn execute_while(
    condition: &str,
    args: &[serde_json::Value],
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    // Check condition first
    if !context.evaluate_condition(condition, args)? {
        // Condition is false, while loop exits successfully
        return Ok(NodeResult::Success);
    }

    // Execute children as a sequence
    let result = execute_sequence(children, context, delta_time)?;
    match result {
        NodeResult::Running => Ok(NodeResult::Running),
        NodeResult::Failure => Ok(NodeResult::Failure),
        NodeResult::Success => {
            // Reset children for next iteration and continue while loop
            for child in children.iter_mut() {
                child.reset();
            }
            // Continue the while loop by returning Running
            // The condition will be checked again on the next tick
            Ok(NodeResult::Running)
        }
        NodeResult::Ready => Ok(NodeResult::Running),
    }
}

fn execute_until(
    condition: &str,
    args: &[serde_json::Value],
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    // Check condition first
    if context.evaluate_condition(condition, args)? {
        // Condition is true, until loop exits successfully
        return Ok(NodeResult::Success);
    }

    // Execute children as a sequence
    let result = execute_sequence(children, context, delta_time)?;
    match result {
        NodeResult::Running => Ok(NodeResult::Running),
        NodeResult::Failure => Ok(NodeResult::Failure),
        NodeResult::Success => {
            // Reset children for next iteration and continue until loop
            for child in children.iter_mut() {
                child.reset();
            }
            // Continue the until loop by returning Running
            // The condition will be checked again on the next tick
            Ok(NodeResult::Running)
        }
        NodeResult::Ready => Ok(NodeResult::Running),
    }
}

fn execute_while_all(
    condition: &str,
    args: &[serde_json::Value],
    children: &mut [Node],
    context: &TreeContext,
    delta_time: f64,
) -> Result<NodeResult> {
    // WhileAll semantics: Execute ALL children as a sequence first
    let result = execute_sequence(children, context, delta_time)?;
    match result {
        NodeResult::Running => Ok(NodeResult::Running),
        NodeResult::Failure => Ok(NodeResult::Failure),
        NodeResult::Success => {
            // All children succeeded, now check condition
            if context.evaluate_condition(condition, args)? {
                // Condition is true, reset children and continue loop
                for child in children.iter_mut() {
                    child.reset();
                }
                // Continue the while loop by returning Running
                Ok(NodeResult::Running)
            } else {
                // Condition is false, while loop exits successfully
                Ok(NodeResult::Success)
            }
        }
        NodeResult::Ready => Ok(NodeResult::Running),
    }
}

fn evaluate_guard(guard: &Guard, context: &TreeContext) -> Result<bool> {
    match guard {
        Guard::While { condition, args } => context.evaluate_condition(condition, args),
        Guard::Until { condition, args } => {
            let result = context.evaluate_condition(condition, args)?;
            Ok(!result)
        }
    }
}