swarm-engine-core 0.1.6

//! Termination Judge - Single Source of Truth for termination decisions
//!
//! Centralizes all termination logic to prevent race conditions and ensure
//! consistent behavior across the orchestrator.

use tracing::{debug, info};

use crate::types::WorkerId;

use super::config::{ExhaustedBehavior, MultiWorkerStrategy, TerminationConfig};
use super::reason::{FailureReason, SuccessReason, TerminationVerdict};
use super::state::CompletionState;

/// External termination request
#[derive(Debug, Clone)]
pub struct TerminationRequest {
    /// Reason for termination
    pub reason: String,
    /// Tick when requested
    pub tick: u64,
}

/// Termination Judge - Single Source of Truth
///
/// All termination-related decisions go through this struct.
/// This prevents race conditions between different completion signals.
#[derive(Debug)]
pub struct TerminationJudge {
    /// Configuration
    config: TerminationConfig,

    /// Completion state
    state: CompletionState,

    /// External termination request
    external_request: Option<TerminationRequest>,

    /// Current tick (updated each tick)
    current_tick: u64,

    /// Total worker count
    worker_count: usize,

    /// Consecutive error count
    consecutive_errors: u64,
}

impl TerminationJudge {
    /// Create a new termination judge
    pub fn new(config: TerminationConfig, worker_count: usize) -> Self {
        Self {
            config,
            state: CompletionState::new(),
            external_request: None,
            current_tick: 0,
            worker_count,
            consecutive_errors: 0,
        }
    }

    /// Update the current tick
    pub fn set_tick(&mut self, tick: u64) {
        self.current_tick = tick;
    }

    /// Get the current tick
    pub fn current_tick(&self) -> u64 {
        self.current_tick
    }

    // =========================================================================
    // Notification Methods - Called by Orchestrator
    // =========================================================================

    /// Notify that a worker completed
    ///
    /// Called when WorkResult::Done is received.
    pub fn notify_worker_done(
        &mut self,
        worker_id: WorkerId,
        success: bool,
        message: Option<String>,
    ) {
        info!(
            worker_id = worker_id.0,
            success = success,
            message = ?message,
            tick = self.current_tick,
            "TerminationJudge: worker done notification"
        );

        self.state
            .record_worker_done(worker_id, success, message.clone(), self.current_tick);

        // Reset consecutive errors on success
        if success {
            self.consecutive_errors = 0;
        }

        // Reevaluate termination
        self.reevaluate();
    }

    /// Notify that exploration completed or exhausted
    ///
    /// Called when ExplorationSpace signals completion.
    pub fn notify_exploration_complete(&mut self, exhausted: bool) {
        info!(
            exhausted = exhausted,
            tick = self.current_tick,
            "TerminationJudge: exploration complete notification"
        );

        self.state.mark_exploration_done(exhausted);

        // Reevaluate termination
        self.reevaluate();
    }

    /// Notify that an error occurred
    pub fn notify_error(&mut self) {
        self.consecutive_errors += 1;

        if let Some(max) = self.config.max_consecutive_errors {
            if self.consecutive_errors >= max {
                info!(
                    errors = self.consecutive_errors,
                    max = max,
                    "TerminationJudge: max consecutive errors exceeded"
                );
                self.state.set_verdict(TerminationVerdict::Failure {
                    reason: FailureReason::MaxErrorsExceeded {
                        count: self.consecutive_errors,
                        limit: max,
                    },
                });
            }
        }
    }

    /// Request external termination
    pub fn request_terminate(&mut self, reason: impl Into<String>) {
        let reason = reason.into();
        info!(
            reason = %reason,
            tick = self.current_tick,
            "TerminationJudge: external termination requested"
        );

        self.external_request = Some(TerminationRequest {
            reason: reason.clone(),
            tick: self.current_tick,
        });

        // Set verdict immediately
        self.state
            .set_verdict(TerminationVerdict::ExternalStop { reason });
    }

    // =========================================================================
    // Query Methods - Called by Orchestrator
    // =========================================================================

    /// Check if the orchestrator should terminate
    ///
    /// This is the ONLY method that should be used to decide loop termination.
    pub fn should_terminate(&self) -> bool {
        // External request always terminates immediately
        if self.external_request.is_some() {
            return true;
        }

        // Verdict determined
        if self.state.has_verdict() {
            return true;
        }

        // Max ticks exceeded
        if self.config.max_ticks > 0 && self.current_tick >= self.config.max_ticks {
            return true;
        }

        false
    }

    /// Check if guidance generation should be skipped
    ///
    /// Returns true if we're in a terminal state and shouldn't generate new work.
    pub fn should_skip_guidance(&self) -> bool {
        // External request - stop immediately
        if self.external_request.is_some() {
            return true;
        }

        // Verdict already determined
        if self.state.has_verdict() {
            return true;
        }

        // Environment done (worker completed successfully)
        if self.state.is_environment_done() {
            return true;
        }

        // Exploration done
        if self.state.is_exploration_done() {
            return true;
        }

        false
    }

    /// Get the final verdict
    ///
    /// Call this after should_terminate() returns true.
    pub fn verdict(&self) -> TerminationVerdict {
        // Already determined
        if let Some(verdict) = self.state.verdict() {
            return verdict.clone();
        }

        // External request
        if let Some(ref req) = self.external_request {
            return TerminationVerdict::ExternalStop {
                reason: req.reason.clone(),
            };
        }

        // Timeout
        if self.config.max_ticks > 0 && self.current_tick >= self.config.max_ticks {
            let partial_success = self.state.any_worker_succeeded();
            return TerminationVerdict::Timeout { partial_success };
        }

        // Should not reach here if called correctly
        TerminationVerdict::Failure {
            reason: FailureReason::InternalError {
                message: "verdict() called without termination condition".to_string(),
            },
        }
    }

    /// Check if environment is done
    pub fn is_environment_done(&self) -> bool {
        self.state.is_environment_done()
    }

    /// Get completion state reference
    pub fn completion_state(&self) -> &CompletionState {
        &self.state
    }

    // =========================================================================
    // Internal Methods
    // =========================================================================

    /// Reevaluate termination based on current state
    fn reevaluate(&mut self) {
        // Already have a verdict
        if self.state.has_verdict() {
            return;
        }

        // Check worker completion based on strategy
        if let Some(verdict) = self.evaluate_worker_completion() {
            debug!(verdict = ?verdict, "TerminationJudge: verdict from worker completion");
            self.state.set_verdict(verdict);
            return;
        }

        // Check exploration exhaustion
        if self.state.is_exploration_exhausted() {
            let verdict = self.evaluate_exhaustion();
            debug!(verdict = ?verdict, "TerminationJudge: verdict from exploration exhaustion");
            self.state.set_verdict(verdict);
        }
    }

    /// Evaluate worker completion based on multi-worker strategy
    fn evaluate_worker_completion(&self) -> Option<TerminationVerdict> {
        let completed = self.state.completed_workers();

        match self.config.multi_worker_strategy {
            MultiWorkerStrategy::FirstSuccess => {
                // First successful worker triggers success
                if let Some((worker_id, result)) = self.state.first_success() {
                    return Some(TerminationVerdict::Success {
                        reason: SuccessReason::WorkerDone {
                            worker_id: worker_id.0,
                            message: result.message.clone(),
                        },
                    });
                }
                None
            }

            MultiWorkerStrategy::AllComplete => {
                // All workers must complete
                if completed.len() >= self.worker_count {
                    if self.state.any_worker_succeeded() {
                        return Some(TerminationVerdict::Success {
                            reason: SuccessReason::ConditionsMet,
                        });
                    } else {
                        return Some(TerminationVerdict::Failure {
                            reason: FailureReason::WorkerFailed {
                                worker_id: 0,
                                message: Some("All workers failed".to_string()),
                            },
                        });
                    }
                }
                None
            }

            MultiWorkerStrategy::AllSuccess => {
                // All workers must succeed
                if completed.len() >= self.worker_count {
                    if self.state.all_completed_workers_succeeded() {
                        return Some(TerminationVerdict::Success {
                            reason: SuccessReason::ConditionsMet,
                        });
                    } else {
                        return Some(TerminationVerdict::Failure {
                            reason: FailureReason::WorkerFailed {
                                worker_id: 0,
                                message: Some("Not all workers succeeded".to_string()),
                            },
                        });
                    }
                }
                None
            }

            MultiWorkerStrategy::Conditions => {
                // Let external conditions handle it
                None
            }
        }
    }

    /// Evaluate exhaustion based on config
    fn evaluate_exhaustion(&self) -> TerminationVerdict {
        match self.config.on_exhausted {
            ExhaustedBehavior::Fail => TerminationVerdict::Failure {
                reason: FailureReason::ExplorationExhausted,
            },

            ExhaustedBehavior::Success => TerminationVerdict::Success {
                reason: SuccessReason::ExplorationComplete,
            },

            ExhaustedBehavior::CheckConditions => {
                // If any worker succeeded, consider it success
                if self.state.any_worker_succeeded() {
                    TerminationVerdict::Success {
                        reason: SuccessReason::ExplorationComplete,
                    }
                } else {
                    TerminationVerdict::Failure {
                        reason: FailureReason::ExplorationExhausted,
                    }
                }
            }
        }
    }
}

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

    fn default_judge() -> TerminationJudge {
        TerminationJudge::new(TerminationConfig::default(), 1)
    }

    #[test]
    fn test_initial_state() {
        let judge = default_judge();
        assert!(!judge.should_terminate());
        assert!(!judge.should_skip_guidance());
    }

    #[test]
    fn test_worker_done_success() {
        let mut judge = default_judge();
        judge.set_tick(5);

        judge.notify_worker_done(WorkerId(0), true, Some("Done!".to_string()));

        assert!(judge.should_terminate());
        assert!(judge.should_skip_guidance());
        assert!(judge.verdict().is_success());
    }

    #[test]
    fn test_external_termination() {
        let mut judge = default_judge();

        judge.request_terminate("User requested stop");

        assert!(judge.should_terminate());
        assert!(judge.should_skip_guidance());
        assert!(matches!(
            judge.verdict(),
            TerminationVerdict::ExternalStop { .. }
        ));
    }

    #[test]
    fn test_max_ticks_timeout() {
        let config = TerminationConfig::with_max_ticks(100);
        let mut judge = TerminationJudge::new(config, 1);

        judge.set_tick(100);

        assert!(judge.should_terminate());
        assert!(matches!(
            judge.verdict(),
            TerminationVerdict::Timeout { .. }
        ));
    }

    #[test]
    fn test_exploration_exhausted() {
        let mut judge = default_judge();

        judge.notify_exploration_complete(true); // exhausted

        assert!(judge.should_terminate());
        assert!(matches!(
            judge.verdict(),
            TerminationVerdict::Failure {
                reason: FailureReason::ExplorationExhausted
            }
        ));
    }

    #[test]
    fn test_all_success_strategy() {
        let config =
            TerminationConfig::default().multi_worker_strategy(MultiWorkerStrategy::AllSuccess);
        let mut judge = TerminationJudge::new(config, 2);

        // First worker succeeds
        judge.notify_worker_done(WorkerId(0), true, None);
        assert!(!judge.should_terminate()); // Not yet

        // Second worker succeeds
        judge.notify_worker_done(WorkerId(1), true, None);
        assert!(judge.should_terminate());
        assert!(judge.verdict().is_success());
    }

    #[test]
    fn test_first_success_strategy() {
        let mut judge = default_judge();

        // First worker succeeds - should terminate immediately
        judge.notify_worker_done(WorkerId(0), true, None);

        assert!(judge.should_terminate());
        assert!(judge.verdict().is_success());
    }
}