//! Graceful Shutdown Coordination

//!

//! AWP-07: Two-phase shutdown with active operation tracking.

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use std::time::{Duration, Instant};

// ----------------------------------------------------------------------------

// AWP-07: Graceful Shutdown

// ----------------------------------------------------------------------------

/// Result of a graceful shutdown operation.

#[derive(Debug, Clone, PartialEq, Eq)]

pub enum ShutdownResult {

    /// All operations completed cleanly.

    Clean,

    /// Timeout reached with operations still active.

    Timeout {

        /// Number of operations still active.

        remaining: usize,

    },

}

/// Graceful shutdown coordinator.

///

/// # Example

/// ```rust

/// use trueno::brick::GracefulShutdown;

/// use std::time::Duration;

///

/// let shutdown = GracefulShutdown::new(Duration::from_secs(5));

///

/// // Register an operation

/// let guard = shutdown.register().ok_or("shutdown in progress")?;

///

/// // ... do work ...

///

/// drop(guard);  // Operation complete

///

/// // Initiate shutdown

/// let result = shutdown.shutdown();

/// assert_eq!(result, trueno::brick::ShutdownResult::Clean);

/// # Ok::<(), &'static str>(())

/// ```

pub struct GracefulShutdown {

    /// Flag indicating shutdown has been requested.

    shutdown_requested: AtomicBool,

    /// Number of active operations.

    active_count: AtomicUsize,

    /// Shutdown timeout.

    timeout: Duration,

}

impl GracefulShutdown {

    /// Create a new shutdown coordinator.

    pub fn new(timeout: Duration) -> Self {

        Self {

            shutdown_requested: AtomicBool::new(false),

            active_count: AtomicUsize::new(0),

            timeout,

        }

    }

    /// Check if shutdown has been requested.

    #[must_use]

    pub fn is_shutdown_requested(&self) -> bool {

        self.shutdown_requested.load(Ordering::Acquire)

    }

    /// Get the current active operation count.

    #[must_use]

    pub fn active_count(&self) -> usize {

        self.active_count.load(Ordering::Acquire)

    }

    /// Register an active operation.

    ///

    /// Returns `None` if shutdown has already been requested.

    pub fn register(&self) -> Option<ShutdownGuard<'_>> {

        if self.is_shutdown_requested() {

            return None; // Reject new operations during shutdown

        }

        self.active_count.fetch_add(1, Ordering::AcqRel);

        Some(ShutdownGuard { shutdown: self })

    }

    /// Initiate graceful shutdown.

    ///

    /// This will:

    /// 1. Stop accepting new operations

    /// 2. Wait for in-flight operations to complete (up to timeout)

    /// 3. Return the result

    pub fn shutdown(&self) -> ShutdownResult {

        // Phase 1: Stop accepting new operations

        self.shutdown_requested.store(true, Ordering::Release);

        // Phase 2: Wait for in-flight operations

        let deadline = Instant::now() + self.timeout;

        loop {

            let active = self.active_count.load(Ordering::Acquire);

            if active == 0 {

                return ShutdownResult::Clean;

            }

            if Instant::now() >= deadline {

                return ShutdownResult::Timeout { remaining: active };

            }

            std::thread::sleep(Duration::from_millis(10));

        }

    }

    /// Reset the shutdown coordinator for reuse.

    pub fn reset(&self) {

        self.shutdown_requested.store(false, Ordering::Release);

        // Note: active_count should already be 0 if shutdown completed cleanly

    }

}

impl Default for GracefulShutdown {

    fn default() -> Self {

        Self::new(Duration::from_secs(30))

    }

}

/// Guard that decrements active count on drop.

pub struct ShutdownGuard<'a> {

    shutdown: &'a GracefulShutdown,

}

impl Drop for ShutdownGuard<'_> {

    fn drop(&mut self) {

        self.shutdown.active_count.fetch_sub(1, Ordering::AcqRel);

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_shutdown_result_eq() {

        assert_eq!(ShutdownResult::Clean, ShutdownResult::Clean);

        assert_eq!(

            ShutdownResult::Timeout { remaining: 5 },

            ShutdownResult::Timeout { remaining: 5 }

        );

        assert_ne!(ShutdownResult::Clean, ShutdownResult::Timeout { remaining: 0 });

    }

    #[test]

    fn test_graceful_shutdown_new() {

        let shutdown = GracefulShutdown::new(Duration::from_secs(5));

        assert!(!shutdown.is_shutdown_requested());

        assert_eq!(shutdown.active_count(), 0);

    }

    #[test]

    fn test_graceful_shutdown_default() {

        let shutdown = GracefulShutdown::default();

        assert!(!shutdown.is_shutdown_requested());

        assert_eq!(shutdown.timeout, Duration::from_secs(30));

    }

    #[test]

    fn test_graceful_shutdown_register() {

        let shutdown = GracefulShutdown::new(Duration::from_secs(5));

        let guard1 = shutdown.register();

        assert!(guard1.is_some());

        assert_eq!(shutdown.active_count(), 1);

        let guard2 = shutdown.register();

        assert!(guard2.is_some());

        assert_eq!(shutdown.active_count(), 2);

        drop(guard1);

        assert_eq!(shutdown.active_count(), 1);

        drop(guard2);

        assert_eq!(shutdown.active_count(), 0);

    }

    #[test]

    fn test_graceful_shutdown_clean() {

        let shutdown = GracefulShutdown::new(Duration::from_millis(100));

        // Register and immediately drop

        let guard = shutdown.register();

        drop(guard);

        let result = shutdown.shutdown();

        assert_eq!(result, ShutdownResult::Clean);

    }

    #[test]

    fn test_graceful_shutdown_rejects_after_shutdown() {

        let shutdown = GracefulShutdown::new(Duration::from_millis(10));

        // Initiate shutdown

        let result = shutdown.shutdown();

        assert_eq!(result, ShutdownResult::Clean);

        // Try to register after shutdown

        let guard = shutdown.register();

        assert!(guard.is_none(), "Should reject new operations after shutdown");

    }

    #[test]

    fn test_graceful_shutdown_reset() {

        let shutdown = GracefulShutdown::new(Duration::from_millis(10));

        // Shutdown

        shutdown.shutdown();

        assert!(shutdown.is_shutdown_requested());

        // Reset

        shutdown.reset();

        assert!(!shutdown.is_shutdown_requested());

        // Should accept new operations again

        let guard = shutdown.register();

        assert!(guard.is_some());

    }

    /// FALSIFICATION TEST: Verify timeout behavior

    ///

    /// The shutdown coordinator MUST return Timeout when operations

    /// don't complete within the specified duration.

    #[test]

    fn test_falsify_shutdown_timeout() {

        let shutdown = GracefulShutdown::new(Duration::from_millis(50));

        // Register an operation that we WON'T drop

        let _guard = shutdown.register();

        assert_eq!(shutdown.active_count(), 1);

        // Initiate shutdown - should timeout since operation is still active

        let start = Instant::now();

        let result = shutdown.shutdown();

        let elapsed = start.elapsed();

        // CRITICAL ASSERTIONS

        match result {

            ShutdownResult::Timeout { remaining } => {

                assert_eq!(remaining, 1, "Should report exactly 1 remaining operation");

            }

            ShutdownResult::Clean => {

                panic!("FALSIFICATION FAILED: Shutdown reported Clean with active operation");

            }

        }

        // Verify we actually waited (within tolerance)

        assert!(

            elapsed >= Duration::from_millis(40),

            "FALSIFICATION FAILED: Shutdown returned too early ({:?} < 40ms)",

            elapsed

        );

    }

    /// FALSIFICATION TEST: Verify guard drop decrements count

    ///

    /// If the guard drop doesn't work correctly, the system will

    /// either hang forever or report wrong active counts.

    #[test]

    fn test_falsify_guard_drop_semantics() {

        let shutdown = GracefulShutdown::new(Duration::from_secs(1));

        // Create many guards

        let mut guards: Vec<_> = (0..100).filter_map(|_| shutdown.register()).collect();

        assert_eq!(

            shutdown.active_count(),

            100,

            "FALSIFICATION FAILED: Not all guards registered"

        );

        // Drop half by truncating the vector (keeps first 50, drops last 50)

        guards.truncate(50);

        // CRITICAL: Active count must be exactly 50

        // If fetch_sub has a bug (e.g., fetch_add by mistake), this catches it

        assert_eq!(

            shutdown.active_count(),

            50,

            "FALSIFICATION FAILED: Guard drop did not correctly decrement count"

        );

        // Drop remaining 50

        drop(guards);

        assert_eq!(

            shutdown.active_count(),

            0,

            "FALSIFICATION FAILED: Final drop did not clear all guards"

        );

    }

    /// FALSIFICATION TEST: Verify concurrent registration rejection

    ///

    /// Once shutdown starts, NO new operations should be accepted,

    /// even if they race with the shutdown call.

    #[test]

    fn test_falsify_rejection_after_shutdown_flag() {

        let shutdown = GracefulShutdown::new(Duration::from_secs(5));

        // Manually set shutdown flag (simulating shutdown start)

        shutdown.shutdown_requested.store(true, Ordering::Release);

        // Try to register many operations

        let mut accepted = 0;

        for _ in 0..100 {

            if shutdown.register().is_some() {

                accepted += 1;

            }

        }

        // CRITICAL: NO operations should have been accepted

        assert_eq!(

            accepted, 0,

            "FALSIFICATION FAILED: {} operations accepted after shutdown flag set",

            accepted

        );

    }

}