//! Metrics Collection Module
//!
//! This module provides real-time metrics collection for cluster monitoring,
//! including GPU utilization, memory usage, and bandwidth monitoring.

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

use anyhow::Result;
use chrono::Utc;
use clusterviz_shared::{ClusterMetrics, ClusterTopology, NodeStatus};
use tracing::{debug, trace};
use uuid::Uuid;

/// Collects real-time metrics from the cluster
pub struct MetricsCollector {
    /// Utilization tracker
    utilization_tracker: UtilizationTracker,
    /// Memory tracker
    memory_tracker: MemoryTracker,
    /// Bandwidth monitor
    bandwidth_monitor: BandwidthMonitor,
    /// Collection interval
    collection_interval: Duration,
    /// Last collection time
    last_collection: Option<Instant>,
}

impl MetricsCollector {
    /// Create a new metrics collector
    #[must_use]
    pub fn new() -> Self {
        Self {
            utilization_tracker: UtilizationTracker::new(),
            memory_tracker: MemoryTracker::new(),
            bandwidth_monitor: BandwidthMonitor::new(),
            collection_interval: Duration::from_millis(100),
            last_collection: None,
        }
    }

    /// Set the collection interval
    #[must_use]
    pub fn with_interval(mut self, interval: Duration) -> Self {
        self.collection_interval = interval;
        self
    }

    /// Collect metrics from the topology
    pub async fn collect(&mut self, topology: &ClusterTopology) -> Result<ClusterMetrics> {
        debug!("Collecting metrics for {} nodes", topology.nodes.len());

        let mut metrics = ClusterMetrics::new();
        metrics.timestamp = Utc::now();

        // Collect GPU utilization for each node
        for node in &topology.nodes {
            let util = self.utilization_tracker.get_utilization(node.id).await?;
            metrics.gpu_utilization.push((node.id, util));
        }

        // Collect memory usage for each node
        for node in &topology.nodes {
            let mem_used = self
                .memory_tracker
                .get_memory_used(node.id, node.memory_gb)
                .await?;
            metrics.memory_used.push((node.id, mem_used));
        }

        // Collect bandwidth usage for each link
        for link in &topology.links {
            let bandwidth = self.bandwidth_monitor.get_bandwidth_usage(link.id).await?;
            metrics.bandwidth_usage.push((link.id, bandwidth));
        }

        // Calculate aggregate metrics
        metrics.total_tflops = self.calculate_total_tflops(topology);
        metrics.active_collectives = self.count_active_collectives();
        metrics.data_transfer_rate_gb = self.calculate_transfer_rate(&metrics);

        self.last_collection = Some(Instant::now());

        trace!(
            "Collected metrics: GPU util {:.1}%, Memory {:.1} GB",
            metrics.average_gpu_utilization() * 100.0,
            metrics.total_memory_used_gb()
        );

        Ok(metrics)
    }

    /// Calculate total TFLOPS across all nodes
    fn calculate_total_tflops(&self, topology: &ClusterTopology) -> f64 {
        // Estimate TFLOPS based on GPU model
        // M3 Max: ~14 TFLOPS, M4 Ultra: ~28 TFLOPS (estimated)
        topology
            .nodes
            .iter()
            .filter(|n| n.status == NodeStatus::Healthy)
            .map(|n| {
                if n.gpu_model.contains("Ultra") {
                    28.0
                } else if n.gpu_model.contains("Max") {
                    14.0
                } else if n.gpu_model.contains("Pro") {
                    8.0
                } else {
                    4.0
                }
            })
            .sum()
    }

    /// Count active collective operations (simulated)
    fn count_active_collectives(&self) -> u32 {
        // In a real implementation, this would track actual collectives
        // For simulation, return a random-ish value
        (rand::random::<u32>() % 5) + 1
    }

    /// Calculate data transfer rate from bandwidth metrics
    fn calculate_transfer_rate(&self, metrics: &ClusterMetrics) -> f64 {
        let total_bandwidth: f64 = metrics.bandwidth_usage.iter().map(|(_, bw)| bw).sum();
        // Convert Gbps to GB/s (divide by 8)
        total_bandwidth / 8.0
    }

    /// Get the last collection time
    #[must_use]
    pub fn last_collection_time(&self) -> Option<Instant> {
        self.last_collection
    }
}

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

/// Tracks GPU utilization across nodes
pub struct UtilizationTracker {
    /// Cached utilization values
    cache: HashMap<Uuid, f64>,
    /// Utilization history for trend analysis
    history: HashMap<Uuid, Vec<f64>>,
    /// Maximum history entries per node
    max_history: usize,
    /// Smoothing factor for exponential moving average
    smoothing_factor: f64,
}

impl UtilizationTracker {
    /// Create a new utilization tracker
    #[must_use]
    pub fn new() -> Self {
        Self {
            cache: HashMap::new(),
            history: HashMap::new(),
            max_history: 100,
            smoothing_factor: 0.2,
        }
    }

    /// Set the maximum history size
    #[must_use]
    pub fn with_max_history(mut self, max: usize) -> Self {
        self.max_history = max;
        self
    }

    /// Set the smoothing factor
    #[must_use]
    pub fn with_smoothing(mut self, factor: f64) -> Self {
        self.smoothing_factor = factor.clamp(0.0, 1.0);
        self
    }

    /// Get current utilization for a node
    pub async fn get_utilization(&mut self, node_id: Uuid) -> Result<f64> {
        // Simulate utilization reading
        // In real implementation, would query the GPU driver
        let raw_util = self.read_raw_utilization(node_id).await?;

        // Apply exponential moving average smoothing
        let smoothed = if let Some(&prev) = self.cache.get(&node_id) {
            self.smoothing_factor * raw_util + (1.0 - self.smoothing_factor) * prev
        } else {
            raw_util
        };

        // Update cache and history
        self.cache.insert(node_id, smoothed);
        let history = self.history.entry(node_id).or_default();
        history.push(smoothed);
        if history.len() > self.max_history {
            history.remove(0);
        }

        Ok(smoothed)
    }

    /// Read raw utilization (simulated)
    async fn read_raw_utilization(&self, node_id: Uuid) -> Result<f64> {
        // Simulate GPU utilization based on node ID for consistency
        // In real implementation, would query IOKit or Metal framework
        let base = (node_id.as_bytes()[0] as f64 / 255.0) * 0.3 + 0.5;
        let noise = (rand::random::<f64>() - 0.5) * 0.1;
        Ok((base + noise).clamp(0.0, 1.0))
    }

    /// Get utilization trend (positive = increasing, negative = decreasing)
    #[must_use]
    pub fn get_trend(&self, node_id: Uuid) -> Option<f64> {
        let history = self.history.get(&node_id)?;
        if history.len() < 2 {
            return None;
        }

        let recent = &history[history.len().saturating_sub(10)..];
        if recent.len() < 2 {
            return None;
        }

        let first_half: f64 =
            recent[..recent.len() / 2].iter().sum::<f64>() / (recent.len() / 2) as f64;
        let second_half: f64 =
            recent[recent.len() / 2..].iter().sum::<f64>() / ((recent.len() + 1) / 2) as f64;

        Some(second_half - first_half)
    }

    /// Get average utilization over history
    #[must_use]
    pub fn get_average(&self, node_id: Uuid) -> Option<f64> {
        let history = self.history.get(&node_id)?;
        if history.is_empty() {
            return None;
        }
        Some(history.iter().sum::<f64>() / history.len() as f64)
    }

    /// Clear cached data for a node
    pub fn clear_node(&mut self, node_id: Uuid) {
        self.cache.remove(&node_id);
        self.history.remove(&node_id);
    }

    /// Clear all cached data
    pub fn clear_all(&mut self) {
        self.cache.clear();
        self.history.clear();
    }
}

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

/// Tracks memory usage across nodes
pub struct MemoryTracker {
    /// Cached memory usage values (in GB)
    cache: HashMap<Uuid, f64>,
    /// Memory usage history
    history: HashMap<Uuid, Vec<f64>>,
    /// Maximum history entries
    max_history: usize,
}

impl MemoryTracker {
    /// Create a new memory tracker
    #[must_use]
    pub fn new() -> Self {
        Self {
            cache: HashMap::new(),
            history: HashMap::new(),
            max_history: 100,
        }
    }

    /// Set the maximum history size
    #[must_use]
    pub fn with_max_history(mut self, max: usize) -> Self {
        self.max_history = max;
        self
    }

    /// Get current memory used for a node
    pub async fn get_memory_used(&mut self, node_id: Uuid, total_gb: u32) -> Result<f64> {
        // Simulate memory reading
        let raw_usage = self.read_raw_memory(node_id, total_gb).await?;

        // Update cache and history
        self.cache.insert(node_id, raw_usage);
        let history = self.history.entry(node_id).or_default();
        history.push(raw_usage);
        if history.len() > self.max_history {
            history.remove(0);
        }

        Ok(raw_usage)
    }

    /// Read raw memory usage (simulated)
    async fn read_raw_memory(&self, node_id: Uuid, total_gb: u32) -> Result<f64> {
        // Simulate memory usage based on node ID
        let base_ratio = (node_id.as_bytes()[1] as f64 / 255.0) * 0.3 + 0.4;
        let noise = (rand::random::<f64>() - 0.5) * 0.05;
        let ratio = (base_ratio + noise).clamp(0.0, 1.0);
        Ok(ratio * total_gb as f64)
    }

    /// Get peak memory usage from history
    #[must_use]
    pub fn get_peak(&self, node_id: Uuid) -> Option<f64> {
        self.history.get(&node_id)?.iter().copied().reduce(f64::max)
    }

    /// Get memory usage trend
    #[must_use]
    pub fn get_trend(&self, node_id: Uuid) -> Option<f64> {
        let history = self.history.get(&node_id)?;
        if history.len() < 2 {
            return None;
        }

        let recent = &history[history.len().saturating_sub(10)..];
        if recent.len() < 2 {
            return None;
        }

        let first = recent.first()?;
        let last = recent.last()?;
        Some(last - first)
    }

    /// Clear cached data
    pub fn clear_all(&mut self) {
        self.cache.clear();
        self.history.clear();
    }
}

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

/// Monitors bandwidth usage across links
pub struct BandwidthMonitor {
    /// Current bandwidth readings per link (Gbps)
    current_bandwidth: HashMap<Uuid, f64>,
    /// Bandwidth history per link
    history: HashMap<Uuid, Vec<f64>>,
    /// Maximum history entries
    max_history: usize,
    /// Moving average window size
    window_size: usize,
}

impl BandwidthMonitor {
    /// Create a new bandwidth monitor
    #[must_use]
    pub fn new() -> Self {
        Self {
            current_bandwidth: HashMap::new(),
            history: HashMap::new(),
            max_history: 100,
            window_size: 10,
        }
    }

    /// Set the maximum history size
    #[must_use]
    pub fn with_max_history(mut self, max: usize) -> Self {
        self.max_history = max;
        self
    }

    /// Set the moving average window size
    #[must_use]
    pub fn with_window_size(mut self, size: usize) -> Self {
        self.window_size = size.max(1);
        self
    }

    /// Get current bandwidth usage for a link
    pub async fn get_bandwidth_usage(&mut self, link_id: Uuid) -> Result<f64> {
        let raw_bandwidth = self.read_raw_bandwidth(link_id).await?;

        // Apply moving average
        let history = self.history.entry(link_id).or_default();
        history.push(raw_bandwidth);
        if history.len() > self.max_history {
            history.remove(0);
        }

        let window = &history[history.len().saturating_sub(self.window_size)..];
        let smoothed = window.iter().sum::<f64>() / window.len() as f64;

        self.current_bandwidth.insert(link_id, smoothed);
        Ok(smoothed)
    }

    /// Read raw bandwidth (simulated)
    async fn read_raw_bandwidth(&self, link_id: Uuid) -> Result<f64> {
        // Simulate bandwidth usage
        // Thunderbolt 5 max: 120 Gbps, simulate 40-80% utilization
        let base = (link_id.as_bytes()[0] as f64 / 255.0) * 40.0 + 48.0;
        let noise = (rand::random::<f64>() - 0.5) * 10.0;
        Ok((base + noise).clamp(0.0, 120.0))
    }

    /// Get peak bandwidth for a link
    #[must_use]
    pub fn get_peak(&self, link_id: Uuid) -> Option<f64> {
        self.history.get(&link_id)?.iter().copied().reduce(f64::max)
    }

    /// Get average bandwidth for a link
    #[must_use]
    pub fn get_average(&self, link_id: Uuid) -> Option<f64> {
        let history = self.history.get(&link_id)?;
        if history.is_empty() {
            return None;
        }
        Some(history.iter().sum::<f64>() / history.len() as f64)
    }

    /// Calculate bandwidth efficiency (actual / theoretical max)
    #[must_use]
    pub fn get_efficiency(&self, link_id: Uuid, max_bandwidth: f64) -> Option<f64> {
        let current = self.current_bandwidth.get(&link_id)?;
        Some(current / max_bandwidth)
    }

    /// Get all link bandwidths
    #[must_use]
    pub fn get_all_bandwidths(&self) -> &HashMap<Uuid, f64> {
        &self.current_bandwidth
    }

    /// Clear all data
    pub fn clear_all(&mut self) {
        self.current_bandwidth.clear();
        self.history.clear();
    }
}

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

/// Aggregated cluster statistics
#[derive(Debug, Clone)]
pub struct ClusterStats {
    /// Average GPU utilization
    pub avg_gpu_util: f64,
    /// Maximum GPU utilization
    pub max_gpu_util: f64,
    /// Total memory used (GB)
    pub total_memory_used: f64,
    /// Total memory capacity (GB)
    pub total_memory_capacity: f64,
    /// Average bandwidth utilization
    pub avg_bandwidth_util: f64,
    /// Total compute capacity (TFLOPS)
    pub total_tflops: f64,
    /// Number of healthy nodes
    pub healthy_nodes: usize,
    /// Total number of nodes
    pub total_nodes: usize,
}

impl ClusterStats {
    /// Calculate from topology and metrics
    #[must_use]
    pub fn from_topology_and_metrics(topology: &ClusterTopology, metrics: &ClusterMetrics) -> Self {
        let healthy_nodes = topology
            .nodes
            .iter()
            .filter(|n| n.status == NodeStatus::Healthy)
            .count();

        let avg_gpu_util = metrics.average_gpu_utilization();
        let max_gpu_util = metrics
            .gpu_utilization
            .iter()
            .map(|(_, u)| *u)
            .fold(0.0, f64::max);

        let total_memory_used = metrics.total_memory_used_gb();
        let total_memory_capacity = topology.total_memory_gb() as f64;

        let avg_bandwidth_util = metrics.average_bandwidth_usage();

        Self {
            avg_gpu_util,
            max_gpu_util,
            total_memory_used,
            total_memory_capacity,
            avg_bandwidth_util,
            total_tflops: metrics.total_tflops,
            healthy_nodes,
            total_nodes: topology.nodes.len(),
        }
    }

    /// Calculate memory utilization ratio
    #[must_use]
    pub fn memory_utilization(&self) -> f64 {
        if self.total_memory_capacity > 0.0 {
            self.total_memory_used / self.total_memory_capacity
        } else {
            0.0
        }
    }

    /// Calculate cluster health score (0.0 - 1.0)
    #[must_use]
    pub fn health_score(&self) -> f64 {
        if self.total_nodes == 0 {
            return 0.0;
        }

        let node_health = self.healthy_nodes as f64 / self.total_nodes as f64;
        let util_health = 1.0 - (self.max_gpu_util - 0.5).abs() * 0.5; // Optimal around 50-70%
        let memory_health = 1.0 - self.memory_utilization().max(0.8).powi(2);

        (node_health * 0.4 + util_health * 0.3 + memory_health * 0.3).clamp(0.0, 1.0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use clusterviz_shared::{LinkInfo, NodeInfo, TransportType};

    fn create_test_topology() -> ClusterTopology {
        let mut topology = ClusterTopology::new("test-cluster");

        let node1 = NodeInfo::new("node-1", "Apple M3 Max", 128);
        let node2 = NodeInfo::new("node-2", "Apple M3 Max", 128);
        let id1 = node1.id;
        let id2 = node2.id;

        topology.add_node(node1);
        topology.add_node(node2);
        topology.add_link(LinkInfo::new(id1, id2, TransportType::Thunderbolt));

        topology
    }

    #[tokio::test]
    async fn test_metrics_collector_creation() {
        let collector = MetricsCollector::new();
        assert!(collector.last_collection.is_none());
    }

    #[tokio::test]
    async fn test_collect_metrics() {
        let mut collector = MetricsCollector::new();
        let topology = create_test_topology();

        let metrics = collector.collect(&topology).await.unwrap();

        assert_eq!(metrics.gpu_utilization.len(), 2);
        assert_eq!(metrics.memory_used.len(), 2);
        assert_eq!(metrics.bandwidth_usage.len(), 1);
        assert!(collector.last_collection.is_some());
    }

    #[tokio::test]
    async fn test_utilization_tracker() {
        let mut tracker = UtilizationTracker::new();
        let node_id = Uuid::new_v4();

        let util = tracker.get_utilization(node_id).await.unwrap();

        assert!(util >= 0.0 && util <= 1.0);
        assert!(tracker.get_average(node_id).is_some());
    }

    #[tokio::test]
    async fn test_utilization_smoothing() {
        let mut tracker = UtilizationTracker::new().with_smoothing(0.5);
        let node_id = Uuid::new_v4();

        // Get multiple readings to test smoothing
        for _ in 0..5 {
            let _ = tracker.get_utilization(node_id).await.unwrap();
        }

        let trend = tracker.get_trend(node_id);
        // Trend might be None if not enough data
        assert!(trend.is_none() || trend.is_some());
    }

    #[tokio::test]
    async fn test_memory_tracker() {
        let mut tracker = MemoryTracker::new();
        let node_id = Uuid::new_v4();

        let mem = tracker.get_memory_used(node_id, 128).await.unwrap();

        assert!(mem >= 0.0 && mem <= 128.0);
    }

    #[tokio::test]
    async fn test_memory_peak() {
        let mut tracker = MemoryTracker::new();
        let node_id = Uuid::new_v4();

        for _ in 0..5 {
            let _ = tracker.get_memory_used(node_id, 128).await.unwrap();
        }

        let peak = tracker.get_peak(node_id);
        assert!(peak.is_some());
        assert!(peak.unwrap() <= 128.0);
    }

    #[tokio::test]
    async fn test_bandwidth_monitor() {
        let mut monitor = BandwidthMonitor::new();
        let link_id = Uuid::new_v4();

        let bandwidth = monitor.get_bandwidth_usage(link_id).await.unwrap();

        assert!(bandwidth >= 0.0 && bandwidth <= 120.0);
    }

    #[tokio::test]
    async fn test_bandwidth_efficiency() {
        let mut monitor = BandwidthMonitor::new();
        let link_id = Uuid::new_v4();

        let _ = monitor.get_bandwidth_usage(link_id).await.unwrap();
        let efficiency = monitor.get_efficiency(link_id, 120.0);

        assert!(efficiency.is_some());
        assert!(efficiency.unwrap() >= 0.0 && efficiency.unwrap() <= 1.0);
    }

    #[test]
    fn test_cluster_stats() {
        let topology = create_test_topology();
        let metrics = clusterviz_shared::sample_metrics(&topology);

        let stats = ClusterStats::from_topology_and_metrics(&topology, &metrics);

        assert_eq!(stats.total_nodes, 2);
        assert!(stats.avg_gpu_util >= 0.0 && stats.avg_gpu_util <= 1.0);
        assert!(stats.health_score() >= 0.0 && stats.health_score() <= 1.0);
    }

    #[test]
    fn test_cluster_stats_memory_utilization() {
        let topology = create_test_topology();
        let metrics = clusterviz_shared::sample_metrics(&topology);

        let stats = ClusterStats::from_topology_and_metrics(&topology, &metrics);
        let mem_util = stats.memory_utilization();

        assert!(mem_util >= 0.0 && mem_util <= 1.0);
    }

    #[tokio::test]
    async fn test_utilization_tracker_clear() {
        let mut tracker = UtilizationTracker::new();
        let node_id = Uuid::new_v4();

        let _ = tracker.get_utilization(node_id).await.unwrap();
        assert!(tracker.cache.contains_key(&node_id));

        tracker.clear_node(node_id);
        assert!(!tracker.cache.contains_key(&node_id));
    }

    #[tokio::test]
    async fn test_bandwidth_moving_average() {
        let mut monitor = BandwidthMonitor::new().with_window_size(5);
        let link_id = Uuid::new_v4();

        let mut readings = Vec::new();
        for _ in 0..10 {
            let bw = monitor.get_bandwidth_usage(link_id).await.unwrap();
            readings.push(bw);
        }

        // Later readings should be smoother due to moving average
        // Just verify we got valid readings
        assert_eq!(readings.len(), 10);
        assert!(readings.iter().all(|&r| r >= 0.0 && r <= 120.0));
    }
}