trustformers-models 0.1.1

//! System memory collection utilities
//!
//! This module provides platform-specific utilities for collecting
//! memory information from the system and current process.

use anyhow::Result;
use std::sync::atomic::{AtomicU64, Ordering};

use super::types::*;

/// System memory information
#[derive(Debug, Clone)]
pub struct GpuMemoryInfo {
    pub total_mb: f64,
    pub used_mb: f64,
    pub free_mb: f64,
}

/// Process memory information for internal tracking
#[derive(Debug, Clone)]
pub struct ProcessMemoryInfo {
    pub total_mb: f64,
    pub heap_mb: f64,
    pub stack_mb: f64,
    pub peak_mb: f64,
    pub allocated_objects: u64,
    pub deallocated_objects: u64,
    pub active_allocations: u64,
    pub gc_collections: u64,
    pub gc_time_ms: f64,
}

impl super::profiler::MemoryProfiler {
    /// Collect current memory metrics
    pub async fn collect_memory_metrics() -> Result<MemoryMetrics> {
        // Use system APIs to collect memory information
        let _memory_info = Self::get_system_memory_info()?;
        let process_info = Self::get_process_memory_info()?;
        let gpu_info = Self::get_gpu_memory_info().await.ok();

        Ok(MemoryMetrics {
            timestamp: std::time::SystemTime::now(),
            total_memory_mb: process_info.total_mb,
            heap_memory_mb: process_info.heap_mb,
            stack_memory_mb: process_info.stack_mb,
            gpu_memory_mb: gpu_info.map(|info| info.used_mb),
            peak_memory_mb: process_info.peak_mb,
            allocated_objects: process_info.allocated_objects,
            deallocated_objects: process_info.deallocated_objects,
            active_allocations: process_info.active_allocations,
            memory_fragmentation_ratio: Self::calculate_fragmentation_ratio(&process_info),
            gc_collections: process_info.gc_collections,
            gc_time_ms: process_info.gc_time_ms,
            memory_growth_rate_mb_per_sec: 0.0, // Will be calculated later
        })
    }

    /// Get system memory information
    pub fn get_system_memory_info() -> Result<SystemMemoryInfo> {
        // This would use platform-specific APIs in a real implementation
        // For now, return mock data that represents realistic system values
        Ok(SystemMemoryInfo {
            total_memory: 16 * 1024 * 1024 * 1024,    // 16GB
            available_memory: 8 * 1024 * 1024 * 1024, // 8GB
            used_memory: 8 * 1024 * 1024 * 1024,      // 8GB
            free_memory: 4 * 1024 * 1024 * 1024,      // 4GB
            cached_memory: 2 * 1024 * 1024 * 1024,    // 2GB
            buffer_memory: 2 * 1024 * 1024 * 1024,    // 2GB
        })
    }

    /// Get process memory information
    pub fn get_process_memory_info() -> Result<ProcessMemoryInfo> {
        // This would use platform-specific APIs to get real memory info
        // For now, return mock data that simulates realistic values
        static COUNTER: AtomicU64 = AtomicU64::new(0);
        let count = COUNTER.fetch_add(1, Ordering::Relaxed);

        // Simulate some realistic memory growth patterns
        let base_memory = 512.0;
        let growth = (count as f64 * 0.1).min(500.0);
        let noise = (count as f64 * 0.05).sin() * 10.0; // Add some realistic noise

        Ok(ProcessMemoryInfo {
            total_mb: base_memory + growth + noise,
            heap_mb: (base_memory + growth + noise) * 0.8,
            stack_mb: 64.0 + (count as f64 * 0.01).min(32.0),
            peak_mb: base_memory + growth + noise + 100.0,
            allocated_objects: 1000000 + count * 1000,
            deallocated_objects: 950000 + count * 950,
            active_allocations: 50000 + count * 50,
            gc_collections: 10 + count / 100,
            gc_time_ms: 150.0 + (count as f64 * 0.1),
        })
    }

    /// Get GPU memory information
    pub async fn get_gpu_memory_info() -> Result<GpuMemoryInfo> {
        // This would use CUDA/ROCm APIs in a real implementation
        // For now, return mock data
        static GPU_COUNTER: AtomicU64 = AtomicU64::new(0);
        let count = GPU_COUNTER.fetch_add(1, Ordering::Relaxed);

        let base_usage = 2048.0;
        let usage_growth = (count as f64 * 0.5).min(2048.0);

        Ok(GpuMemoryInfo {
            total_mb: 8192.0,
            used_mb: base_usage + usage_growth,
            free_mb: 8192.0 - (base_usage + usage_growth),
        })
    }

    /// Calculate memory fragmentation ratio
    pub fn calculate_fragmentation_ratio(info: &ProcessMemoryInfo) -> f64 {
        // Simple fragmentation estimation based on allocation patterns
        // In reality, this would be more sophisticated and platform-specific
        let theoretical_optimal = info.active_allocations as f64 * 64.0 / 1024.0 / 1024.0; // Assume avg 64 bytes per allocation
        if theoretical_optimal > 0.0 {
            ((info.total_mb - theoretical_optimal) / info.total_mb).clamp(0.0, 1.0)
        } else {
            0.0
        }
    }

    /// Update adaptive thresholds based on current metrics
    pub async fn update_adaptive_thresholds(
        current_metrics: &MemoryMetrics,
        adaptive_thresholds: &std::sync::Arc<
            std::sync::Mutex<super::analytics::AdaptiveThresholds>,
        >,
    ) {
        // Update thresholds every 5 minutes
        let should_update = {
            let thresholds = match adaptive_thresholds.lock() {
                Ok(guard) => guard,
                Err(_) => return, // Skip update on lock failure
            };
            current_metrics
                .timestamp
                .duration_since(thresholds.last_updated)
                .unwrap_or_default()
                .as_secs()
                > 300
        };

        if should_update {
            let mut thresholds = match adaptive_thresholds.lock() {
                Ok(guard) => guard,
                Err(_) => return, // Skip update on lock failure
            };
            // Simple adaptive update based on current usage
            let adaptation_factor = 0.1;
            thresholds.base_memory_threshold = thresholds.base_memory_threshold
                * (1.0 - adaptation_factor)
                + current_metrics.total_memory_mb * 1.2 * adaptation_factor;
            thresholds.growth_rate_threshold = thresholds.growth_rate_threshold
                * (1.0 - adaptation_factor)
                + current_metrics.memory_growth_rate_mb_per_sec.abs() * 2.0 * adaptation_factor;
            thresholds.fragmentation_threshold = thresholds.fragmentation_threshold
                * (1.0 - adaptation_factor)
                + current_metrics.memory_fragmentation_ratio * 1.5 * adaptation_factor;
            thresholds.last_updated = current_metrics.timestamp;
        }
    }

    /// Detect memory leaks using heuristics
    pub async fn detect_memory_leaks(
        current_metrics: &MemoryMetrics,
        previous_metrics: &Option<MemoryMetrics>,
        leak_detection: &std::sync::Arc<
            std::sync::Mutex<super::analytics::LeakDetectionHeuristics>,
        >,
        alerts: &std::sync::Arc<std::sync::Mutex<Vec<MemoryAlert>>>,
        cached_recommendations: &super::analytics::AlertRecommendations,
    ) {
        if let Some(_prev) = previous_metrics {
            let growth_rate = current_metrics.memory_growth_rate_mb_per_sec;
            let detection = match leak_detection.lock() {
                Ok(guard) => guard,
                Err(_) => return, // Skip detection on lock failure
            };

            // Check for sustained growth
            if growth_rate > detection.sustained_growth_threshold {
                let mut alerts_vec = match alerts.lock() {
                    Ok(guard) => guard,
                    Err(_) => return, // Skip alert on lock failure
                };
                alerts_vec.push(MemoryAlert {
                    id: uuid::Uuid::new_v4(),
                    timestamp: current_metrics.timestamp,
                    alert_type: MemoryAlertType::MemoryLeak,
                    severity: if growth_rate > detection.sustained_growth_threshold * 2.0 {
                        AlertSeverity::Critical
                    } else {
                        AlertSeverity::Warning
                    },
                    message: format!(
                        "Potential memory leak detected: {:.2} MB/sec growth",
                        growth_rate
                    ),
                    details: std::collections::HashMap::new(),
                    recommendations: cached_recommendations.memory_leak.clone(),
                });
            }
        }
    }

    /// Update memory prediction
    pub async fn update_memory_prediction(
        _current_metrics: &MemoryMetrics,
        memory_predictor: &std::sync::Arc<std::sync::Mutex<super::analytics::MemoryPredictor>>,
        metrics_history: &std::sync::Arc<
            std::sync::Mutex<std::collections::VecDeque<MemoryMetrics>>,
        >,
    ) {
        let history = match metrics_history.lock() {
            Ok(guard) => guard,
            Err(_) => return, // Skip prediction on lock failure
        };
        if history.len() >= 10 {
            let recent_metrics: Vec<MemoryMetrics> = history.iter().cloned().collect();
            drop(history);

            let mut predictor = match memory_predictor.lock() {
                Ok(guard) => guard,
                Err(_) => return, // Skip prediction on lock failure
            };
            let _prediction = predictor.predict_memory_usage(&recent_metrics, Some(300));
            // 5 minutes
            // Prediction result would be stored or used for alerts in a real implementation
        }
    }

    /// Analyze for alerts with adaptive thresholds
    pub async fn analyze_for_alerts_adaptive(
        current: &MemoryMetrics,
        _previous: &Option<MemoryMetrics>,
        alerts: &std::sync::Arc<std::sync::Mutex<Vec<MemoryAlert>>>,
        _config: &super::types::ProfilerConfig,
        cached_recommendations: &super::analytics::AlertRecommendations,
        adaptive_thresholds: &std::sync::Arc<
            std::sync::Mutex<super::analytics::AdaptiveThresholds>,
        >,
    ) {
        let thresholds = match adaptive_thresholds.lock() {
            Ok(guard) => guard,
            Err(_) => return, // Skip analysis on lock failure
        };
        let mut new_alerts = Vec::new();

        // High memory usage alert using adaptive threshold
        if current.total_memory_mb > thresholds.base_memory_threshold {
            new_alerts.push(MemoryAlert {
                id: uuid::Uuid::new_v4(),
                timestamp: current.timestamp,
                alert_type: MemoryAlertType::HighMemoryUsage,
                severity: if current.total_memory_mb > thresholds.base_memory_threshold * 1.5 {
                    AlertSeverity::Critical
                } else {
                    AlertSeverity::Warning
                },
                message: format!(
                    "Memory usage is {:.2}MB, exceeding adaptive threshold of {:.2}MB",
                    current.total_memory_mb, thresholds.base_memory_threshold
                ),
                details: std::collections::HashMap::new(),
                recommendations: cached_recommendations.high_memory.clone(),
            });
        }

        // Rapid growth alert
        if current.memory_growth_rate_mb_per_sec > thresholds.growth_rate_threshold {
            new_alerts.push(MemoryAlert {
                id: uuid::Uuid::new_v4(),
                timestamp: current.timestamp,
                alert_type: MemoryAlertType::RapidGrowth,
                severity: AlertSeverity::Warning,
                message: format!(
                    "Rapid memory growth detected: {:.2}MB/sec",
                    current.memory_growth_rate_mb_per_sec
                ),
                details: std::collections::HashMap::new(),
                recommendations: cached_recommendations.rapid_growth.clone(),
            });
        }

        // Fragmentation alert
        if current.memory_fragmentation_ratio > thresholds.fragmentation_threshold {
            new_alerts.push(MemoryAlert {
                id: uuid::Uuid::new_v4(),
                timestamp: current.timestamp,
                alert_type: MemoryAlertType::FragmentationHigh,
                severity: AlertSeverity::Info,
                message: format!(
                    "High memory fragmentation: {:.2}%",
                    current.memory_fragmentation_ratio * 100.0
                ),
                details: std::collections::HashMap::new(),
                recommendations: cached_recommendations.fragmentation.clone(),
            });
        }

        // Store alerts if any were generated
        if !new_alerts.is_empty() {
            let mut alerts_vec = match alerts.lock() {
                Ok(guard) => guard,
                Err(_) => {
                    drop(thresholds);
                    return; // Skip storing alerts on lock failure
                },
            };
            alerts_vec.extend(new_alerts);

            // Keep only recent alerts (last 1000)
            if alerts_vec.len() > 1000 {
                let excess = alerts_vec.len() - 1000;
                alerts_vec.drain(0..excess);
            }
        }
        drop(thresholds);
    }

    /// Update memory patterns
    pub async fn update_patterns(
        current_metrics: &MemoryMetrics,
        patterns: &std::sync::Arc<std::sync::Mutex<Vec<MemoryPattern>>>,
    ) {
        let mut patterns_vec = match patterns.lock() {
            Ok(guard) => guard,
            Err(_) => return, // Skip pattern update on lock failure
        };

        // Simple pattern detection - in practice this would be much more sophisticated
        if current_metrics.total_memory_mb > 1000.0 {
            // Check if we already have a large allocation pattern
            let has_large_pattern = patterns_vec
                .iter()
                .any(|p| matches!(p.pattern_type, PatternType::LargeAllocations));

            if !has_large_pattern {
                patterns_vec.push(MemoryPattern {
                    pattern_type: PatternType::LargeAllocations,
                    frequency: 1,
                    average_size: current_metrics.total_memory_mb,
                    total_size: current_metrics.total_memory_mb,
                    locations: vec!["system".to_string()],
                    trend: PatternTrend::Increasing,
                });
            } else {
                // Update existing pattern
                if let Some(pattern) = patterns_vec
                    .iter_mut()
                    .find(|p| matches!(p.pattern_type, PatternType::LargeAllocations))
                {
                    pattern.frequency += 1;
                    pattern.average_size =
                        (pattern.average_size + current_metrics.total_memory_mb) / 2.0;
                    pattern.total_size += current_metrics.total_memory_mb;
                }
            }
        }

        // Keep patterns list bounded
        if patterns_vec.len() > 100 {
            let excess = patterns_vec.len() - 100;
            patterns_vec.drain(0..excess);
        }
    }
}