use tst::TST;
use std::time::Instant;
use std::io::{self, Write};

fn main() -> Result<(), tst::Error> {
    // Create a temporary directory for the database
    let db_path = std::env::temp_dir().join("tst_prefix_example");
    std::fs::create_dir_all(&db_path)?;
    
    println!("Creating ternary search tree at: {}", db_path.display());
    
    // Create a new TST
    let mut tree = TST::new(db_path.to_str().unwrap(), true)?;
    
    // Insert a variety of keys with different prefixes
    println!("Inserting data with various prefixes...");
    
    // Names
    let names = [
        "Alice", "Alexander", "Amanda", "Andrew", "Amy",
        "Bob", "Barbara", "Benjamin", "Brenda", "Brian",
        "Charlie", "Catherine", "Christopher", "Cynthia", "Carl",
        "David", "Diana", "Daniel", "Deborah", "Donald",
        "Edward", "Elizabeth", "Eric", "Emily", "Ethan"
    ];
    
    for (i, name) in names.iter().enumerate() {
        let value = format!("person-{}", i).into_bytes();
        tree.set(name, value)?;
    }
    
    // Cities
    let cities = [
        "New York", "Los Angeles", "Chicago", "Houston", "Phoenix",
        "Philadelphia", "San Antonio", "San Diego", "Dallas", "San Jose",
        "Austin", "Jacksonville", "Fort Worth", "Columbus", "San Francisco",
        "Charlotte", "Indianapolis", "Seattle", "Denver", "Washington"
    ];
    
    for (i, city) in cities.iter().enumerate() {
        let value = format!("city-{}", i).into_bytes();
        tree.set(city, value)?;
    }
    
    // Countries
    let countries = [
        "United States", "Canada", "Mexico", "Brazil", "Argentina",
        "United Kingdom", "France", "Germany", "Italy", "Spain",
        "China", "Japan", "India", "Australia", "Russia"
    ];
    
    for (i, country) in countries.iter().enumerate() {
        let value = format!("country-{}", i).into_bytes();
        tree.set(country, value)?;
    }
    
    println!("Total items inserted: {}", names.len() + cities.len() + countries.len());
    
    // Test prefix operations
    test_prefix(&mut tree, "A")?;
    test_prefix(&mut tree, "B")?;
    test_prefix(&mut tree, "C")?;
    test_prefix(&mut tree, "San")?;
    test_prefix(&mut tree, "United")?;
    
    // Test non-existent prefix
    test_prefix(&mut tree, "Z")?;
    
    // Test empty prefix (should return all keys)
    println!("\nTesting empty prefix (should return all keys):");
    let start = Instant::now();
    let all_keys = tree.list("")?;
    let duration = start.elapsed();
    
    println!("Found {} keys with empty prefix in {:?}", all_keys.len(), duration);
    println!("First 5 keys (alphabetically):");
    for key in all_keys.iter().take(5) {
        println!("  {}", key);
    }
    
    // Clean up (optional)
    if std::env::var("KEEP_DB").is_err() {
        std::fs::remove_dir_all(&db_path)?;
        println!("\nCleaned up database directory");
    } else {
        println!("\nDatabase kept at: {}", db_path.display());
    }
    
    Ok(())
}

fn test_prefix(tree: &mut TST, prefix: &str) -> Result<(), tst::Error> {
    println!("\nTesting prefix '{}':", prefix);
    
    // Test list operation
    let start = Instant::now();
    let keys = tree.list(prefix)?;
    let list_duration = start.elapsed();
    
    println!("Found {} keys with prefix '{}' in {:?}", keys.len(), prefix, list_duration);
    
    if !keys.is_empty() {
        println!("Keys:");
        for key in &keys {
            println!("  {}", key);
        }
        
        // Test getall operation
        let start = Instant::now();
        let values = tree.getall(prefix)?;
        let getall_duration = start.elapsed();
        
        println!("Retrieved {} values in {:?}", values.len(), getall_duration);
        println!("First value: {}", 
                 if !values.is_empty() { 
                     String::from_utf8_lossy(&values[0]) 
                 } else { 
                     "None".into() 
                 });
    }
    
    Ok(())
}