//! Topology discovery and optimization for distributed training
//!
//! This module provides functionality to discover network topology,
//! optimize communication patterns, and adapt to hardware characteristics.

use crate::error::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Network topology information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopologyInfo {
    /// Number of nodes in the cluster
    pub num_nodes: usize,
    /// Number of GPUs per node
    pub gpus_per_node: usize,
    /// Network interconnect type (e.g., "infiniband", "ethernet")
    pub interconnect: String,
    /// Bandwidth information (Gbps)
    pub bandwidth: TopologyBandwidth,
    /// Node-to-node latency matrix (microseconds)
    pub latency_matrix: Vec<Vec<f32>>,
    /// GPU-to-GPU topology within nodes
    pub intra_node_topology: Vec<GpuTopology>,
    /// Inter-node network topology
    pub inter_node_topology: NetworkTopology,
}

/// Bandwidth information for different communication types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopologyBandwidth {
    /// Intra-node GPU-to-GPU bandwidth (NVLink, etc.)
    pub intra_node_gpu: f32,
    /// Intra-node CPU-GPU bandwidth (PCIe)
    pub cpu_gpu: f32,
    /// Inter-node network bandwidth
    pub inter_node: f32,
    /// Memory bandwidth per GPU
    pub memory: f32,
}

/// GPU topology within a single node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuTopology {
    /// Node ID
    pub node_id: usize,
    /// GPU devices in this node
    pub gpus: Vec<GpuDevice>,
    /// GPU-to-GPU connectivity matrix
    pub gpu_connectivity: Vec<Vec<ConnectionType>>,
}

/// GPU device information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuDevice {
    /// Device ID within the node
    pub device_id: usize,
    /// GPU model/architecture
    pub model: String,
    /// Memory size in GB
    pub memory_gb: f32,
    /// Compute capability
    pub compute_capability: String,
    /// CUDA/ROCm device index
    pub device_index: i32,
}

/// Connection type between GPUs
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ConnectionType {
    /// No direct connection
    None,
    /// PCIe connection
    Pcie,
    /// NVLink connection (NVIDIA)
    NvLink,
    /// Infinity Fabric (AMD)
    InfinityFabric,
    /// Cross-node network connection
    Network,
}

/// Network topology between nodes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkTopology {
    /// Network type (tree, fat-tree, torus, etc.)
    pub topology_type: String,
    /// Number of network switches/hops between nodes
    pub switch_matrix: Vec<Vec<usize>>,
    /// Network interface details per node
    pub network_interfaces: Vec<NetworkInterface>,
}

impl Default for NetworkTopology {
    fn default() -> Self {
        Self {
            topology_type: "single-node".to_string(),
            switch_matrix: vec![vec![0]],
            network_interfaces: vec![],
        }
    }
}

/// Network interface information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkInterface {
    /// Node ID
    pub node_id: usize,
    /// Interface name (e.g., "ib0", "eth0")
    pub interface_name: String,
    /// IP address
    pub ip_address: String,
    /// Bandwidth in Gbps
    pub bandwidth_gbps: f32,
    /// MTU size
    pub mtu: usize,
}

/// Topology optimizer for communication patterns
pub struct TopologyOptimizer {
    topology: TopologyInfo,
    optimization_cache: HashMap<String, OptimizationResult>,
}

/// Result of topology optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationResult {
    /// Recommended communication pattern
    pub pattern: CommunicationPattern,
    /// Expected performance improvement
    pub improvement_factor: f32,
    /// Memory overhead
    pub memory_overhead_mb: f32,
}

/// Communication pattern recommendation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CommunicationPattern {
    /// Ring-based AllReduce
    Ring,
    /// Tree-based AllReduce
    Tree,
    /// Recursive doubling
    RecursiveDoubling,
    /// Hierarchical (intra-node + inter-node)
    Hierarchical,
    /// Custom pattern with specific routing
    Custom { routing: Vec<CommunicationRoute> },
}

/// Communication route specification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommunicationRoute {
    /// Source rank
    pub src: i32,
    /// Destination rank
    pub dst: i32,
    /// Intermediate hops
    pub hops: Vec<i32>,
    /// Expected bandwidth utilization
    pub bandwidth_utilization: f32,
}

/// AllReduce communication pattern
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AllReducePattern {
    /// Algorithm used (e.g., "ring", "tree", "hierarchical")
    pub algorithm: String,
    /// Communication steps
    pub steps: Vec<CommunicationStep>,
    /// Estimated completion time in milliseconds
    pub estimated_time_ms: f64,
    /// Expected bandwidth in GB/s
    pub expected_bandwidth_gbps: f32,
}

/// Individual communication step
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommunicationStep {
    /// Operation name
    pub operation: String,
    /// Participating ranks
    pub participants: Vec<usize>,
    /// Estimated time for this step
    pub estimated_time_ms: f64,
}

impl TopologyInfo {
    /// Discover topology information from the current environment
    pub fn discover() -> Result<Self> {
        // For now, return a simulated topology
        // In a real implementation, this would:
        // 1. Query GPU topology via CUDA/ROCm APIs
        // 2. Discover network interfaces
        // 3. Measure inter-node bandwidth/latency
        // 4. Build topology matrices

        Ok(Self::default_topology())
    }

    /// Create default topology for testing
    pub fn default_topology() -> Self {
        let num_nodes = 2;
        let gpus_per_node = 4;

        Self {
            num_nodes,
            gpus_per_node,
            interconnect: "infiniband".to_string(),
            bandwidth: TopologyBandwidth {
                intra_node_gpu: 600.0, // 600 GB/s NVLink
                cpu_gpu: 32.0,         // 32 GB/s PCIe 4.0 x16
                inter_node: 200.0,     // 200 Gbps InfiniBand
                memory: 2000.0,        // 2 TB/s HBM
            },
            latency_matrix: vec![
                vec![0.0, 150.0], // Node 0 to [Node 0, Node 1] in μs
                vec![150.0, 0.0], // Node 1 to [Node 0, Node 1] in μs
            ],
            intra_node_topology: (0..num_nodes)
                .map(|node_id| GpuTopology::default_for_node(node_id, gpus_per_node))
                .collect(),
            inter_node_topology: NetworkTopology::default_for_nodes(num_nodes),
        }
    }

    /// Get total number of GPUs in the cluster
    pub fn total_gpus(&self) -> usize {
        self.num_nodes * self.gpus_per_node
    }

    /// Check if two ranks are on the same node
    pub fn same_node(&self, rank1: i32, rank2: i32) -> bool {
        let node1 = rank1 as usize / self.gpus_per_node;
        let node2 = rank2 as usize / self.gpus_per_node;
        node1 == node2
    }

    /// Get node ID for a given rank
    pub fn node_for_rank(&self, rank: i32) -> usize {
        rank as usize / self.gpus_per_node
    }

    /// Get GPU ID within node for a given rank
    pub fn gpu_for_rank(&self, rank: i32) -> usize {
        rank as usize % self.gpus_per_node
    }

    /// Get total device count across all nodes
    pub fn device_count(&self) -> usize {
        self.num_nodes * self.gpus_per_node
    }

    /// Check if topology has NVLink connections
    pub fn has_nvlink_connections(&self) -> bool {
        self.intra_node_topology.iter().any(|node| {
            node.gpu_connectivity
                .iter()
                .any(|row| row.contains(&ConnectionType::NvLink))
        })
    }

    /// Get bandwidth matrix for all connections
    pub fn bandwidth_matrix(&self) -> Vec<Vec<f32>> {
        let total_devices = self.device_count();
        let mut matrix = vec![vec![0.0; total_devices]; total_devices];

        // Fill with intra-node bandwidth
        for node in &self.intra_node_topology {
            for (i, row) in node.gpu_connectivity.iter().enumerate() {
                for (j, &conn) in row.iter().enumerate() {
                    let global_i = node.node_id * self.gpus_per_node + i;
                    let global_j = node.node_id * self.gpus_per_node + j;

                    matrix[global_i][global_j] = match conn {
                        ConnectionType::NvLink => self.bandwidth.intra_node_gpu,
                        ConnectionType::Pcie => self.bandwidth.cpu_gpu,
                        ConnectionType::InfinityFabric => self.bandwidth.intra_node_gpu,
                        ConnectionType::Network => self.bandwidth.inter_node,
                        ConnectionType::None => 0.0,
                    };
                }
            }
        }

        // Add inter-node bandwidth
        for i in 0..total_devices {
            for j in 0..total_devices {
                if self.node_for_rank(i as i32) != self.node_for_rank(j as i32) {
                    matrix[i][j] = self.bandwidth.inter_node;
                }
            }
        }

        matrix
    }

    /// Count NVLink connections
    pub fn nvlink_connection_count(&self) -> usize {
        self.intra_node_topology
            .iter()
            .map(|node| {
                node.gpu_connectivity
                    .iter()
                    .map(|row| {
                        row.iter()
                            .filter(|&&conn| conn == ConnectionType::NvLink)
                            .count()
                    })
                    .sum::<usize>()
            })
            .sum()
    }
}

impl GpuTopology {
    /// Create default GPU topology for a node
    pub fn default_for_node(node_id: usize, gpus_per_node: usize) -> Self {
        let gpus = (0..gpus_per_node)
            .map(|gpu_id| GpuDevice {
                device_id: gpu_id,
                model: "RTX-5090".to_string(),
                memory_gb: 32.0,
                compute_capability: "9.0".to_string(),
                device_index: (node_id * gpus_per_node + gpu_id) as i32,
            })
            .collect();

        // Create fully connected NVLink topology
        let mut gpu_connectivity = vec![vec![ConnectionType::None; gpus_per_node]; gpus_per_node];
        for i in 0..gpus_per_node {
            for j in 0..gpus_per_node {
                gpu_connectivity[i][j] = if i == j {
                    ConnectionType::None
                } else {
                    ConnectionType::NvLink
                };
            }
        }

        Self {
            node_id,
            gpus,
            gpu_connectivity,
        }
    }
}

impl NetworkTopology {
    /// Create default network topology for nodes
    pub fn default_for_nodes(num_nodes: usize) -> Self {
        // Simple fully-connected topology
        let switch_matrix = vec![vec![1; num_nodes]; num_nodes];

        let network_interfaces = (0..num_nodes)
            .map(|node_id| NetworkInterface {
                node_id,
                interface_name: format!("ib{node_id}"),
                ip_address: format!("10.0.0.{}", node_id + 1),
                bandwidth_gbps: 200.0,
                mtu: 4096,
            })
            .collect();

        Self {
            topology_type: "fully_connected".to_string(),
            switch_matrix,
            network_interfaces,
        }
    }
}

impl TopologyOptimizer {
    /// Create new topology optimizer
    pub fn new(topology: TopologyInfo) -> Self {
        Self {
            topology,
            optimization_cache: HashMap::new(),
        }
    }

    /// Optimize communication pattern for AllReduce
    pub fn optimize_allreduce(
        &mut self,
        world_size: i32,
        tensor_size_mb: f32,
    ) -> Result<OptimizationResult> {
        let cache_key = format!("allreduce_{world_size}_{tensor_size_mb}");

        if let Some(cached) = self.optimization_cache.get(&cache_key) {
            return Ok(cached.clone());
        }

        let pattern = if world_size <= 8 {
            CommunicationPattern::Ring
        } else if tensor_size_mb < 1.0 {
            CommunicationPattern::Tree
        } else {
            CommunicationPattern::Hierarchical
        };

        let improvement_factor = self.estimate_improvement(&pattern, world_size, tensor_size_mb);
        let memory_overhead = self.estimate_memory_overhead(&pattern, tensor_size_mb);

        let result = OptimizationResult {
            pattern,
            improvement_factor,
            memory_overhead_mb: memory_overhead,
        };

        self.optimization_cache.insert(cache_key, result.clone());
        Ok(result)
    }

    /// Estimate performance improvement for a communication pattern
    fn estimate_improvement(
        &self,
        pattern: &CommunicationPattern,
        world_size: i32,
        tensor_size_mb: f32,
    ) -> f32 {
        match pattern {
            CommunicationPattern::Ring => {
                // Ring is good for large tensors
                if tensor_size_mb > 10.0 { 1.2 } else { 0.9 }
            }
            CommunicationPattern::Tree => {
                // Tree is good for small tensors
                if tensor_size_mb < 1.0 { 1.4 } else { 0.8 }
            }
            CommunicationPattern::Hierarchical => {
                // Hierarchical is good for multi-node
                if world_size as usize > self.topology.gpus_per_node {
                    1.5
                } else {
                    1.0
                }
            }
            _ => 1.0,
        }
    }

    /// Estimate memory overhead for a communication pattern
    fn estimate_memory_overhead(&self, pattern: &CommunicationPattern, tensor_size_mb: f32) -> f32 {
        match pattern {
            CommunicationPattern::Ring => tensor_size_mb * 0.1, // 10% overhead
            CommunicationPattern::Tree => tensor_size_mb * 0.2, // 20% overhead
            CommunicationPattern::Hierarchical => tensor_size_mb * 0.15, // 15% overhead
            _ => tensor_size_mb * 0.05,                         // 5% default overhead
        }
    }

    /// Create new topology optimizer with default topology
    pub fn with_default() -> Self {
        // Create a default topology for discovery
        let topology = TopologyInfo::default_topology();
        Self {
            topology,
            optimization_cache: HashMap::new(),
        }
    }

    /// Discover topology from device list
    pub fn discover_topology(&self, devices: &[rtx_runtime::Device]) -> Result<TopologyInfo> {
        // Create topology info based on device list
        let device_count = devices.len();
        let nodes = device_count.div_ceil(8); // Assume 8 GPUs per node
        let gpus_per_node = if nodes == 1 { device_count } else { 8 };

        // Create GPU devices
        let mut gpus = Vec::new();
        for (i, device) in devices.iter().enumerate() {
            gpus.push(GpuDevice {
                device_id: i,
                model: format!("GPU-{i}"), // Simplified model name
                memory_gb: (device.properties().total_memory / (1024 * 1024 * 1024)) as f32,
                compute_capability: format!(
                    "{}.{}",
                    device.properties().compute_capability.0,
                    device.properties().compute_capability.1
                ),
                device_index: i as i32,
            });
        }

        // Create GPU topology for single node
        let mut intra_node_topology = Vec::new();
        if !gpus.is_empty() {
            let mut gpu_connectivity = vec![vec![ConnectionType::None; device_count]; device_count];

            // Simulate NVLink connections for RTX 5090
            for i in 0..device_count {
                for j in 0..device_count {
                    if i != j && (i as i32 - j as i32).abs() <= 2 {
                        gpu_connectivity[i][j] = ConnectionType::NvLink;
                    } else if i != j {
                        gpu_connectivity[i][j] = ConnectionType::Pcie;
                    }
                }
            }

            intra_node_topology.push(GpuTopology {
                node_id: 0,
                gpus,
                gpu_connectivity,
            });
        }

        let topology = TopologyInfo {
            num_nodes: nodes,
            gpus_per_node,
            interconnect: "nvlink+infiniband".to_string(),
            bandwidth: TopologyBandwidth {
                intra_node_gpu: 900.0, // RTX 5090 NVLink bandwidth (GB/s)
                cpu_gpu: 32.0,         // PCIe Gen4 x16
                inter_node: 200.0,     // InfiniBand HDR
                memory: 1000.0,        // RTX 5090 memory bandwidth
            },
            latency_matrix: vec![vec![0.1; nodes]; nodes], // 0.1ms latency
            intra_node_topology,
            inter_node_topology: NetworkTopology {
                topology_type: "fat_tree".to_string(),
                switch_matrix: vec![vec![1; nodes]; nodes], // Simple switch matrix
                network_interfaces: vec![
                    NetworkInterface {
                        node_id: 0,
                        interface_name: "ib0".to_string(),
                        ip_address: "192.168.1.1".to_string(),
                        bandwidth_gbps: 200.0,
                        mtu: 9000,
                    };
                    nodes
                ],
            },
        };

        Ok(topology)
    }

    /// Optimize AllReduce pattern for given topology
    pub fn optimize_allreduce_pattern(&self, topology: &TopologyInfo) -> Result<AllReducePattern> {
        let total_gpus = topology.total_gpus();

        // Choose pattern based on topology
        let algorithm = if total_gpus <= 8 && topology.num_nodes == 1 {
            "nvlink_ring".to_string()
        } else if topology.interconnect.contains("nvlink") {
            "hierarchical_ring".to_string()
        } else {
            "recursive_doubling".to_string()
        };

        // Estimate completion time
        let estimated_time_ms = if topology.num_nodes == 1 {
            5.0 // Single node should be very fast with NVLink
        } else {
            topology.num_nodes as f64 * 10.0 // Scale with node count
        };

        let steps = if algorithm == "hierarchical_ring" {
            vec![
                CommunicationStep {
                    operation: "intra_node_ring".to_string(),
                    participants: (0..topology.gpus_per_node).collect(),
                    estimated_time_ms: estimated_time_ms * 0.3,
                },
                CommunicationStep {
                    operation: "inter_node_allreduce".to_string(),
                    participants: vec![0], // Node masters
                    estimated_time_ms: estimated_time_ms * 0.7,
                },
            ]
        } else {
            vec![CommunicationStep {
                operation: algorithm.clone(),
                participants: (0..total_gpus).collect(),
                estimated_time_ms,
            }]
        };

        Ok(AllReducePattern {
            algorithm,
            steps,
            estimated_time_ms,
            expected_bandwidth_gbps: topology.bandwidth.intra_node_gpu * 0.8, // 80% efficiency
        })
    }
}

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

    #[test]
    fn test_topology_discovery() {
        let topology = TopologyInfo::discover().unwrap();
        assert_eq!(topology.num_nodes, 2);
        assert_eq!(topology.gpus_per_node, 4);
        assert_eq!(topology.total_gpus(), 8);
    }

    #[test]
    fn test_rank_mapping() {
        let topology = TopologyInfo::default_topology();

        assert_eq!(topology.node_for_rank(0), 0);
        assert_eq!(topology.node_for_rank(3), 0);
        assert_eq!(topology.node_for_rank(4), 1);
        assert_eq!(topology.node_for_rank(7), 1);

        assert_eq!(topology.gpu_for_rank(0), 0);
        assert_eq!(topology.gpu_for_rank(3), 3);
        assert_eq!(topology.gpu_for_rank(4), 0);
        assert_eq!(topology.gpu_for_rank(7), 3);

        assert!(topology.same_node(0, 3));
        assert!(!topology.same_node(3, 4));
    }

    #[test]
    fn test_topology_optimizer() {
        let topology = TopologyInfo::default_topology();
        let mut optimizer = TopologyOptimizer::new(topology);

        let result = optimizer.optimize_allreduce(8, 10.0).unwrap();
        assert!(result.improvement_factor > 0.0);
        assert!(result.memory_overhead_mb > 0.0);

        // Test caching
        let result2 = optimizer.optimize_allreduce(8, 10.0).unwrap();
        assert_eq!(result.improvement_factor, result2.improvement_factor);
    }

    #[test]
    fn test_gpu_topology() {
        let gpu_topo = GpuTopology::default_for_node(0, 4);
        assert_eq!(gpu_topo.gpus.len(), 4);
        assert_eq!(gpu_topo.gpu_connectivity.len(), 4);

        // Check NVLink connectivity
        assert_eq!(gpu_topo.gpu_connectivity[0][1], ConnectionType::NvLink);
        assert_eq!(gpu_topo.gpu_connectivity[0][0], ConnectionType::None);
    }

    #[test]
    fn test_connection_type() {
        assert_eq!(ConnectionType::NvLink, ConnectionType::NvLink);
        assert_ne!(ConnectionType::NvLink, ConnectionType::Pcie);
    }
}