prompthive 0.2.8

//! Signal handling and graceful shutdown for PromptHive
//!
//! This module provides cross-platform signal handling to ensure clean shutdown
//! of PromptHive processes, including:
//! - SIGTERM/SIGINT handling for graceful shutdown
//! - Resource cleanup and state preservation
//! - Active operation completion before termination
//! - Proper cleanup of temporary files and locks

use anyhow::Result;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{debug, info, warn};

/// Global shutdown signal flag
static SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false);

/// Shutdown handler configuration
#[derive(Debug, Clone)]
pub struct ShutdownConfig {
    /// Maximum time to wait for graceful shutdown
    pub grace_period: Duration,
    /// Whether to force cleanup temporary files
    pub cleanup_temp_files: bool,
    /// Whether to save current state before shutdown
    pub save_state: bool,
}

impl Default for ShutdownConfig {
    fn default() -> Self {
        Self {
            grace_period: Duration::from_secs(30),
            cleanup_temp_files: true,
            save_state: true,
        }
    }
}

/// Shutdown coordinator for managing graceful application termination
pub struct ShutdownHandler {
    config: ShutdownConfig,
    shutdown_flag: Arc<AtomicBool>,
    start_time: Instant,
}

impl Default for ShutdownHandler {
    fn default() -> Self {
        Self::new()
    }
}

impl ShutdownHandler {
    /// Create a new shutdown handler with default configuration
    pub fn new() -> Self {
        Self::with_config(ShutdownConfig::default())
    }

    /// Create a new shutdown handler with custom configuration
    pub fn with_config(config: ShutdownConfig) -> Self {
        Self {
            config,
            shutdown_flag: Arc::new(AtomicBool::new(false)),
            start_time: Instant::now(),
        }
    }

    /// Initialize signal handlers for graceful shutdown
    pub fn setup_signal_handlers(&self) -> Result<()> {
        let shutdown_flag = Arc::clone(&self.shutdown_flag);

        #[cfg(unix)]
        {
            use signal_hook::{consts::SIGINT, consts::SIGTERM, iterator::Signals};
            use std::thread;

            let mut signals = Signals::new([SIGINT, SIGTERM])?;
            let shutdown_flag_clone = Arc::clone(&shutdown_flag);

            thread::spawn(move || {
                for sig in signals.forever() {
                    match sig {
                        SIGINT => {
                            info!("Received SIGINT (Ctrl+C), initiating graceful shutdown...");
                            shutdown_flag_clone.store(true, Ordering::SeqCst);
                            SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
                            break;
                        }
                        SIGTERM => {
                            info!("Received SIGTERM, initiating graceful shutdown...");
                            shutdown_flag_clone.store(true, Ordering::SeqCst);
                            SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
                            break;
                        }
                        _ => {
                            warn!("Received unexpected signal: {}", sig);
                        }
                    }
                }
            });
        }

        #[cfg(windows)]
        {
            use std::os::raw::c_ulong;
            use windows_sys::Win32::System::Console::{
                SetConsoleCtrlHandler, CTRL_BREAK_EVENT, CTRL_C_EVENT,
            };

            unsafe extern "system" fn ctrl_handler(ctrl_type: c_ulong) -> i32 {
                match ctrl_type {
                    CTRL_C_EVENT | CTRL_BREAK_EVENT => {
                        info!("Received Ctrl+C/Break, initiating graceful shutdown...");
                        SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
                        1 // TRUE - we handled the signal
                    }
                    _ => 0, // FALSE - let the default handler handle it
                }
            }

            unsafe {
                if SetConsoleCtrlHandler(Some(ctrl_handler), 1) == 0 {
                    return Err(anyhow::anyhow!(
                        "Failed to set Windows console control handler"
                    ));
                }
            }
        }

        info!("Signal handlers installed successfully");
        Ok(())
    }

    /// Check if shutdown has been requested
    pub fn is_shutdown_requested(&self) -> bool {
        self.shutdown_flag.load(Ordering::SeqCst) || SHUTDOWN_REQUESTED.load(Ordering::SeqCst)
    }

    /// Wait for shutdown signal or timeout
    pub fn wait_for_shutdown(&self) -> Result<()> {
        let start = Instant::now();

        info!(
            "Waiting for shutdown signal (grace period: {:?})",
            self.config.grace_period
        );

        while !self.is_shutdown_requested() {
            if start.elapsed() > self.config.grace_period {
                warn!("Grace period exceeded, forcing shutdown");
                break;
            }

            std::thread::sleep(Duration::from_millis(100));
        }

        if self.is_shutdown_requested() {
            info!("Shutdown signal received, beginning graceful shutdown");
        }

        Ok(())
    }

    /// Perform graceful shutdown cleanup
    pub fn shutdown(&self, storage: Option<&crate::Storage>) -> Result<()> {
        let shutdown_start = Instant::now();
        info!("Beginning graceful shutdown...");

        // 1. Stop accepting new operations
        debug!("Stopping new operations...");

        // 2. Wait for active operations to complete
        self.wait_for_active_operations()?;

        // 3. Save current state if requested
        if self.config.save_state {
            if let Some(storage) = storage {
                self.save_application_state(storage)?;
            }
        }

        // 4. Cleanup temporary files
        if self.config.cleanup_temp_files {
            self.cleanup_temporary_files()?;
        }

        // 5. Flush logs and cleanup resources
        self.cleanup_resources()?;

        let shutdown_duration = shutdown_start.elapsed();
        info!(
            "Graceful shutdown completed in {:?} (total uptime: {:?})",
            shutdown_duration,
            self.start_time.elapsed()
        );

        Ok(())
    }

    /// Wait for active operations to complete
    fn wait_for_active_operations(&self) -> Result<()> {
        let timeout = Duration::from_secs(10);
        let start = Instant::now();

        debug!("Waiting for active operations to complete...");

        // In a more complex application, you would check for:
        // - Active HTTP requests
        // - Running background tasks
        // - Open file handles
        // - Database connections

        // For PromptHive, most operations are quick and synchronous
        // So we just do a short wait for any in-flight operations
        std::thread::sleep(Duration::from_millis(100));

        if start.elapsed() >= timeout {
            warn!("Timeout waiting for active operations to complete");
        } else {
            debug!("All active operations completed");
        }

        Ok(())
    }

    /// Save current application state
    fn save_application_state(&self, _storage: &crate::Storage) -> Result<()> {
        debug!("Saving application state...");

        // Save current state, such as:
        // - In-memory caches
        // - Unsaved changes
        // - Configuration updates
        // - Session information

        // For PromptHive, most state is already persisted
        // This could include saving telemetry data, caches, etc.

        debug!("Application state saved successfully");
        Ok(())
    }

    /// Cleanup temporary files and directories
    fn cleanup_temporary_files(&self) -> Result<()> {
        debug!("Cleaning up temporary files...");

        use std::fs;
        use std::path::PathBuf;

        // Common temporary directories to clean
        let temp_dirs = [
            std::env::temp_dir().join("prompthive"),
            PathBuf::from("/tmp").join("prompthive"),
        ];

        for temp_dir in &temp_dirs {
            if temp_dir.exists() {
                match fs::remove_dir_all(temp_dir) {
                    Ok(()) => debug!("Cleaned temporary directory: {:?}", temp_dir),
                    Err(e) => warn!("Failed to clean temporary directory {:?}: {}", temp_dir, e),
                }
            }
        }

        debug!("Temporary file cleanup completed");
        Ok(())
    }

    /// Cleanup system resources
    fn cleanup_resources(&self) -> Result<()> {
        debug!("Cleaning up system resources...");

        // Flush any pending log writes
        // In a real implementation, you might:
        // - Close database connections
        // - Flush write buffers
        // - Release file locks
        // - Cleanup network connections

        debug!("System resource cleanup completed");
        Ok(())
    }
}

/// Global function to check if shutdown has been requested
pub fn is_shutdown_requested() -> bool {
    SHUTDOWN_REQUESTED.load(Ordering::SeqCst)
}

/// Global function to request shutdown
pub fn request_shutdown() {
    SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
    info!("Shutdown requested programmatically");
}

/// Utility trait for shutdown-aware operations
pub trait ShutdownAware {
    /// Check if the operation should continue or abort due to shutdown
    fn should_continue(&self) -> bool {
        !is_shutdown_requested()
    }

    /// Sleep with shutdown awareness - returns early if shutdown is requested
    fn interruptible_sleep(&self, duration: Duration) -> bool {
        let start = Instant::now();
        let check_interval = Duration::from_millis(100);

        while start.elapsed() < duration {
            if is_shutdown_requested() {
                return false; // Interrupted by shutdown
            }

            let remaining = duration - start.elapsed();
            let sleep_time = check_interval.min(remaining);
            std::thread::sleep(sleep_time);
        }

        true // Completed full duration
    }
}

// Implement ShutdownAware for common types
impl ShutdownAware for () {}

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

    #[test]
    fn test_shutdown_handler_creation() {
        let handler = ShutdownHandler::new();
        assert!(!handler.is_shutdown_requested());
    }

    #[test]
    fn test_shutdown_config() {
        let config = ShutdownConfig {
            grace_period: Duration::from_secs(5),
            cleanup_temp_files: true,
            save_state: false,
        };

        let handler = ShutdownHandler::with_config(config.clone());
        assert_eq!(handler.config.grace_period, Duration::from_secs(5));
        assert!(handler.config.cleanup_temp_files);
        assert!(!handler.config.save_state);
    }

    #[test]
    fn test_global_shutdown_request() {
        // Reset state
        SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
        assert!(!is_shutdown_requested());

        // Request shutdown
        request_shutdown();
        assert!(is_shutdown_requested());

        // Reset for other tests
        SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
    }

    #[test]
    fn test_shutdown_aware_trait() {
        // Use a lock to ensure test isolation
        use std::sync::Mutex;
        static TEST_LOCK: Mutex<()> = Mutex::new(());
        let _guard = TEST_LOCK.lock().unwrap();
        
        struct TestStruct;
        impl ShutdownAware for TestStruct {}

        let test_obj = TestStruct;

        // Reset shutdown state
        SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
        // Give it a moment to settle
        std::thread::sleep(std::time::Duration::from_millis(1));
        assert!(test_obj.should_continue());

        // Request shutdown
        request_shutdown();
        // Give it a moment to settle
        std::thread::sleep(std::time::Duration::from_millis(1));
        assert!(!test_obj.should_continue());

        // Reset for other tests
        SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
    }

    #[test]
    fn test_interruptible_sleep() {
        // Use a lock to ensure test isolation
        use std::sync::Mutex;
        static TEST_LOCK: Mutex<()> = Mutex::new(());
        let _guard = TEST_LOCK.lock().unwrap();
        
        struct TestStruct;
        impl ShutdownAware for TestStruct {}

        let test_obj = TestStruct;

        // Reset shutdown state
        SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
        // Give it a moment to settle
        std::thread::sleep(std::time::Duration::from_millis(1));

        let start = Instant::now();
        let completed = test_obj.interruptible_sleep(Duration::from_millis(50));
        let elapsed = start.elapsed();

        assert!(completed);
        assert!(elapsed >= Duration::from_millis(40)); // Allow some timing variance

        // Reset for other tests
        SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
    }
}