use radixtree::RadixTree;
use std::time::{Duration, Instant};
use std::io::{self, Write};

// Use much smaller batches to avoid hitting OurDB's size limit
const BATCH_SIZE: usize = 1_000;
const NUM_BATCHES: usize = 1_000;  // Total records: 1,000,000
const PROGRESS_INTERVAL: usize = 100;

fn main() -> Result<(), radixtree::Error> {
    // Overall metrics
    let total_start_time = Instant::now();
    let mut total_records_inserted = 0;
    let mut batch_times = Vec::with_capacity(NUM_BATCHES);
    
    println!("Will insert up to {} records in batches of {}", 
             BATCH_SIZE * NUM_BATCHES, BATCH_SIZE);
    
    // Process in batches to avoid OurDB size limits
    for batch in 0..NUM_BATCHES {
        // Create a new database for each batch
        let batch_path = std::env::temp_dir().join(format!("radixtree_batch_{}", batch));
        
        // Clean up any existing database
        if batch_path.exists() {
            std::fs::remove_dir_all(&batch_path)?;
        }
        std::fs::create_dir_all(&batch_path)?;
        
        println!("\nBatch {}/{}: Creating new radix tree...", batch + 1, NUM_BATCHES);
        let mut tree = RadixTree::new(batch_path.to_str().unwrap(), true)?;
        
        let batch_start_time = Instant::now();
        let mut last_progress_time = Instant::now();
        let mut last_progress_count = 0;
        
        // Insert records for this batch
        for i in 0..BATCH_SIZE {
            let global_index = batch * BATCH_SIZE + i;
            let key = format!("key:{:08}", global_index);
            let value = format!("val{}", global_index).into_bytes();
            
            tree.set(&key, value)?;
            
            // Show progress at intervals
            if (i + 1) % PROGRESS_INTERVAL == 0 || i == BATCH_SIZE - 1 {
                let records_since_last = i + 1 - last_progress_count;
                let time_since_last = last_progress_time.elapsed();
                let records_per_second = records_since_last as f64 / time_since_last.as_secs_f64();
                
                print!("\rProgress: {}/{} records ({:.2}%) - {:.2} records/sec", 
                       i + 1, BATCH_SIZE, 
                       (i + 1) as f64 / BATCH_SIZE as f64 * 100.0,
                       records_per_second);
                io::stdout().flush().unwrap();
                
                last_progress_time = Instant::now();
                last_progress_count = i + 1;
            }
        }
        
        let batch_duration = batch_start_time.elapsed();
        batch_times.push(batch_duration);
        total_records_inserted += BATCH_SIZE;
        
        println!("\nBatch {}/{} completed in {:?} ({:.2} records/sec)", 
                 batch + 1, NUM_BATCHES, 
                 batch_duration,
                 BATCH_SIZE as f64 / batch_duration.as_secs_f64());
        
        // Test random access performance for this batch
        println!("Testing access performance for batch {}...", batch + 1);
        let mut total_get_time = Duration::new(0, 0);
        let num_samples = 100;
        
        // Use a simple distribution pattern
        for i in 0..num_samples {
            // Distribute samples across the batch
            let sample_id = batch * BATCH_SIZE + (i * (BATCH_SIZE / num_samples));
            let key = format!("key:{:08}", sample_id);
            
            let get_start = Instant::now();
            let _ = tree.get(&key)?;
            total_get_time += get_start.elapsed();
        }
        
        println!("Average time to retrieve a record: {:?}", 
                 total_get_time / num_samples as u32);
                 
        // Test prefix search performance
        println!("Testing prefix search performance...");
        let prefix = format!("key:{:02}", batch % 100);
        
        let list_start = Instant::now();
        let keys = tree.list(&prefix)?;
        let list_duration = list_start.elapsed();
        
        println!("Found {} keys with prefix '{}' in {:?}", 
                 keys.len(), prefix, list_duration);
    }
    
    // Overall performance summary
    let total_duration = total_start_time.elapsed();
    println!("\n\nPerformance Summary:");
    println!("Total time to insert {} records: {:?}", total_records_inserted, total_duration);
    println!("Average insertion rate: {:.2} records/second", 
             total_records_inserted as f64 / total_duration.as_secs_f64());
    
    // Show performance trend
    println!("\nPerformance Trend (batch number vs. time):");
    for (i, duration) in batch_times.iter().enumerate() {
        if i % 10 == 0 || i == batch_times.len() - 1 {  // Only show every 10th point
            println!("  Batch {}: {:?} ({:.2} records/sec)", 
                     i + 1, 
                     duration,
                     BATCH_SIZE as f64 / duration.as_secs_f64());
        }
    }
    
    Ok(())
}