use heromodels::db::{Collection, Db};
use heromodels::models::grid4::{Node, NodeDevice, ComputeSlice, StorageSlice};
use heromodels::models::grid4::node::node_index::{nodegroupid, country};
use heromodels_core::Model;

// Helper function to print node details
fn print_node_details(node: &Node) {
    println!("\n--- Node Details ---");
    println!("ID: {}", node.get_id());
    println!("NodeGroup ID: {}", node.nodegroupid);
    println!("Uptime: {}%", node.uptime);
    println!("Country: {}", node.country);
    println!("Birth Time: {}", node.birthtime);
    println!("Public Key: {}", node.pubkey);
    println!("Compute Slices: {}", node.computeslices.len());
    println!("Storage Slices: {}", node.storageslices.len());
    println!("Created At: {}", node.base_data.created_at);
    println!("Modified At: {}", node.base_data.modified_at);
    
    // Print capacity details
    println!("  Capacity:");
    println!("    Storage: {:.1} GB", node.capacity.storage_gb);
    println!("    Memory: {:.1} GB", node.capacity.mem_gb);
    println!("    GPU Memory: {:.1} GB", node.capacity.mem_gb_gpu);
    println!("    Passmark: {}", node.capacity.passmark);
    println!("    vCores: {}", node.capacity.vcores);
    
    // Print device info
    println!("  Devices:");
    println!("    Vendor: {}", node.devices.vendor);
    println!("    CPUs: {}", node.devices.cpu.len());
    println!("    GPUs: {}", node.devices.gpu.len());
    println!("    Memory: {}", node.devices.memory.len());
    println!("    Storage: {}", node.devices.storage.len());
    println!("    Network: {}", node.devices.network.len());
    
    // Print compute slices
    if !node.computeslices.is_empty() {
        println!("  Compute Slices:");
        for (i, slice) in node.computeslices.iter().enumerate() {
            println!("    {}. ID: {}, Memory: {:.1}GB, Storage: {:.1}GB, vCores: {}, GPUs: {}", 
                i + 1, slice.id, slice.mem_gb, slice.storage_gb, slice.vcores, slice.gpus);
        }
    }
    
    // Print storage slices
    if !node.storageslices.is_empty() {
        println!("  Storage Slices:");
        for (i, slice) in node.storageslices.iter().enumerate() {
            println!("    {}. ID: {}", i + 1, slice.id);
        }
    }
}

fn main() {
    // Create a new DB instance in /tmp/grid4_nodes_db, and reset before every run
    let db = heromodels::db::hero::OurDB::new("/tmp/grid4_nodes_db", true).expect("Can create DB");

    println!("Grid4 Node Models - Basic Usage Example");
    println!("======================================");

    // Create device components for nodes
    
    // CPU devices
    let cpu1 = CPUDevice {
        id: "cpu_intel_i7_12700k".to_string(),
        cores: 12,
        passmark: 28500,
        description: "Intel Core i7-12700K".to_string(),
        cpu_brand: "Intel".to_string(),
        cpu_version: "12th Gen".to_string(),
    };
    
    let cpu2 = CPUDevice {
        id: "cpu_amd_ryzen_9_5900x".to_string(),
        cores: 12,
        passmark: 32000,
        description: "AMD Ryzen 9 5900X".to_string(),
        cpu_brand: "AMD".to_string(),
        cpu_version: "Zen 3".to_string(),
    };

    // GPU devices
    let gpu1 = GPUDevice {
        id: "gpu_rtx_3080".to_string(),
        cores: 8704,
        memory_gb: 10.0,
        description: "NVIDIA GeForce RTX 3080".to_string(),
        gpu_brand: "NVIDIA".to_string(),
        gpu_version: "RTX 30 Series".to_string(),
    };
    
    let gpu2 = GPUDevice {
        id: "gpu_rtx_4090".to_string(),
        cores: 16384,
        memory_gb: 24.0,
        description: "NVIDIA GeForce RTX 4090".to_string(),
        gpu_brand: "NVIDIA".to_string(),
        gpu_version: "RTX 40 Series".to_string(),
    };

    // Memory devices
    let memory1 = MemoryDevice {
        id: "mem_ddr4_32gb".to_string(),
        size_gb: 32.0,
        description: "DDR4-3200 32GB Kit".to_string(),
    };
    
    let memory2 = MemoryDevice {
        id: "mem_ddr5_64gb".to_string(),
        size_gb: 64.0,
        description: "DDR5-5600 64GB Kit".to_string(),
    };

    // Storage devices
    let storage1 = StorageDevice {
        id: "ssd_nvme_1tb".to_string(),
        size_gb: 1000.0,
        description: "NVMe SSD 1TB".to_string(),
    };
    
    let storage2 = StorageDevice {
        id: "hdd_sata_4tb".to_string(),
        size_gb: 4000.0,
        description: "SATA HDD 4TB".to_string(),
    };

    // Network devices
    let network1 = NetworkDevice {
        id: "eth_1gbit".to_string(),
        speed_mbps: 1000,
        description: "Gigabit Ethernet".to_string(),
    };
    
    let network2 = NetworkDevice {
        id: "eth_10gbit".to_string(),
        speed_mbps: 10000,
        description: "10 Gigabit Ethernet".to_string(),
    };

    // Create device info configurations
    let devices1 = DeviceInfo {
        vendor: "Dell".to_string(),
        cpu: vec![cpu1.clone()],
        gpu: vec![gpu1.clone()],
        memory: vec![memory1.clone()],
        storage: vec![storage1.clone(), storage2.clone()],
        network: vec![network1.clone()],
    };
    
    let devices2 = DeviceInfo {
        vendor: "HP".to_string(),
        cpu: vec![cpu2.clone()],
        gpu: vec![gpu2.clone()],
        memory: vec![memory2.clone()],
        storage: vec![storage1.clone()],
        network: vec![network2.clone()],
    };

    // Create node capacities
    let capacity1 = NodeCapacity {
        storage_gb: 5000.0,
        mem_gb: 32.0,
        mem_gb_gpu: 10.0,
        passmark: 28500,
        vcores: 24,
    };
    
    let capacity2 = NodeCapacity {
        storage_gb: 1000.0,
        mem_gb: 64.0,
        mem_gb_gpu: 24.0,
        passmark: 32000,
        vcores: 24,
    };

    // Create compute slices
    let compute_slice1 = ComputeSlice::new()
        .id(1)
        .mem_gb(4.0)
        .storage_gb(100.0)
        .passmark(3000)
        .vcores(2)
        .cpu_oversubscription(150)
        .storage_oversubscription(120)
        .gpus(0);
    
    let compute_slice2 = ComputeSlice::new()
        .id(2)
        .mem_gb(8.0)
        .storage_gb(200.0)
        .passmark(6000)
        .vcores(4)
        .cpu_oversubscription(130)
        .storage_oversubscription(110)
        .gpus(1);
    
    let compute_slice3 = ComputeSlice::new()
        .id(1)
        .mem_gb(16.0)
        .storage_gb(400.0)
        .passmark(12000)
        .vcores(8)
        .cpu_oversubscription(110)
        .storage_oversubscription(100)
        .gpus(1);

    // Create storage slices
    let storage_slice1 = StorageSlice::new().id(1);
    let storage_slice2 = StorageSlice::new().id(2);
    let storage_slice3 = StorageSlice::new().id(3);

    // Create nodes with different configurations

    // Node 1 - Web hosting node
    let node1 = Node::new()
        .nodegroupid(1001)
        .uptime(98)
        .add_compute_slice(compute_slice1.clone())
        .add_compute_slice(compute_slice2.clone())
        .add_storage_slice(storage_slice1.clone())
        .add_storage_slice(storage_slice2.clone())
        .devices(devices1.clone())
        .country("US".to_string())
        .capacity(capacity1.clone())
        .birthtime(1640995200) // 2022-01-01
        .pubkey("node1_pubkey_abc123xyz789".to_string())
        .signature_node("node1_signature_def456".to_string())
        .signature_farmer("farmer1_signature_ghi789".to_string());

    // Node 2 - High-performance computing node
    let node2 = Node::new()
        .nodegroupid(1002)
        .uptime(99)
        .add_compute_slice(compute_slice3.clone())
        .add_storage_slice(storage_slice3.clone())
        .devices(devices2.clone())
        .country("DE".to_string())
        .capacity(capacity2.clone())
        .birthtime(1672531200) // 2023-01-01
        .pubkey("node2_pubkey_jkl012mno345".to_string())
        .signature_node("node2_signature_pqr678".to_string())
        .signature_farmer("farmer2_signature_stu901".to_string());

    // Node 3 - Storage-focused node
    let node3 = Node::new()
        .nodegroupid(1001)
        .uptime(95)
        .add_storage_slice(storage_slice1.clone())
        .add_storage_slice(storage_slice2.clone())
        .add_storage_slice(storage_slice3.clone())
        .devices(devices1.clone())
        .country("NL".to_string())
        .capacity(capacity1.clone())
        .birthtime(1704067200) // 2024-01-01
        .pubkey("node3_pubkey_vwx234yzab567".to_string())
        .signature_node("node3_signature_cde890".to_string())
        .signature_farmer("farmer1_signature_fgh123".to_string());

    // Save all nodes to database and get their assigned IDs and updated models
    let (node1_id, db_node1) = db
        .collection()
        .expect("can open node collection")
        .set(&node1)
        .expect("can set node");
    let (node2_id, db_node2) = db
        .collection()
        .expect("can open node collection")
        .set(&node2)
        .expect("can set node");
    let (node3_id, db_node3) = db
        .collection()
        .expect("can open node collection")
        .set(&node3)
        .expect("can set node");

    println!("Node 1 assigned ID: {}", node1_id);
    println!("Node 2 assigned ID: {}", node2_id);
    println!("Node 3 assigned ID: {}", node3_id);

    // Print all nodes retrieved from database
    println!("\n--- Nodes Retrieved from Database ---");
    println!("\n1. Web hosting node:");
    print_node_details(&db_node1);

    println!("\n2. High-performance computing node:");
    print_node_details(&db_node2);

    println!("\n3. Storage-focused node:");
    print_node_details(&db_node3);

    // Demonstrate different ways to retrieve nodes from the database

    // 1. Retrieve by nodegroup ID index
    println!("\n--- Retrieving Nodes by Different Methods ---");
    println!("\n1. By NodeGroup ID Index (NodeGroup 1001):");
    let nodegroup_nodes = db
        .collection::<Node>()
        .expect("can open node collection")
        .get::<nodegroupid, _>(&1001i32)
        .expect("can load nodes by nodegroup");

    assert_eq!(nodegroup_nodes.len(), 2);
    for (i, node) in nodegroup_nodes.iter().enumerate() {
        println!("  Node {}: ID {}, Country: {}, Uptime: {}%", 
            i + 1, node.get_id(), node.country, node.uptime);
    }

    // 2. Retrieve by country index
    println!("\n2. By Country Index (Germany - DE):");
    let country_nodes = db
        .collection::<Node>()
        .expect("can open node collection")
        .get::<country, _>("DE")
        .expect("can load nodes by country");

    assert_eq!(country_nodes.len(), 1);
    print_node_details(&country_nodes[0]);

    // 3. Update node uptime
    println!("\n3. Updating Node Uptime:");
    let mut updated_node = db_node1.clone();
    updated_node.uptime = 99;

    let (_, uptime_updated_node) = db
        .collection::<Node>()
        .expect("can open node collection")
        .set(&updated_node)
        .expect("can update node");

    println!("Updated node uptime to 99%:");
    println!("  Node ID: {}, New Uptime: {}%", uptime_updated_node.get_id(), uptime_updated_node.uptime);

    // Show all nodes and calculate analytics
    let all_nodes = db
        .collection::<Node>()
        .expect("can open node collection")
        .get_all()
        .expect("can load all nodes");

    println!("\n--- Node Analytics ---");
    println!("Total Nodes: {}", all_nodes.len());

    // Calculate total capacity
    let total_storage_gb: f64 = all_nodes.iter().map(|n| n.capacity.storage_gb).sum();
    let total_memory_gb: f64 = all_nodes.iter().map(|n| n.capacity.mem_gb).sum();
    let total_gpu_memory_gb: f64 = all_nodes.iter().map(|n| n.capacity.mem_gb_gpu).sum();
    let total_vcores: i32 = all_nodes.iter().map(|n| n.capacity.vcores).sum();
    let avg_uptime: f64 = all_nodes.iter().map(|n| n.uptime as f64).sum::<f64>() / all_nodes.len() as f64;

    println!("Total Capacity:");
    println!("  Storage: {:.1} GB", total_storage_gb);
    println!("  Memory: {:.1} GB", total_memory_gb);
    println!("  GPU Memory: {:.1} GB", total_gpu_memory_gb);
    println!("  vCores: {}", total_vcores);
    println!("  Average Uptime: {:.1}%", avg_uptime);

    // Count nodes by country
    let mut country_counts = std::collections::HashMap::new();
    for node in &all_nodes {
        *country_counts.entry(&node.country).or_insert(0) += 1;
    }

    println!("\nNodes by Country:");
    for (country, count) in country_counts {
        println!("  {}: {}", country, count);
    }

    // Count total slices
    let total_compute_slices: usize = all_nodes.iter().map(|n| n.computeslices.len()).sum();
    let total_storage_slices: usize = all_nodes.iter().map(|n| n.storageslices.len()).sum();

    println!("\nTotal Slices:");
    println!("  Compute Slices: {}", total_compute_slices);
    println!("  Storage Slices: {}", total_storage_slices);

    // Vendor distribution
    let mut vendor_counts = std::collections::HashMap::new();
    for node in &all_nodes {
        *vendor_counts.entry(&node.devices.vendor).or_insert(0) += 1;
    }

    println!("\nNodes by Vendor:");
    for (vendor, count) in vendor_counts {
        println!("  {}: {}", vendor, count);
    }

    println!("\n--- Model Information ---");
    println!("Node DB Prefix: {}", Node::db_prefix());
}