//! ClusterViz - Real-Time Cluster Monitor
//!
//! This crate provides a complete system for monitoring and visualizing
//! Thunderbolt 5 Mac cluster topology and data flow in real-time.
//!
//! # Features
//!
//! - **Topology Discovery**: Automatic detection of cluster nodes and links
//! - **Real-Time Metrics**: GPU utilization, memory usage, bandwidth monitoring
//! - **Bottleneck Detection**: Identify performance issues in real-time
//! - **Collective Tracing**: Visualize distributed operations (AllReduce, Broadcast, etc.)
//!
//! # Example
//!
//! ```rust,no_run
//! use rtx_clusterviz_demo::ClusterViz;
//!
//! #[tokio::main]
//! async fn main() -> anyhow::Result<()> {
//!     let mut clusterviz = ClusterViz::new();
//!
//!     // Discover cluster topology
//!     clusterviz.discover_topology().await?;
//!
//!     // Collect metrics
//!     let metrics = clusterviz.collect_metrics().await?;
//!     println!("Average GPU utilization: {:.1}%", metrics.average_gpu_utilization() * 100.0);
//!
//!     // Check for bottlenecks
//!     let bottlenecks = clusterviz.detect_bottlenecks().await?;
//!     for bn in bottlenecks {
//!         println!("Bottleneck: {} - {}", bn.bottleneck_type, bn.description);
//!     }
//!
//!     Ok(())
//! }
//! ```

pub mod metrics;
pub mod sample_data;
pub mod topology;

use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use clusterviz_shared::{
    AlertConfig, BottleneckInfo, BottleneckType, ClusterMetrics, ClusterTopology, CollectiveTrace,
    Severity,
};
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

use crate::metrics::MetricsCollector;
use crate::topology::TopologyDiscoverer;

/// Error types specific to ClusterViz
#[derive(Debug, thiserror::Error)]
pub enum ClusterVizError {
    /// Topology discovery failed
    #[error("Failed to discover cluster topology: {0}")]
    DiscoveryFailed(String),

    /// Metrics collection failed
    #[error("Failed to collect metrics: {0}")]
    MetricsCollectionFailed(String),

    /// No topology available
    #[error("No cluster topology available - run discover_topology() first")]
    NoTopology,

    /// Node not found
    #[error("Node not found: {0}")]
    NodeNotFound(String),

    /// Link not found
    #[error("Link not found: {0}")]
    LinkNotFound(String),

    /// Configuration error
    #[error("Configuration error: {0}")]
    ConfigError(String),
}

/// Main ClusterViz system for real-time cluster monitoring
pub struct ClusterViz {
    /// Current cluster topology
    topology: Option<Arc<RwLock<ClusterTopology>>>,
    /// Metrics collector
    metrics_collector: MetricsCollector,
    /// Topology discoverer
    topology_discoverer: TopologyDiscoverer,
    /// Alert configuration
    alert_config: AlertConfig,
    /// Historical metrics (for trend analysis)
    metrics_history: Vec<ClusterMetrics>,
    /// Maximum history size
    max_history_size: usize,
    /// Active collective traces
    active_traces: Vec<CollectiveTrace>,
    /// Detected bottlenecks
    detected_bottlenecks: Vec<BottleneckInfo>,
}

impl ClusterViz {
    /// Create a new ClusterViz instance
    #[must_use]
    pub fn new() -> Self {
        Self {
            topology: None,
            metrics_collector: MetricsCollector::new(),
            topology_discoverer: TopologyDiscoverer::new(),
            alert_config: AlertConfig::new(),
            metrics_history: Vec::new(),
            max_history_size: 1000,
            active_traces: Vec::new(),
            detected_bottlenecks: Vec::new(),
        }
    }

    /// Create with custom alert configuration
    #[must_use]
    pub fn with_alert_config(mut self, config: AlertConfig) -> Self {
        self.alert_config = config;
        self
    }

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

    /// Discover the cluster topology automatically
    ///
    /// This method probes the network to find all available nodes
    /// and their interconnections. Returns a clone of the discovered topology.
    pub async fn discover_topology(&mut self) -> Result<ClusterTopology> {
        info!("Starting cluster topology discovery");

        let topology = self
            .topology_discoverer
            .discover()
            .await
            .context("Failed to discover topology")?;

        info!(
            "Discovered {} nodes and {} links",
            topology.node_count(),
            topology.link_count()
        );

        let topology_clone = topology.clone();
        self.topology = Some(Arc::new(RwLock::new(topology)));

        Ok(topology_clone)
    }

    /// Set the topology manually (for testing or pre-configured clusters)
    pub fn set_topology(&mut self, topology: ClusterTopology) {
        self.topology = Some(Arc::new(RwLock::new(topology)));
    }

    /// Get a reference to the current topology
    #[must_use]
    pub fn get_topology(&self) -> Option<&Arc<RwLock<ClusterTopology>>> {
        self.topology.as_ref()
    }

    /// Collect real-time metrics from the cluster
    pub async fn collect_metrics(&mut self) -> Result<ClusterMetrics> {
        let topology = self.topology.as_ref().ok_or(ClusterVizError::NoTopology)?;

        let topology_read = topology.read().await;

        debug!(
            "Collecting metrics for {} nodes",
            topology_read.node_count()
        );

        let metrics = self
            .metrics_collector
            .collect(&topology_read)
            .await
            .context("Failed to collect metrics")?;

        // Store in history
        self.metrics_history.push(metrics.clone());
        if self.metrics_history.len() > self.max_history_size {
            self.metrics_history.remove(0);
        }

        Ok(metrics)
    }

    /// Detect bottlenecks in the cluster
    pub async fn detect_bottlenecks(&mut self) -> Result<Vec<BottleneckInfo>> {
        // Collect metrics first (requires mutable borrow)
        let metrics = self.collect_metrics().await?;

        // Then get topology for analysis
        let topology = self.topology.as_ref().ok_or(ClusterVizError::NoTopology)?;

        let topology_read = topology.read().await;

        let mut bottlenecks = Vec::new();

        // Check GPU utilization
        for (node_id, util) in &metrics.gpu_utilization {
            if *util >= self.alert_config.gpu_util_critical {
                let mut bn = BottleneckInfo::new(
                    BottleneckType::ComputeBottleneck,
                    Severity::Critical,
                    format!("GPU utilization at {:.1}%", util * 100.0),
                )
                .with_metrics(*util, self.alert_config.gpu_util_critical);

                bn.add_affected_node(*node_id);
                bn.remediation =
                    Some("Consider redistributing workload or adding more nodes".to_string());
                bottlenecks.push(bn);
            } else if *util >= self.alert_config.gpu_util_warning {
                let mut bn = BottleneckInfo::new(
                    BottleneckType::ComputeBottleneck,
                    Severity::Warning,
                    format!("GPU utilization at {:.1}%", util * 100.0),
                )
                .with_metrics(*util, self.alert_config.gpu_util_warning);

                bn.add_affected_node(*node_id);
                bottlenecks.push(bn);
            }
        }

        // Check memory usage
        for (node_id, mem_used) in &metrics.memory_used {
            if let Some(node) = topology_read.get_node(*node_id) {
                let usage_ratio = mem_used / node.memory_gb as f64;
                if usage_ratio >= self.alert_config.memory_critical {
                    let mut bn = BottleneckInfo::new(
                        BottleneckType::MemoryPressure,
                        Severity::Critical,
                        format!("Memory usage at {:.1}%", usage_ratio * 100.0),
                    )
                    .with_metrics(usage_ratio, self.alert_config.memory_critical);

                    bn.add_affected_node(*node_id);
                    bn.remediation =
                        Some("Reduce batch size or enable gradient checkpointing".to_string());
                    bottlenecks.push(bn);
                } else if usage_ratio >= self.alert_config.memory_warning {
                    let mut bn = BottleneckInfo::new(
                        BottleneckType::MemoryPressure,
                        Severity::Warning,
                        format!("Memory usage at {:.1}%", usage_ratio * 100.0),
                    )
                    .with_metrics(usage_ratio, self.alert_config.memory_warning);

                    bn.add_affected_node(*node_id);
                    bottlenecks.push(bn);
                }
            }
        }

        // Check bandwidth saturation
        for (link_id, bandwidth) in &metrics.bandwidth_usage {
            if let Some(link) = topology_read.links.iter().find(|l| l.id == *link_id) {
                let saturation = bandwidth / link.transport_type.max_bandwidth_gbps();
                if saturation >= self.alert_config.bandwidth_critical {
                    let mut bn = BottleneckInfo::new(
                        BottleneckType::BandwidthSaturation,
                        Severity::Critical,
                        format!("Link bandwidth at {:.1}%", saturation * 100.0),
                    )
                    .with_metrics(saturation, self.alert_config.bandwidth_critical);

                    bn.add_affected_link(*link_id);
                    bn.remediation =
                        Some("Add parallel links or optimize communication patterns".to_string());
                    bottlenecks.push(bn);
                } else if saturation >= self.alert_config.bandwidth_warning {
                    let mut bn = BottleneckInfo::new(
                        BottleneckType::BandwidthSaturation,
                        Severity::Warning,
                        format!("Link bandwidth at {:.1}%", saturation * 100.0),
                    )
                    .with_metrics(saturation, self.alert_config.bandwidth_warning);

                    bn.add_affected_link(*link_id);
                    bottlenecks.push(bn);
                }

                // Check latency
                if link.latency_us >= self.alert_config.latency_critical_us {
                    let mut bn = BottleneckInfo::new(
                        BottleneckType::HighLatency,
                        Severity::Critical,
                        format!("Link latency at {:.1}us", link.latency_us),
                    )
                    .with_metrics(link.latency_us, self.alert_config.latency_critical_us);

                    bn.add_affected_link(*link_id);
                    bottlenecks.push(bn);
                } else if link.latency_us >= self.alert_config.latency_warning_us {
                    let mut bn = BottleneckInfo::new(
                        BottleneckType::HighLatency,
                        Severity::Warning,
                        format!("Link latency at {:.1}us", link.latency_us),
                    )
                    .with_metrics(link.latency_us, self.alert_config.latency_warning_us);

                    bn.add_affected_link(*link_id);
                    bottlenecks.push(bn);
                }
            }
        }

        // Check load imbalance
        if !metrics.gpu_utilization.is_empty() {
            let utils: Vec<f64> = metrics.gpu_utilization.iter().map(|(_, u)| *u).collect();
            let mean: f64 = utils.iter().sum::<f64>() / utils.len() as f64;
            let variance: f64 =
                utils.iter().map(|u| (u - mean).powi(2)).sum::<f64>() / utils.len() as f64;
            let std_dev = variance.sqrt();
            let coefficient_of_variation = if mean > 0.0 { std_dev / mean } else { 0.0 };

            if coefficient_of_variation >= self.alert_config.load_imbalance_warning {
                let mut bn = BottleneckInfo::new(
                    BottleneckType::LoadImbalance,
                    Severity::Warning,
                    format!(
                        "Load imbalance detected (CV: {:.2})",
                        coefficient_of_variation
                    ),
                )
                .with_metrics(
                    coefficient_of_variation,
                    self.alert_config.load_imbalance_warning,
                );

                bn.remediation =
                    Some("Review data partitioning and workload distribution".to_string());
                bottlenecks.push(bn);
            }
        }

        // Store detected bottlenecks
        self.detected_bottlenecks = bottlenecks.clone();

        if !bottlenecks.is_empty() {
            warn!("Detected {} bottlenecks", bottlenecks.len());
        }

        Ok(bottlenecks)
    }

    /// Add an active collective trace
    pub fn add_trace(&mut self, trace: CollectiveTrace) {
        self.active_traces.push(trace);
    }

    /// Get active traces
    #[must_use]
    pub fn get_active_traces(&self) -> &[CollectiveTrace] {
        &self.active_traces
    }

    /// Clear completed traces
    pub fn clear_traces(&mut self) {
        self.active_traces.clear();
    }

    /// Get metrics history
    #[must_use]
    pub fn get_metrics_history(&self) -> &[ClusterMetrics] {
        &self.metrics_history
    }

    /// Get detected bottlenecks
    #[must_use]
    pub fn get_bottlenecks(&self) -> &[BottleneckInfo] {
        &self.detected_bottlenecks
    }

    /// Get current alert configuration
    #[must_use]
    pub fn get_alert_config(&self) -> &AlertConfig {
        &self.alert_config
    }

    /// Update alert configuration
    pub fn set_alert_config(&mut self, config: AlertConfig) {
        self.alert_config = config;
    }

    /// Run the demo with sample data
    pub async fn run_demo(&mut self) -> Result<()> {
        info!("Starting ClusterViz Demo");

        // Use sample topology
        let topology = sample_data::four_node_cluster();
        self.set_topology(topology);

        info!("Loaded 4-node cluster topology");

        // Collect metrics
        let metrics = self.collect_metrics().await?;
        info!(
            "Collected metrics - Average GPU utilization: {:.1}%",
            metrics.average_gpu_utilization() * 100.0
        );
        info!(
            "Total memory used: {:.1} GB",
            metrics.total_memory_used_gb()
        );

        // Add sample traces
        let node_ids: Vec<_> = {
            let topology = self.topology.as_ref().unwrap().read().await;
            topology.nodes.iter().map(|n| n.id).collect()
        };

        let allreduce_trace = clusterviz_shared::sample_allreduce_trace(&node_ids);
        self.add_trace(allreduce_trace);

        let broadcast_trace = clusterviz_shared::sample_broadcast_trace(&node_ids, 0);
        self.add_trace(broadcast_trace);

        info!("Added {} collective traces", self.active_traces.len());

        // Detect bottlenecks
        let bottlenecks = self.detect_bottlenecks().await?;
        if bottlenecks.is_empty() {
            info!("No bottlenecks detected");
        } else {
            for bn in &bottlenecks {
                warn!(
                    "Bottleneck: {} - {} (severity: {})",
                    bn.bottleneck_type, bn.description, bn.severity
                );
            }
        }

        // Run monitoring loop (3 iterations for demo)
        for i in 1..=3 {
            tokio::time::sleep(Duration::from_millis(500)).await;

            let metrics = self.collect_metrics().await?;
            debug!(
                "Iteration {}: GPU util {:.1}%, Memory {:.1} GB",
                i,
                metrics.average_gpu_utilization() * 100.0,
                metrics.total_memory_used_gb()
            );
        }

        info!("ClusterViz Demo completed successfully");
        Ok(())
    }
}

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

/// Run the demo
pub async fn run_demo() -> Result<()> {
    let mut clusterviz = ClusterViz::new();
    clusterviz.run_demo().await
}

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

    #[tokio::test]
    async fn test_clusterviz_creation() {
        let clusterviz = ClusterViz::new();
        assert!(clusterviz.topology.is_none());
        assert!(clusterviz.metrics_history.is_empty());
    }

    #[tokio::test]
    async fn test_set_topology() {
        let mut clusterviz = ClusterViz::new();
        let topology = sample_data::four_node_cluster();

        clusterviz.set_topology(topology);

        assert!(clusterviz.topology.is_some());
        let topo = clusterviz.topology.as_ref().unwrap().read().await;
        assert_eq!(topo.node_count(), 4);
    }

    #[tokio::test]
    async fn test_collect_metrics() {
        let mut clusterviz = ClusterViz::new();
        let topology = sample_data::four_node_cluster();
        clusterviz.set_topology(topology);

        let metrics = clusterviz.collect_metrics().await.unwrap();

        assert!(!metrics.gpu_utilization.is_empty());
        assert!(!metrics.memory_used.is_empty());
        assert_eq!(clusterviz.metrics_history.len(), 1);
    }

    #[tokio::test]
    async fn test_collect_metrics_no_topology() {
        let mut clusterviz = ClusterViz::new();
        let result = clusterviz.collect_metrics().await;

        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_detect_bottlenecks() {
        let mut clusterviz = ClusterViz::new();
        let topology = sample_data::four_node_cluster();
        clusterviz.set_topology(topology);

        let bottlenecks = clusterviz.detect_bottlenecks().await.unwrap();

        // With default sample data, we might or might not have bottlenecks
        // depending on the random utilization values
        assert!(clusterviz.detected_bottlenecks.len() == bottlenecks.len());
    }

    #[tokio::test]
    async fn test_alert_config() {
        let clusterviz = ClusterViz::new().with_alert_config(AlertConfig::strict());

        assert!(clusterviz.alert_config.gpu_util_warning < 0.80);
    }

    #[tokio::test]
    async fn test_add_trace() {
        let mut clusterviz = ClusterViz::new();
        let nodes = vec![uuid::Uuid::new_v4(), uuid::Uuid::new_v4()];
        let trace = clusterviz_shared::sample_allreduce_trace(&nodes);

        clusterviz.add_trace(trace);

        assert_eq!(clusterviz.active_traces.len(), 1);
    }

    #[tokio::test]
    async fn test_clear_traces() {
        let mut clusterviz = ClusterViz::new();
        let nodes = vec![uuid::Uuid::new_v4(), uuid::Uuid::new_v4()];
        let trace = clusterviz_shared::sample_allreduce_trace(&nodes);

        clusterviz.add_trace(trace);
        assert_eq!(clusterviz.active_traces.len(), 1);

        clusterviz.clear_traces();
        assert!(clusterviz.active_traces.is_empty());
    }

    #[tokio::test]
    async fn test_max_history() {
        let mut clusterviz = ClusterViz::new().with_max_history(3);
        let topology = sample_data::four_node_cluster();
        clusterviz.set_topology(topology);

        // Collect more metrics than history size
        for _ in 0..5 {
            let _ = clusterviz.collect_metrics().await.unwrap();
        }

        assert_eq!(clusterviz.metrics_history.len(), 3);
    }

    #[tokio::test]
    async fn test_run_demo() {
        let mut clusterviz = ClusterViz::new();
        let result = clusterviz.run_demo().await;

        assert!(result.is_ok());
        assert!(clusterviz.topology.is_some());
        assert!(!clusterviz.metrics_history.is_empty());
    }
}