swarm-engine-core 0.1.6

//! Completion state tracking
//!
//! Tracks the completion state of workers, exploration, and environment.

use std::collections::HashMap;

use crate::types::WorkerId;

use super::reason::TerminationVerdict;

/// Completion state - Single Source of Truth for completion tracking
#[derive(Debug, Clone, Default)]
pub struct CompletionState {
    /// Worker completion results (worker_id -> (success, message))
    worker_results: HashMap<WorkerId, WorkerResult>,

    /// Exploration completion flag (frontier exhausted or explicitly marked)
    exploration_done: bool,

    /// Exploration was exhausted (no more frontiers) vs explicitly completed
    exploration_exhausted: bool,

    /// Environment signaled done (from WorkResult::Done)
    environment_done: bool,

    /// Final verdict (set once determined)
    final_verdict: Option<TerminationVerdict>,

    /// Tick when completion was first detected
    completion_tick: Option<u64>,
}

/// Result of a worker's completion
#[derive(Debug, Clone)]
pub struct WorkerResult {
    /// Success flag
    pub success: bool,
    /// Optional message
    pub message: Option<String>,
    /// Tick when completed
    pub tick: u64,
}

impl CompletionState {
    /// Create new empty completion state
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a worker's completion
    pub fn record_worker_done(
        &mut self,
        worker_id: WorkerId,
        success: bool,
        message: Option<String>,
        tick: u64,
    ) {
        self.worker_results.insert(
            worker_id,
            WorkerResult {
                success,
                message,
                tick,
            },
        );

        if success {
            self.environment_done = true;
            if self.completion_tick.is_none() {
                self.completion_tick = Some(tick);
            }
        }
    }

    /// Mark exploration as done
    pub fn mark_exploration_done(&mut self, exhausted: bool) {
        self.exploration_done = true;
        self.exploration_exhausted = exhausted;
    }

    /// Set the final verdict
    pub fn set_verdict(&mut self, verdict: TerminationVerdict) {
        if self.final_verdict.is_none() {
            self.final_verdict = Some(verdict);
        }
    }

    /// Get the final verdict if determined
    pub fn verdict(&self) -> Option<&TerminationVerdict> {
        self.final_verdict.as_ref()
    }

    /// Check if environment is done (any worker completed successfully)
    pub fn is_environment_done(&self) -> bool {
        self.environment_done
    }

    /// Check if exploration is done
    pub fn is_exploration_done(&self) -> bool {
        self.exploration_done
    }

    /// Check if exploration was exhausted (not explicitly completed)
    pub fn is_exploration_exhausted(&self) -> bool {
        self.exploration_exhausted
    }

    /// Check if a verdict has been determined
    pub fn has_verdict(&self) -> bool {
        self.final_verdict.is_some()
    }

    /// Get all completed workers
    pub fn completed_workers(&self) -> &HashMap<WorkerId, WorkerResult> {
        &self.worker_results
    }

    /// Check if any worker completed successfully
    pub fn any_worker_succeeded(&self) -> bool {
        self.worker_results.values().any(|r| r.success)
    }

    /// Check if all workers that completed were successful
    pub fn all_completed_workers_succeeded(&self) -> bool {
        !self.worker_results.is_empty() && self.worker_results.values().all(|r| r.success)
    }

    /// Get the tick when completion was first detected
    pub fn completion_tick(&self) -> Option<u64> {
        self.completion_tick
    }

    /// Get the first successful worker result
    pub fn first_success(&self) -> Option<(&WorkerId, &WorkerResult)> {
        self.worker_results.iter().find(|(_, r)| r.success)
    }

    /// Reset completion state (for reuse)
    pub fn reset(&mut self) {
        self.worker_results.clear();
        self.exploration_done = false;
        self.exploration_exhausted = false;
        self.environment_done = false;
        self.final_verdict = None;
        self.completion_tick = None;
    }
}

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

    #[test]
    fn test_completion_state_default() {
        let state = CompletionState::new();
        assert!(!state.is_environment_done());
        assert!(!state.is_exploration_done());
        assert!(!state.has_verdict());
    }

    #[test]
    fn test_record_worker_done() {
        let mut state = CompletionState::new();

        state.record_worker_done(WorkerId(0), true, Some("Done!".to_string()), 10);

        assert!(state.is_environment_done());
        assert!(state.any_worker_succeeded());
        assert_eq!(state.completion_tick(), Some(10));
    }

    #[test]
    fn test_exploration_exhausted() {
        let mut state = CompletionState::new();

        state.mark_exploration_done(true);

        assert!(state.is_exploration_done());
        assert!(state.is_exploration_exhausted());
    }

    #[test]
    fn test_exploration_completed() {
        let mut state = CompletionState::new();

        state.mark_exploration_done(false);

        assert!(state.is_exploration_done());
        assert!(!state.is_exploration_exhausted());
    }
}