foxhunt/ml/src/operations.rs

//! Safe operations module for ML models
//!
//! This module provides safety wrappers for all ML operations to ensure
//! production-grade reliability and error handling.

use crate::{Decimal, MLError, MLResult};
use tracing::{debug, error, warn};

/// Safe ML operations manager
#[derive(Debug, Clone)]
pub struct SafeMLOperations {
    max_tensor_size: usize,
    timeout_ms: u64,
}

impl Default for SafeMLOperations {
    fn default() -> Self {
        Self {
            max_tensor_size: 1_000_000, // 1M elements max
            timeout_ms: 5000,           // 5 second timeout
        }
    }
}

impl SafeMLOperations {
    /// Create a new safe ML operations manager
    pub fn new(max_tensor_size: usize, timeout_ms: u64) -> Self {
        Self {
            max_tensor_size,
            timeout_ms,
        }
    }

    /// Safely validate tensor dimensions with comprehensive error context
    pub fn validate_tensor_dims(&self, dims: &[usize], operation: &str) -> MLResult<()> {
        let total_size = dims.iter().product::<usize>();

        // Log validation attempt for monitoring
        debug!(
            operation = operation,
            dims = ?dims,
            total_size = total_size,
            max_size = self.max_tensor_size,
            "Validating tensor dimensions"
        );

        if total_size > self.max_tensor_size {
            error!(
                operation = operation,
                dims = ?dims,
                total_size = total_size,
                max_size = self.max_tensor_size,
                "Tensor size validation failed - exceeds maximum allowed size"
            );
            return Err(MLError::ResourceLimit {
                resource: format!("tensor_size_for_{}", operation),
                limit: self.max_tensor_size,
            });
        }

        if dims.iter().any(|&d| d == 0) {
            error!(
                "Zero dimension found in tensor for operation: {}",
                operation
            );
            return Err(MLError::DimensionMismatch {
                expected: 1,
                actual: 0,
            });
        }

        debug!("Tensor dimensions validated for {}: {:?}", operation, dims);
        Ok(())
    }

    /// Safely perform mathematical operations with NaN/infinity checking
    pub fn safe_math_op<T, F>(&self, operation: &str, func: F) -> MLResult<T>
    where
        F: FnOnce() -> MLResult<T>,
    {
        debug!("Starting safe math operation: {}", operation);

        let result = func();

        match result {
            Ok(val) => {
                debug!("Safe math operation {} completed successfully", operation);
                Ok(val)
            },
            Err(e) => {
                error!("Safe math operation {} failed: {}", operation, e);
                Err(e)
            },
        }
    }

    /// Safely validate financial values
    pub fn validate_financial_value(&self, value: Decimal, field: &str) -> MLResult<()> {
        if value.is_sign_negative() && !field.contains("return") && !field.contains("diff") {
            warn!(
                "Negative value {} for field {} (may be valid for returns/diffs)",
                value, field
            );
        }

        if value.is_zero() && field.contains("price") {
            error!("Zero price value for field: {}", field);
            return Err(MLError::ValidationError {
                message: format!("Invalid zero price for field: {}", field),
            });
        }

        debug!("Financial value {} validated for field: {}", value, field);
        Ok(())
    }

    /// Safely allocate memory for ML operations
    pub fn safe_allocate<T>(&self, size: usize, operation: &str) -> MLResult<Vec<T>>
    where
        T: Default + Clone,
    {
        if size > self.max_tensor_size {
            error!(
                "Allocation size {} exceeds maximum {} for operation: {}",
                size, self.max_tensor_size, operation
            );
            return Err(MLError::ResourceLimit {
                resource: "memory_allocation".to_string(),
                limit: self.max_tensor_size,
            });
        }

        let vec = vec![T::default(); size];
        debug!(
            "Successfully allocated {} elements for operation: {}",
            size, operation
        );
        Ok(vec)
    }

    /// Get maximum tensor size
    pub fn max_tensor_size(&self) -> usize {
        self.max_tensor_size
    }

    /// Get timeout in milliseconds
    pub fn timeout_ms(&self) -> u64 {
        self.timeout_ms
    }
}

/// Global safe operations instance
static GLOBAL_SAFE_OPS: std::sync::OnceLock<SafeMLOperations> = std::sync::OnceLock::new();

/// Get the global safe operations instance
pub fn get_safe_operations() -> &'static SafeMLOperations {
    GLOBAL_SAFE_OPS.get_or_init(SafeMLOperations::default)
}

/// Initialize safe operations with custom configuration
pub fn initialize_safe_operations(config: SafeMLOperations) -> MLResult<()> {
    GLOBAL_SAFE_OPS
        .set(config)
        .map_err(|_| MLError::ConfigError {
            reason: "Safe operations already initialized".to_string(),
        })?;
    Ok(())
}

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

    #[test]
    fn test_validate_tensor_dims() {
        let ops = SafeMLOperations::default();

        // Valid dimensions
        assert!(ops.validate_tensor_dims(&[10, 20, 30], "test").is_ok());

        // Zero dimension should fail
        assert!(ops.validate_tensor_dims(&[10, 0, 30], "test").is_err());

        // Too large should fail
        assert!(ops.validate_tensor_dims(&[10000, 10000], "test").is_err());
    }

    #[test]
    fn test_safe_math_op() {
        let ops = SafeMLOperations::default();

        let result = ops.safe_math_op("add", || Ok(2 + 2));
        assert_eq!(result.unwrap(), 4);

        let error_result: MLResult<i32> = ops.safe_math_op("error", || {
            Err(MLError::ValidationError {
                message: "test error".to_string(),
            })
        });
        assert!(error_result.is_err());
    }

    #[test]
    fn test_validate_financial_value() {
        let ops = SafeMLOperations::default();

        // Valid positive price
        assert!(ops
            .validate_financial_value(Decimal::try_from(100.50).unwrap_or(Decimal::ZERO), "price")
            .is_ok());

        // Valid negative return
        assert!(ops
            .validate_financial_value(Decimal::try_from(-0.05).unwrap_or(Decimal::ZERO), "return")
            .is_ok());

        // Invalid zero price should fail
        assert!(ops
            .validate_financial_value(Decimal::ZERO, "price")
            .is_err());
    }

    #[test]
    fn test_safe_allocate() {
        let ops = SafeMLOperations::default();

        // Valid allocation
        let vec: Vec<f32> = ops.safe_allocate(100, "test").unwrap();
        assert_eq!(vec.len(), 100);

        // Too large allocation should fail
        assert!(ops.safe_allocate::<f32>(2_000_000, "test").is_err());
    }
}