//! Sample Data Module
//!
//! This module provides sample cluster configurations, workloads, and traces
//! for demonstration and testing purposes.

use chrono::Utc;
use clusterviz_shared::{
    ClusterTopology, CollectiveOp, CollectiveTrace, DataFlowStep, LinkInfo, NodeInfo, NodeStatus,
    TransportType,
};
use uuid::Uuid;

/// Create a 4-node M3 Max cluster configuration
///
/// This represents a typical Mac Studio cluster connected via Thunderbolt 5.
#[must_use]
pub fn four_node_cluster() -> ClusterTopology {
    let mut topology = ClusterTopology::new("M3-Max-Training-Cluster");

    // Create 4 Mac Studio nodes with M3 Max
    let node1 = NodeInfo::new("mac-studio-alpha", "Apple M3 Max", 128).as_coordinator();
    let node2 = NodeInfo::new("mac-studio-beta", "Apple M3 Max", 128);
    let node3 = NodeInfo::new("mac-studio-gamma", "Apple M3 Max", 96);
    let node4 = NodeInfo::new("mac-studio-delta", "Apple M3 Max", 96);

    let ids = [node1.id, node2.id, node3.id, node4.id];

    topology.add_node(node1);
    topology.add_node(node2);
    topology.add_node(node3);
    topology.add_node(node4);

    // Create Thunderbolt 5 mesh topology
    // Full mesh for maximum bandwidth
    for i in 0..4 {
        for j in (i + 1)..4 {
            let mut link = LinkInfo::new(ids[i], ids[j], TransportType::Thunderbolt);
            // Slightly vary bandwidth to simulate real conditions
            link.bandwidth_gbps = 96.0 + (i as f64 * 2.0); // 96-102 Gbps
            link.latency_us = 2.5 + (j as f64 * 0.1); // 2.5-2.8 us
            topology.add_link(link);
        }
    }

    topology
}

/// Create an 8-node M4 Ultra cluster configuration
///
/// This represents a high-end cluster with next-generation hardware.
#[must_use]
pub fn eight_node_cluster() -> ClusterTopology {
    let mut topology = ClusterTopology::new("M4-Ultra-HPC-Cluster");

    // Create 8 Mac Pro nodes with M4 Ultra (theoretical)
    let gpu_configs = [
        ("mac-pro-01", "Apple M4 Ultra", 512),
        ("mac-pro-02", "Apple M4 Ultra", 512),
        ("mac-pro-03", "Apple M4 Ultra", 384),
        ("mac-pro-04", "Apple M4 Ultra", 384),
        ("mac-pro-05", "Apple M4 Max", 192),
        ("mac-pro-06", "Apple M4 Max", 192),
        ("mac-pro-07", "Apple M4 Max", 128),
        ("mac-pro-08", "Apple M4 Max", 128),
    ];

    let mut nodes = Vec::new();
    for (i, (name, gpu, mem)) in gpu_configs.iter().enumerate() {
        let mut node = NodeInfo::new(*name, *gpu, *mem);
        if i == 0 {
            node = node.as_coordinator();
        }
        nodes.push(node);
    }

    let ids: Vec<Uuid> = nodes.iter().map(|n| n.id).collect();

    for node in nodes {
        topology.add_node(node);
    }

    // Fat-tree topology for 8 nodes
    // Tier 1: Pairs (4 pairs)
    for i in (0..8).step_by(2) {
        let mut link = LinkInfo::new(ids[i], ids[i + 1], TransportType::Thunderbolt);
        link.bandwidth_gbps = 100.0;
        link.latency_us = 2.0;
        topology.add_link(link);
    }

    // Tier 2: Connect adjacent pairs
    for i in 0..4 {
        let src = i * 2;
        let dst = ((i + 1) % 4) * 2;
        let mut link = LinkInfo::new(ids[src], ids[dst], TransportType::Thunderbolt);
        link.bandwidth_gbps = 96.0;
        link.latency_us = 2.5;
        topology.add_link(link);
    }

    // Cross-tier links for redundancy
    topology.add_link(LinkInfo::new(ids[0], ids[4], TransportType::Thunderbolt));
    topology.add_link(LinkInfo::new(ids[2], ids[6], TransportType::Thunderbolt));
    topology.add_link(LinkInfo::new(ids[1], ids[5], TransportType::Thunderbolt));
    topology.add_link(LinkInfo::new(ids[3], ids[7], TransportType::Thunderbolt));

    topology
}

/// Create a mixed-generation cluster
///
/// Demonstrates heterogeneous hardware configurations.
#[must_use]
pub fn mixed_generation_cluster() -> ClusterTopology {
    let mut topology = ClusterTopology::new("Mixed-Generation-Cluster");

    // Mix of M3 and M4 generation hardware
    let configs = [
        ("workstation-01", "Apple M4 Max", 128, true),
        ("workstation-02", "Apple M3 Max", 128, false),
        ("workstation-03", "Apple M3 Pro", 64, false),
        ("workstation-04", "Apple M4 Pro", 96, false),
    ];

    let mut nodes = Vec::new();
    for (name, gpu, mem, is_coord) in &configs {
        let mut node = NodeInfo::new(*name, *gpu, *mem);
        if *is_coord {
            node = node.as_coordinator();
        }
        nodes.push(node);
    }

    let ids: Vec<Uuid> = nodes.iter().map(|n| n.id).collect();

    for node in nodes {
        topology.add_node(node);
    }

    // Star topology with coordinator at center
    for i in 1..4 {
        let mut link = LinkInfo::new(ids[0], ids[i], TransportType::Thunderbolt);
        // Bandwidth varies by GPU generation
        link.bandwidth_gbps = if i <= 2 { 96.0 } else { 80.0 };
        link.latency_us = 2.5;
        topology.add_link(link);
    }

    // Additional ring connection for redundancy
    for i in 1..4 {
        let next = if i == 3 { 1 } else { i + 1 };
        let mut link = LinkInfo::new(ids[i], ids[next], TransportType::Thunderbolt);
        link.bandwidth_gbps = 80.0;
        link.latency_us = 3.0;
        topology.add_link(link);
    }

    topology
}

/// Configure a topology for training workload
///
/// Sets utilization patterns typical for distributed training.
pub fn training_workload(topology: &mut ClusterTopology) {
    for node in &mut topology.nodes {
        // Training typically uses high GPU and memory
        node.utilization = 0.85 + (rand::random::<f64>() * 0.10);
        node.status = NodeStatus::Healthy;
        node.last_heartbeat = Utc::now();
    }

    for link in &mut topology.links {
        // Training uses moderate bandwidth for gradient sync
        link.bandwidth_gbps = link.transport_type.max_bandwidth_gbps() * 0.60;
        link.is_active = true;
    }
}

/// Configure a topology for inference workload
///
/// Sets utilization patterns typical for distributed inference.
pub fn inference_workload(topology: &mut ClusterTopology) {
    for (i, node) in topology.nodes.iter_mut().enumerate() {
        // Inference has variable utilization based on request load
        let base_util = 0.30 + (i as f64 * 0.05);
        node.utilization = base_util + (rand::random::<f64>() * 0.20);
        node.status = NodeStatus::Healthy;
        node.last_heartbeat = Utc::now();
    }

    for link in &mut topology.links {
        // Inference uses lower bandwidth (mostly input data)
        link.bandwidth_gbps = link.transport_type.max_bandwidth_gbps() * 0.25;
        link.is_active = rand::random::<bool>();
    }
}

/// Create sample collective traces for a set of nodes
#[must_use]
pub fn sample_collective_traces(node_ids: &[Uuid]) -> Vec<CollectiveTrace> {
    let mut traces = Vec::new();

    // AllReduce trace (gradient synchronization)
    let mut allreduce = CollectiveTrace::new(CollectiveOp::AllReduce, node_ids.to_vec());
    allreduce.description = Some("Gradient synchronization - FP16 gradients".to_string());

    let chunk_size: u64 = 128 * 1024 * 1024; // 128 MB
    let n = node_ids.len();

    // Ring all-reduce pattern
    for phase in 0..(n - 1) {
        for i in 0..n {
            let src = node_ids[i];
            let dst = node_ids[(i + 1) % n];
            let mut step = DataFlowStep::new(src, dst, chunk_size);
            step.sequence = (phase * n + i) as u32;
            step.duration_us = 1000 + (phase as u64 * 50);
            allreduce.add_step(step);
        }
    }
    allreduce.duration_us = (n - 1) as u64 * 1200;
    traces.push(allreduce);

    // Broadcast trace (model distribution)
    let mut broadcast = CollectiveTrace::new(CollectiveOp::Broadcast, node_ids.to_vec());
    broadcast.description = Some("Model weight broadcast - BF16 weights".to_string());

    let model_size: u64 = 2 * 1024 * 1024 * 1024; // 2 GB

    if !node_ids.is_empty() {
        let root = node_ids[0];
        for (i, &dst) in node_ids.iter().enumerate().skip(1) {
            let mut step = DataFlowStep::new(root, dst, model_size);
            step.sequence = i as u32;
            step.duration_us = 15000; // ~15ms at 120 Gbps
            broadcast.add_step(step);
        }
    }
    broadcast.duration_us = 15000;
    traces.push(broadcast);

    // AllGather trace (activation gathering)
    let mut allgather = CollectiveTrace::new(CollectiveOp::AllGather, node_ids.to_vec());
    allgather.description = Some("Activation gathering for tensor parallelism".to_string());

    let activation_size: u64 = 64 * 1024 * 1024; // 64 MB per node

    for (i, &src) in node_ids.iter().enumerate() {
        for (j, &dst) in node_ids.iter().enumerate() {
            if i != j {
                let mut step = DataFlowStep::new(src, dst, activation_size);
                step.sequence = (i * n + j) as u32;
                step.duration_us = 500;
                allgather.add_step(step);
            }
        }
    }
    allgather.duration_us = 500 * (n - 1) as u64;
    traces.push(allgather);

    // ReduceScatter trace
    let mut reduce_scatter = CollectiveTrace::new(CollectiveOp::ReduceScatter, node_ids.to_vec());
    reduce_scatter.description = Some("Reduce-scatter for pipeline parallelism".to_string());

    let scatter_size: u64 = 32 * 1024 * 1024; // 32 MB

    for phase in 0..(n - 1) {
        for i in 0..n {
            let src = node_ids[i];
            let dst = node_ids[(i + n - 1) % n];
            let mut step = DataFlowStep::new(src, dst, scatter_size);
            step.sequence = (phase * n + i) as u32;
            step.duration_us = 300;
            reduce_scatter.add_step(step);
        }
    }
    reduce_scatter.duration_us = 300 * (n - 1) as u64;
    traces.push(reduce_scatter);

    traces
}

/// Create traces for a training step
#[must_use]
pub fn training_step_traces(node_ids: &[Uuid], batch_size: usize) -> Vec<CollectiveTrace> {
    let mut traces = Vec::new();

    // Forward pass activation sync
    let mut forward = CollectiveTrace::new(CollectiveOp::AllGather, node_ids.to_vec());
    forward.description = Some(format!("Forward pass - batch size {}", batch_size));

    let activation_per_layer: u64 = (batch_size * 4096 * 4) as u64; // 4 bytes per element

    for &src in node_ids {
        for &dst in node_ids {
            if src != dst {
                let step = DataFlowStep::new(src, dst, activation_per_layer);
                forward.add_step(step);
            }
        }
    }
    forward.duration_us = 800;
    traces.push(forward);

    // Backward pass gradient sync
    let mut backward = CollectiveTrace::new(CollectiveOp::AllReduce, node_ids.to_vec());
    backward.description = Some("Backward pass - gradient synchronization".to_string());

    let gradient_size: u64 = 512 * 1024 * 1024; // 512 MB total gradients

    let n = node_ids.len();
    for i in 0..n {
        let src = node_ids[i];
        let dst = node_ids[(i + 1) % n];
        let mut step = DataFlowStep::new(src, dst, gradient_size / n as u64);
        step.duration_us = 2000;
        backward.add_step(step);
    }
    backward.duration_us = 2000;
    traces.push(backward);

    // Parameter update broadcast
    let mut update = CollectiveTrace::new(CollectiveOp::Broadcast, node_ids.to_vec());
    update.description = Some("Parameter update broadcast".to_string());

    if !node_ids.is_empty() {
        let root = node_ids[0];
        let param_size: u64 = 100 * 1024 * 1024; // 100 MB updated params

        for &dst in &node_ids[1..] {
            let step = DataFlowStep::new(root, dst, param_size);
            update.add_step(step);
        }
    }
    update.duration_us = 1000;
    traces.push(update);

    traces
}

/// Create traces for inference batching
#[must_use]
pub fn inference_batch_traces(
    node_ids: &[Uuid],
    requests_per_batch: usize,
) -> Vec<CollectiveTrace> {
    let mut traces = Vec::new();

    // Input distribution
    let mut input_dist = CollectiveTrace::new(CollectiveOp::Scatter, node_ids.to_vec());
    input_dist.description = Some(format!("Input scatter - {} requests", requests_per_batch));

    if !node_ids.is_empty() {
        let root = node_ids[0];
        let input_size: u64 = (requests_per_batch * 1024 * 4) as u64; // 4KB per request

        for &dst in &node_ids[1..] {
            let step = DataFlowStep::new(root, dst, input_size / node_ids.len() as u64);
            input_dist.add_step(step);
        }
    }
    input_dist.duration_us = 100;
    traces.push(input_dist);

    // Output gathering
    let mut output_gather = CollectiveTrace::new(CollectiveOp::AllGather, node_ids.to_vec());
    output_gather.description = Some("Output gathering".to_string());

    let output_size: u64 = (requests_per_batch * 4096 * 4) as u64; // Larger output

    for &src in node_ids {
        for &dst in node_ids {
            if src != dst {
                let mut step = DataFlowStep::new(src, dst, output_size / node_ids.len() as u64);
                step.duration_us = 200;
                output_gather.add_step(step);
            }
        }
    }
    output_gather.duration_us = 200;
    traces.push(output_gather);

    traces
}

/// Generate a stress test topology with high utilization
#[must_use]
pub fn stress_test_cluster() -> ClusterTopology {
    let mut topology = four_node_cluster();

    // Set all nodes to high utilization
    for node in &mut topology.nodes {
        node.utilization = 0.95 + (rand::random::<f64>() * 0.05);
        node.status = NodeStatus::HighLoad;
    }

    // Set all links to high bandwidth usage
    for link in &mut topology.links {
        link.bandwidth_gbps = link.transport_type.max_bandwidth_gbps() * 0.90;
        link.latency_us = link.transport_type.typical_latency_us() * 1.5;
        link.is_active = true;
    }

    topology
}

/// Generate a cluster with some degraded nodes
#[must_use]
pub fn degraded_cluster() -> ClusterTopology {
    let mut topology = four_node_cluster();

    // Mark one node as degraded
    if let Some(node) = topology.nodes.get_mut(2) {
        node.status = NodeStatus::Degraded;
        node.utilization = 0.3;
    }

    // Mark one node as high load
    if let Some(node) = topology.nodes.get_mut(3) {
        node.status = NodeStatus::HighLoad;
        node.utilization = 0.98;
    }

    topology
}

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

    #[test]
    fn test_four_node_cluster() {
        let topology = four_node_cluster();

        assert_eq!(topology.node_count(), 4);
        assert_eq!(topology.link_count(), 6); // Full mesh: 4*3/2 = 6 links
        assert_eq!(topology.total_memory_gb(), 448); // 128+128+96+96

        // Verify coordinator
        let coordinator_count = topology.nodes.iter().filter(|n| n.is_coordinator).count();
        assert_eq!(coordinator_count, 1);
    }

    #[test]
    fn test_eight_node_cluster() {
        let topology = eight_node_cluster();

        assert_eq!(topology.node_count(), 8);
        assert!(topology.link_count() >= 8); // At least pair + tier connections
        assert_eq!(topology.total_memory_gb(), 2432); // 512*2 + 384*2 + 192*2 + 128*2
    }

    #[test]
    fn test_mixed_generation_cluster() {
        let topology = mixed_generation_cluster();

        assert_eq!(topology.node_count(), 4);
        assert!(topology.link_count() >= 4);

        // Verify we have different GPU models
        let gpu_models: Vec<&str> = topology
            .nodes
            .iter()
            .map(|n| n.gpu_model.as_str())
            .collect();
        assert!(gpu_models.contains(&"Apple M4 Max"));
        assert!(gpu_models.contains(&"Apple M3 Max"));
    }

    #[test]
    fn test_training_workload() {
        let mut topology = four_node_cluster();
        training_workload(&mut topology);

        for node in &topology.nodes {
            assert!(node.utilization >= 0.85);
            assert!(node.utilization <= 0.95);
            assert_eq!(node.status, NodeStatus::Healthy);
        }

        for link in &topology.links {
            assert!(link.is_active);
        }
    }

    #[test]
    fn test_inference_workload() {
        let mut topology = four_node_cluster();
        inference_workload(&mut topology);

        for node in &topology.nodes {
            // Inference workload: base 0.30 + (index * 0.05) + random(0..0.20)
            // For 4 nodes, max is 0.30 + 0.15 + 0.20 = 0.65
            assert!(node.utilization >= 0.30);
            assert!(node.utilization <= 0.70);
        }
    }

    #[test]
    fn test_sample_collective_traces() {
        let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
        let traces = sample_collective_traces(&node_ids);

        assert_eq!(traces.len(), 4);

        // Verify trace types
        let ops: Vec<CollectiveOp> = traces.iter().map(|t| t.operation).collect();
        assert!(ops.contains(&CollectiveOp::AllReduce));
        assert!(ops.contains(&CollectiveOp::Broadcast));
        assert!(ops.contains(&CollectiveOp::AllGather));
        assert!(ops.contains(&CollectiveOp::ReduceScatter));
    }

    #[test]
    fn test_training_step_traces() {
        let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
        let traces = training_step_traces(&node_ids, 32);

        assert_eq!(traces.len(), 3); // Forward, backward, update
    }

    #[test]
    fn test_inference_batch_traces() {
        let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
        let traces = inference_batch_traces(&node_ids, 64);

        assert_eq!(traces.len(), 2); // Input scatter, output gather
    }

    #[test]
    fn test_stress_test_cluster() {
        let topology = stress_test_cluster();

        for node in &topology.nodes {
            assert!(node.utilization >= 0.95);
            assert_eq!(node.status, NodeStatus::HighLoad);
        }

        for link in &topology.links {
            assert!(link.is_active);
            assert!(link.bandwidth_gbps >= link.transport_type.max_bandwidth_gbps() * 0.85);
        }
    }

    #[test]
    fn test_degraded_cluster() {
        let topology = degraded_cluster();

        let degraded_count = topology
            .nodes
            .iter()
            .filter(|n| n.status == NodeStatus::Degraded)
            .count();
        assert_eq!(degraded_count, 1);

        let high_load_count = topology
            .nodes
            .iter()
            .filter(|n| n.status == NodeStatus::HighLoad)
            .count();
        assert_eq!(high_load_count, 1);
    }

    #[test]
    fn test_trace_data_flow_populated() {
        let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
        let traces = sample_collective_traces(&node_ids);

        for trace in &traces {
            assert!(!trace.data_flow.is_empty());
            assert!(trace.total_bytes > 0);
            assert!(trace.duration_us > 0);
        }
    }

    #[test]
    fn test_cluster_link_bandwidth_reasonable() {
        let topology = four_node_cluster();

        for link in &topology.links {
            // All links should be Thunderbolt with reasonable bandwidth
            assert_eq!(link.transport_type, TransportType::Thunderbolt);
            assert!(link.bandwidth_gbps > 0.0);
            assert!(link.bandwidth_gbps <= TransportType::Thunderbolt.max_bandwidth_gbps());
            assert!(link.latency_us > 0.0);
        }
    }
}