//! Comprehensive ML Safety Framework
//!
//! This module provides enterprise-grade safety controls for all ML operations
//! in the Foxhunt trading system. All ML components MUST use these safety
//! controls to prevent trading system failures.

use common::types::Price;
use std;

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

// candle_core removed — Tensor/Device replaced with opaque wrappers.
// GPU tensor safety is now handled by cuda_autograd bounds checks.
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

// IntegerPrice replaced with common::Price

// Re-export safety modules
pub mod bounds_checker;
pub mod drift_detector;
pub mod drift_responder;
pub mod financial_validator;
pub mod gradient_safety;
pub mod math_ops;
pub mod memory_manager;
pub mod tensor_ops;
pub mod timeout_manager;

use bounds_checker::BoundsChecker;
use drift_detector::ModelDriftDetector;
use financial_validator::FinancialValidator;
// DO NOT RE-EXPORT - Use explicit imports at usage sites
use math_ops::SafeMathOps;
use memory_manager::SafeMemoryManager;
use tensor_ops::SafeTensorOps;
use timeout_manager::TimeoutManager;

// Re-export gradient safety types for external usage
pub use gradient_safety::{GradientSafetyConfig, GradientSafetyManager, GradientStatistics};

/// Safety enforcement level for DQN training
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum SafetyLevel {
    /// Stop training immediately on ANY anomaly (production default)
    Strict,
    /// Warn on anomalies, continue training (development default)
    Normal,
    /// Log only, never stop training (debugging only)
    Permissive,
}

/// Global safety configuration for all ML operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLSafetyConfig {
    /// Enable comprehensive safety checks (never disable in production)
    pub safety_enabled: bool,
    /// Maximum tensor size in elements (prevent OOM)
    pub max_tensor_elements: usize,
    /// Maximum model inference timeout in milliseconds
    pub max_inference_timeout_ms: u64,
    /// Maximum `GPU` memory per operation in bytes
    pub max_gpu_memory_bytes: usize,
    /// Drift detection sensitivity (0.0 = disabled, 1.0 = highest)
    pub drift_sensitivity: f64,
    /// Financial validation precision (decimal places)
    pub financial_precision: u8,
    /// Enable NaN/Infinity checks for all operations
    pub nan_infinity_checks: bool,
    /// Maximum allowed model prediction value
    pub max_prediction_value: f64,
    /// Minimum allowed model prediction value
    pub min_prediction_value: f64,
    /// Enable bounds checking for all array/tensor operations
    pub bounds_checking: bool,
    /// Enable automatic fallback mechanisms
    pub auto_fallback: bool,
    /// Maximum number of retries for failed operations
    pub max_retries: u32,
}

impl Default for MLSafetyConfig {
    fn default() -> Self {
        Self {
            safety_enabled: true,
            max_tensor_elements: 100_000_000, // 100M elements
            max_inference_timeout_ms: 5000,   // 5 seconds
            max_gpu_memory_bytes: 8 * 1024 * 1024 * 1024, // 8GB
            drift_sensitivity: 0.7,
            financial_precision: 6,
            nan_infinity_checks: true,
            max_prediction_value: 1e6,
            min_prediction_value: -1e6,
            bounds_checking: true,
            auto_fallback: true,
            max_retries: 3,
        }
    }
}

/// Safety errors for ML operations
#[derive(Error, Debug)]
pub enum MLSafetyError {
    #[error("Mathematical safety violation: {reason}")]
    MathSafety { reason: String },

    #[error("Tensor safety violation: {reason}")]
    TensorSafety { reason: String },

    #[error("Financial validation failed: {reason}")]
    FinancialValidation { reason: String },

    #[error("Bounds check failed: index {index} >= length {length}")]
    BoundsCheck { index: usize, length: usize },

    #[error("Memory safety violation: {reason}")]
    MemorySafety { reason: String },

    #[error("Timeout exceeded: {timeout_ms}ms")]
    Timeout { timeout_ms: u64 },

    #[error("Model drift detected: {drift_score:.3} > threshold {threshold:.3}")]
    ModelDrift { drift_score: f64, threshold: f64 },

    #[error("GPU operation failed: {reason}")]
    GPUFailure { reason: String },

    #[error("NaN or Infinity detected in operation: {operation}")]
    InvalidFloat { operation: String },

    #[error("Prediction value out of bounds: {value} not in [{min}, {max}]")]
    PredictionOutOfBounds { value: f64, min: f64, max: f64 },

    #[error("Resource unavailable: {resource}")]
    ResourceUnavailable { resource: String },

    #[error("Compute framework error: {0}")]
    ComputeError(String),

    #[error("System resource exhausted: {resource}")]
    ResourceExhausted { resource: String },

    #[error("Validation error: {message}")]
    ValidationError { message: String },
}

// NOTE: From<ProductionTrainingError> impl is in the `ml` crate (safety_conversions module)
// because training_pipeline stays in ml and ml-core cannot depend on ml.

pub type SafetyResult<T> = Result<T, MLSafetyError>;

/// Safety status for operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SafetyStatus {
    Safe,
    Warning { reason: String },
    Danger { reason: String },
    Critical { reason: String },
}

/// Comprehensive ML Safety Manager
///
/// This is the central safety coordinator for all ML operations.
/// ALL ML operations MUST go through this manager.
#[derive(Debug)]
pub struct MLSafetyManager {
    config: MLSafetyConfig,
    math_ops: SafeMathOps,
    tensor_ops: SafeTensorOps,
    bounds_checker: BoundsChecker,
    memory_manager: Arc<RwLock<SafeMemoryManager>>,
    drift_detector: Arc<RwLock<ModelDriftDetector>>,
    financial_validator: FinancialValidator,
    timeout_manager: TimeoutManager,
    operation_history: Arc<RwLock<HashMap<String, Vec<Instant>>>>,
}

impl MLSafetyManager {
    /// Create a new ML safety manager with configuration
    pub fn new(config: MLSafetyConfig) -> Self {
        Self {
            math_ops: SafeMathOps::new(&config),
            tensor_ops: SafeTensorOps::new(&config),
            bounds_checker: BoundsChecker::new(&config),
            memory_manager: Arc::new(RwLock::new(SafeMemoryManager::new(&config))),
            drift_detector: Arc::new(RwLock::new(ModelDriftDetector::new(&config))),
            financial_validator: FinancialValidator::new(&config),
            timeout_manager: TimeoutManager::new(&config),
            operation_history: Arc::new(RwLock::new(HashMap::new())),
            config,
        }
    }

    /// Validate and execute a mathematical operation safely
    pub async fn safe_math_operation<F, T>(
        &self,
        operation_name: &str,
        operation: F,
    ) -> SafetyResult<T>
    where
        F: FnOnce() -> SafetyResult<T> + Send + 'static,
        T: Send + 'static,
    {
        if !self.config.safety_enabled {
            return operation();
        }

        // Record operation start time
        self.record_operation_start(operation_name).await;

        // Execute with timeout
        let result = self
            .timeout_manager
            .execute_with_timeout(
                operation_name,
                operation,
                Duration::from_millis(self.config.max_inference_timeout_ms),
            )
            .await?;

        // Validate result
        self.validate_operation_result(operation_name, &result)
            .await?;

        Ok(result)
    }

    /// Validate data that would go into a tensor (CPU-side check).
    ///
    /// Checks size limits, NaN/Inf, and shape consistency.
    /// Tensor creation itself is now done via `GpuTensor::from_host`.
    pub async fn validate_tensor_data(
        &self,
        data: &[f64],
        shape: &[usize],
        operation_context: &str,
    ) -> SafetyResult<()> {
        if !self.config.safety_enabled {
            return Ok(());
        }

        let total_elements: usize = shape.iter().product();
        if total_elements > self.config.max_tensor_elements {
            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Tensor too large: {} elements > {} limit in {}",
                    total_elements, self.config.max_tensor_elements, operation_context
                ),
            });
        }

        if data.len() != total_elements {
            return Err(MLSafetyError::TensorSafety {
                reason: format!(
                    "Data length {} doesn't match shape size {} in {}",
                    data.len(),
                    total_elements,
                    operation_context
                ),
            });
        }

        if self.config.nan_infinity_checks {
            for (i, &value) in data.iter().enumerate() {
                if !value.is_finite() {
                    return Err(MLSafetyError::InvalidFloat {
                        operation: format!(
                            "{}: Non-finite value {} at index {}",
                            operation_context, value, i
                        ),
                    });
                }
            }
        }

        Ok(())
    }

    /// Validate financial values and convert to safe types
    pub async fn validate_financial_prediction(
        &self,
        prediction: f64,
        context: &str,
    ) -> SafetyResult<Price> {
        if !self.config.safety_enabled {
            return Price::from_f64(prediction).map_err(|e| MLSafetyError::FinancialValidation {
                reason: format!(
                    "Failed to convert prediction to Price in {} (safety disabled): {} - {}",
                    context, prediction, e
                ),
            });
        }

        // Check for NaN/Infinity
        if !prediction.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Financial prediction in {}: {}", context, prediction),
            });
        }

        // Check prediction bounds
        if prediction < self.config.min_prediction_value
            || prediction > self.config.max_prediction_value
        {
            return Err(MLSafetyError::PredictionOutOfBounds {
                value: prediction,
                min: self.config.min_prediction_value,
                max: self.config.max_prediction_value,
            });
        }

        // Validate through financial validator
        self.financial_validator
            .validate_price(prediction, context)
            .await?;

        // Convert to safe financial type with proper error handling
        Price::from_f64(prediction).map_err(|e| MLSafetyError::FinancialValidation {
            reason: format!(
                "Failed to convert validated prediction to Price in {}: {} - {}",
                context, prediction, e
            ),
        })
    }

    /// Validate and convert price with currency support
    pub async fn validate_and_convert_price(
        &self,
        prediction: f64,
        currency: &str,
    ) -> SafetyResult<Price> {
        if !self.config.safety_enabled {
            return Price::from_f64(prediction).map_err(|e| MLSafetyError::FinancialValidation {
                reason: format!(
                    "Failed to convert prediction to Price for {} (safety disabled): {} - {}",
                    currency, prediction, e
                ),
            });
        }

        // Check for NaN/Infinity
        if !prediction.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Price prediction in {}: {}", currency, prediction),
            });
        }

        // Check prediction bounds
        if prediction < self.config.min_prediction_value
            || prediction > self.config.max_prediction_value
        {
            return Err(MLSafetyError::PredictionOutOfBounds {
                value: prediction,
                min: self.config.min_prediction_value,
                max: self.config.max_prediction_value,
            });
        }

        // Validate through financial validator with currency context
        let context = format!("price_prediction_{}_{}", currency, prediction);
        self.financial_validator
            .validate_price(prediction, &context)
            .await?;

        // Convert to safe financial type with proper error handling
        Price::from_f64(prediction).map_err(|e| MLSafetyError::FinancialValidation {
            reason: format!(
                "Failed to convert validated prediction to Price for {}: {} - {}",
                currency, prediction, e
            ),
        })
    }

    /// Validate financial value for safety
    pub fn validate_financial_value(&self, value: f64) -> SafetyResult<()> {
        if !self.config.safety_enabled {
            return Ok(());
        }

        // Check for NaN/Infinity
        if !value.is_finite() {
            return Err(MLSafetyError::InvalidFloat {
                operation: format!("Financial value validation: {}", value),
            });
        }

        // Check prediction bounds
        if value < self.config.min_prediction_value || value > self.config.max_prediction_value {
            return Err(MLSafetyError::PredictionOutOfBounds {
                value,
                min: self.config.min_prediction_value,
                max: self.config.max_prediction_value,
            });
        }

        Ok(())
    }

    /// Check for model drift and update detection
    pub async fn check_model_drift(
        &self,
        model_id: &str,
        predictions: &[f64],
        actual_values: Option<&[f64]>,
    ) -> SafetyResult<SafetyStatus> {
        if !self.config.safety_enabled {
            return Ok(SafetyStatus::Safe);
        }

        let mut drift_detector = self.drift_detector.write().await;
        let drift_score = drift_detector
            .update_and_check(model_id, predictions, actual_values)
            .await?;

        if drift_score > self.config.drift_sensitivity {
            warn!(
                "Model drift detected for {}: score {:.3} > threshold {:.3}",
                model_id, drift_score, self.config.drift_sensitivity
            );

            return Ok(SafetyStatus::Danger {
                reason: format!(
                    "Model drift: score {:.3} > threshold {:.3}",
                    drift_score, self.config.drift_sensitivity
                ),
            });
        }

        if drift_score > self.config.drift_sensitivity * 0.7 {
            return Ok(SafetyStatus::Warning {
                reason: format!("Model drift warning: score {:.3}", drift_score),
            });
        }

        Ok(SafetyStatus::Safe)
    }

    /// Get comprehensive safety status
    pub async fn get_safety_status(&self) -> HashMap<String, SafetyStatus> {
        let mut status = HashMap::new();

        // Memory status
        let memory_manager = self.memory_manager.read().await;
        status.insert("memory".to_owned(), memory_manager.get_status().await);
        drop(memory_manager);

        // Drift detection status
        let drift_detector = self.drift_detector.read().await;
        status.insert(
            "drift_detection".to_owned(),
            drift_detector.get_status().await,
        );
        drop(drift_detector);

        // Operation history status
        let history = self.operation_history.read().await;
        let recent_operations = history.values().map(|ops| ops.len()).sum::<usize>();

        let ops_status = if recent_operations > 10000 {
            SafetyStatus::Warning {
                reason: format!("High operation count: {}", recent_operations),
            }
        } else {
            SafetyStatus::Safe
        };
        status.insert("operations".to_owned(), ops_status);

        status
    }

    /// Emergency shutdown all ML operations
    pub async fn emergency_shutdown(&self, reason: &str) -> SafetyResult<()> {
        error!("ML Safety Manager emergency shutdown: {}", reason);

        // Stop all ongoing operations
        self.timeout_manager.shutdown_all().await;

        // Clear all caches and memory
        let mut memory_manager = self.memory_manager.write().await;
        memory_manager.emergency_cleanup().await?;

        // Reset drift detection
        let mut drift_detector = self.drift_detector.write().await;
        drift_detector.reset_all().await;

        // Clear operation history
        let mut history = self.operation_history.write().await;
        history.clear();

        info!("ML Safety Manager emergency shutdown completed");
        Ok(())
    }

    /// Record operation start for monitoring
    async fn record_operation_start(&self, operation_name: &str) {
        let mut history = self.operation_history.write().await;
        let operations = history
            .entry(operation_name.to_string())
            .or_insert_with(Vec::new);
        operations.push(Instant::now());

        // Keep only recent operations
        let cutoff = Instant::now().checked_sub(Duration::from_secs(300))
            .unwrap_or_else(Instant::now); // 5 minutes
        operations.retain(|&time| time > cutoff);
    }

    /// Validate operation result
    async fn validate_operation_result<T>(
        &self,
        operation_name: &str,
        _result: &T,
    ) -> SafetyResult<()> {
        debug!("Operation {} completed successfully", operation_name);
        Ok(())
    }
}

/// Global ML safety manager instance
static GLOBAL_SAFETY_MANAGER: once_cell::sync::Lazy<MLSafetyManager> =
    once_cell::sync::Lazy::new(|| MLSafetyManager::new(MLSafetyConfig::default()));

/// Get the global ML safety manager
pub fn get_global_safety_manager() -> &'static MLSafetyManager {
    &GLOBAL_SAFETY_MANAGER
}

/// Initialize ML safety with custom configuration
pub fn initialize_ml_safety(_config: MLSafetyConfig) -> &'static MLSafetyManager {
    // Note: This is a simplified version. In production, we would use
    // proper initialization that allows configuration override
    &GLOBAL_SAFETY_MANAGER
}

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

    #[tokio::test]
    async fn test_validate_tensor_data() {
        let manager = MLSafetyManager::new(MLSafetyConfig::default());

        // Valid data
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let shape = &[2, 2];
        let result = manager
            .validate_tensor_data(&data, shape, "test_creation")
            .await;
        assert!(result.is_ok());

        // Invalid data - NaN
        let bad_data = vec![1.0, f64::NAN, 3.0, 4.0];
        let bad_result = manager
            .validate_tensor_data(&bad_data, shape, "test_nan")
            .await;
        assert!(bad_result.is_err());
    }

    #[tokio::test]
    async fn test_financial_validation() {
        let manager = MLSafetyManager::new(MLSafetyConfig::default());

        // Valid prediction
        let valid_pred = manager
            .validate_financial_prediction(123.45, "test_valid")
            .await;
        assert!(valid_pred.is_ok());

        // Invalid prediction - NaN
        let invalid_pred = manager
            .validate_financial_prediction(f64::NAN, "test_nan")
            .await;
        assert!(invalid_pred.is_err());

        // Invalid prediction - out of bounds
        let oob_pred = manager
            .validate_financial_prediction(1e10, "test_oob")
            .await;
        assert!(oob_pred.is_err());
    }

    #[tokio::test]
    async fn test_safety_status() {
        let manager = MLSafetyManager::new(MLSafetyConfig::default());
        let status = manager.get_safety_status().await;

        assert!(status.contains_key("memory"));
        assert!(status.contains_key("drift_detection"));
        assert!(status.contains_key("operations"));
    }
}
