//! Timeout Manager for ML Safety
//!
//! This module provides timeout management for ML operations to prevent
//! hanging operations that could block the HFT system.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::{Mutex, RwLock};
use tokio::time::timeout;
use tracing::{debug, error, warn};

use super::{MLSafetyConfig, MLSafetyError, SafetyResult};

/// Information about a running operation
#[derive(Debug, Clone)]
struct OperationInfo {
    name: String,
    started_at: Instant,
    timeout_duration: Duration,
}

/// Timeout manager for ML operations
#[derive(Debug)]
pub struct TimeoutManager {
    config: MLSafetyConfig,
    active_operations: Arc<RwLock<HashMap<String, OperationInfo>>>,
    operation_counter: Arc<Mutex<u64>>,
    default_timeout: Duration,
    max_concurrent_operations: usize,
}

impl TimeoutManager {
    /// Create new timeout manager
    pub fn new(config: &MLSafetyConfig) -> Self {
        Self {
            config: config.clone(),
            active_operations: Arc::new(RwLock::new(HashMap::new())),
            operation_counter: Arc::new(Mutex::new(0)),
            default_timeout: Duration::from_millis(config.max_inference_timeout_ms),
            max_concurrent_operations: 100, // Safe default for HFT
        }
    }

    /// Execute operation with timeout
    pub async fn execute_with_timeout<F, T>(
        &self,
        operation_name: &str,
        operation: F,
        timeout_duration: Duration,
    ) -> SafetyResult<T>
    where
        F: FnOnce() -> SafetyResult<T> + Send + 'static,
        T: Send + 'static,
    {
        // Check concurrent operation limit
        {
            let operations = self.active_operations.read().await;
            if operations.len() >= self.max_concurrent_operations {
                return Err(MLSafetyError::ResourceExhausted {
                    resource: "Maximum concurrent operations exceeded".to_owned(),
                });
            }
        }

        // Generate unique operation ID
        let operation_id = {
            let mut counter = self.operation_counter.lock().await;
            *counter += 1;
            format!("{}_{}", operation_name, *counter)
        };

        // Register operation
        {
            let mut operations = self.active_operations.write().await;
            operations.insert(
                operation_id.clone(),
                OperationInfo {
                    name: operation_name.to_string(),
                    started_at: Instant::now(),
                    timeout_duration,
                },
            );
        }

        debug!(
            "Starting operation {} with timeout {}ms",
            operation_id,
            timeout_duration.as_millis()
        );

        // Execute with timeout
        let result = timeout(timeout_duration, async move {
            tokio::task::spawn_blocking(operation).await.map_err(|e| {
                MLSafetyError::TensorSafety {
                    reason: format!("Operation panicked: {}", e),
                }
            })?
        })
        .await;

        // Unregister operation
        {
            let mut operations = self.active_operations.write().await;
            if let Some(info) = operations.remove(&operation_id) {
                let elapsed = info.started_at.elapsed();
                debug!(
                    "Completed operation {} in {:.3}s",
                    operation_id,
                    elapsed.as_secs_f64()
                );
            }
        }

        // Handle result
        match result {
            Ok(value) => value,
            Err(_) => {
                error!(
                    "Operation {} timed out after {:.1}s",
                    operation_id,
                    timeout_duration.as_secs_f64()
                );
                Err(MLSafetyError::Timeout {
                    timeout_ms: timeout_duration.as_millis() as u64,
                })
            },
        }
    }

    /// Execute with default timeout
    pub async fn execute_with_default_timeout<F, T>(
        &self,
        operation_name: &str,
        operation: F,
    ) -> SafetyResult<T>
    where
        F: FnOnce() -> SafetyResult<T> + Send + 'static,
        T: Send + 'static,
    {
        self.execute_with_timeout(operation_name, operation, self.default_timeout)
            .await
    }

    /// Get currently active operations
    pub async fn get_active_operations(&self) -> Vec<(String, Duration)> {
        let operations = self.active_operations.read().await;
        operations
            .iter()
            .map(|(id, info)| (id.clone(), info.started_at.elapsed()))
            .collect()
    }

    /// Get operation statistics
    pub async fn get_operation_stats(&self) -> HashMap<String, String> {
        let operations = self.active_operations.read().await;
        let mut stats = HashMap::new();

        stats.insert(
            "active_operations".to_owned(),
            operations.len().to_string(),
        );
        stats.insert(
            "max_concurrent".to_owned(),
            self.max_concurrent_operations.to_string(),
        );
        stats.insert(
            "default_timeout_ms".to_owned(),
            self.default_timeout.as_millis().to_string(),
        );

        if !operations.is_empty() {
            let total_duration: Duration = operations
                .values()
                .map(|info| info.started_at.elapsed())
                .sum();
            let avg_duration = total_duration / operations.len() as u32;
            stats.insert(
                "avg_operation_duration_ms".to_owned(),
                avg_duration.as_millis().to_string(),
            );

            if let Some(longest_running) = operations
                .values()
                .max_by_key(|info| info.started_at.elapsed())
            {
                stats.insert(
                    "longest_running_operation".to_owned(),
                    longest_running.name.clone(),
                );
                stats.insert(
                    "longest_running_duration_ms".to_owned(),
                    longest_running.started_at.elapsed().as_millis().to_string(),
                );
            }
        }

        stats
    }

    /// Check for operations that have exceeded their timeout
    pub async fn check_for_stuck_operations(&self) -> Vec<String> {
        let operations = self.active_operations.read().await;
        let mut stuck_operations = Vec::new();

        for (id, info) in operations.iter() {
            let elapsed = info.started_at.elapsed();
            if elapsed > info.timeout_duration * 2 {
                warn!(
                    "Stuck operation detected: {} running for {:.1}s",
                    id,
                    elapsed.as_secs_f64()
                );
                stuck_operations.push(id.clone());
            }
        }

        stuck_operations
    }

    /// Emergency shutdown all operations
    pub async fn shutdown_all(&self) {
        let mut operations = self.active_operations.write().await;
        if !operations.is_empty() {
            warn!(
                "Emergency shutdown: terminating {} active operations",
                operations.len()
            );
            operations.clear();
        }
    }

    /// Clone for multi-threaded access
    pub fn clone_handle(&self) -> Self {
        Self {
            config: self.config.clone(),
            active_operations: Arc::clone(&self.active_operations),
            operation_counter: Arc::clone(&self.operation_counter),
            default_timeout: self.default_timeout,
            max_concurrent_operations: self.max_concurrent_operations,
        }
    }
}
