//! Benchmarks for the radix tree implementations used in cache-aware routing.
//!
//! This benchmark tests both implementations:
//! - StringTree: Character-based tree for HTTP router (text input)
//! - TokenTree: Token-based tree for gRPC router (pre-tokenized input)
//!
//! Run with: cargo bench --bench radix_tree_benchmark
//!
//! For quick validation: cargo bench --bench radix_tree_benchmark -- benchmark_summary --exact

use std::{
    collections::BTreeMap,
    hint::black_box,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, Mutex,
    },
    thread,
    time::Instant,
};

use criterion::{criterion_group, criterion_main, Criterion};
use kv_index::{StringTree, TokenTree};
use rand::{
    distr::{Alphanumeric, SampleString},
    rng as thread_rng, Rng,
};

// Global results storage for summary
lazy_static::lazy_static! {
    static ref BENCHMARK_RESULTS: Mutex<BTreeMap<String, String>> = Mutex::new(BTreeMap::new());
}

fn add_result(category: &str, result: String) {
    let mut results = BENCHMARK_RESULTS.lock().unwrap();
    let index = results.len();
    let key = format!("{:03}_{}", index, category);
    results.insert(key, result);
}

/// Common conversation prefixes that create shared tree paths
const CONVERSATION_PREFIXES: [&str; 6] = [
    "<|system|>\nYou are a helpful assistant.\n<|user|>\n",
    "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>\n<|im_start|>user\n",
    "[INST] <<SYS>>\nYou are a helpful assistant.\n<</SYS>>\n\n",
    "Human: ",
    "User: ",
    "### Instruction:\n",
];

/// Token ID type
type TokenId = u32;

/// Generate random ASCII strings of given length
fn random_ascii_string(len: usize) -> String {
    Alphanumeric.sample_string(&mut thread_rng(), len)
}

/// Generate fixed-size strings for benchmarks
fn generate_fixed_char_strings(count: usize, char_len: usize) -> Vec<String> {
    let mut rng = thread_rng();
    (0..count)
        .map(|_| {
            let prefix_idx = rng.random_range(0..CONVERSATION_PREFIXES.len());
            let prefix = CONVERSATION_PREFIXES[prefix_idx];
            let remaining = char_len.saturating_sub(prefix.len());
            format!("{}{}", prefix, random_ascii_string(remaining))
        })
        .collect()
}

/// Generate fixed-size token sequences for size-specific benchmarks
fn generate_fixed_token_sequences(count: usize, token_len: usize) -> Vec<Vec<TokenId>> {
    let mut rng = thread_rng();
    (0..count)
        .map(|_| (0..token_len).map(|_| rng.random_range(0..50000)).collect())
        .collect()
}

/// Generate worker endpoint URLs for scaling tests
fn generate_worker_endpoints(count: usize) -> Vec<String> {
    (0..count)
        .map(|i| {
            if i % 4 == 0 {
                format!("grpc://worker-{}.sglang.svc.cluster.local:50051", i)
            } else {
                format!("http://worker-{}.sglang.svc.cluster.local:8000", i)
            }
        })
        .collect()
}

// ============================================================================
// Benchmark Macros
// ============================================================================

/// Macro for INSERT benchmarks with string-based trees (StringTree)
macro_rules! bench_string_insert {
    ($group:expr, $num_workers:expr, $char_size:expr, $workers:expr, $strings:expr,
     $prefix:literal, $category:literal, $tree_new:expr, $insert_method:ident) => {{
        let printed = Arc::new(AtomicBool::new(false));
        let bench_name = format!(concat!($prefix, "_insert_{}w_{}c"), $num_workers, $char_size);
        let workers_clone = $workers.clone();
        let strings_clone = $strings.clone();

        $group.bench_function(&bench_name, |b| {
            let workers = workers_clone.clone();
            let strings = strings_clone.clone();
            let printed = printed.clone();

            b.iter_custom(|iters| {
                let tree = $tree_new;
                let start = Instant::now();
                for i in 0..iters {
                    let tenant = &workers[i as usize % workers.len()];
                    let text = &strings[i as usize % strings.len()];
                    tree.$insert_method(black_box(text), tenant);
                }
                let duration = start.elapsed();

                if !printed.swap(true, Ordering::Relaxed) {
                    let ops_per_sec = iters as f64 / duration.as_secs_f64();
                    let latency_us = duration.as_nanos() as f64 / iters as f64 / 1000.0;
                    let throughput_mb = (ops_per_sec * $char_size as f64) / 1_000_000.0;
                    add_result(
                        $category,
                        format!(
                            "{:>3}w | {:>5} chars | INSERT | {:>8.0} ops/s | {:>6.1} µs | {:>7.1} MB/s",
                            $num_workers, $char_size, ops_per_sec, latency_us, throughput_mb
                        ),
                    );
                }

                duration
            });
        });
    }};
}

/// Macro for MATCH benchmarks with string-based trees
macro_rules! bench_string_match {
    ($group:expr, $num_workers:expr, $char_size:expr, $tree:expr, $strings:expr,
     $prefix:literal, $category:literal, $match_method:ident) => {{
        let printed = Arc::new(AtomicBool::new(false));
        let bench_name = format!(concat!($prefix, "_match_{}w_{}c"), $num_workers, $char_size);
        let tree_clone = $tree.clone();
        let strings_clone = $strings.clone();

        $group.bench_function(&bench_name, |b| {
            let tree = tree_clone.clone();
            let strings = strings_clone.clone();
            let mut idx = 0;
            let printed = printed.clone();

            b.iter_custom(|iters| {
                let start = Instant::now();
                for _ in 0..iters {
                    let result = tree.$match_method(black_box(&strings[idx % strings.len()]));
                    black_box(result);
                    idx += 1;
                }
                let duration = start.elapsed();

                if !printed.swap(true, Ordering::Relaxed) {
                    let ops_per_sec = iters as f64 / duration.as_secs_f64();
                    let latency_us = duration.as_nanos() as f64 / iters as f64 / 1000.0;
                    let throughput_mb = (ops_per_sec * $char_size as f64) / 1_000_000.0;
                    add_result(
                        $category,
                        format!(
                            "{:>3}w | {:>5} chars | MATCH  | {:>8.0} ops/s | {:>6.1} µs | {:>7.1} MB/s",
                            $num_workers, $char_size, ops_per_sec, latency_us, throughput_mb
                        ),
                    );
                }

                duration
            });
        });
    }};
}

/// Macro for CONCURRENT benchmarks with string-based trees
macro_rules! bench_string_concurrent {
    ($group:expr, $num_workers:expr, $workers:expr, $num_threads:expr, $ops_per_thread:expr,
     $prefix:literal, $category:literal, $tree_new:expr, $insert_method:ident, $match_method:ident) => {{
        let printed = Arc::new(AtomicBool::new(false));
        let bench_name = format!(concat!($prefix, "_concurrent_{}w"), $num_workers);
        let workers_clone = $workers.clone();

        $group.bench_function(&bench_name, |b| {
            let printed = printed.clone();
            let workers = workers_clone.clone();

            b.iter_custom(|iters| {
                let start = Instant::now();
                for _ in 0..iters {
                    let tree = Arc::new($tree_new);
                    let workers_ref = &workers;
                    let handles: Vec<_> = (0..$num_threads)
                        .map(|t| {
                            let tree = Arc::clone(&tree);
                            let worker = workers_ref[t % workers_ref.len()].clone();
                            thread::spawn(move || {
                                for i in 0..$ops_per_thread {
                                    let text = format!(
                                        "{}thread{}_request{}_padding{}",
                                        CONVERSATION_PREFIXES[i % CONVERSATION_PREFIXES.len()],
                                        t,
                                        i,
                                        "x".repeat(1000)
                                    );
                                    if i % 3 == 0 {
                                        tree.$match_method(&text);
                                    } else {
                                        tree.$insert_method(&text, &worker);
                                    }
                                }
                            })
                        })
                        .collect();

                    for h in handles {
                        h.join().unwrap();
                    }
                }
                let duration = start.elapsed();

                if !printed.swap(true, Ordering::Relaxed) {
                    let total_ops = iters * $num_threads as u64 * $ops_per_thread as u64;
                    let ops_per_sec = total_ops as f64 / duration.as_secs_f64();
                    add_result(
                        $category,
                        format!(
                            "{:>3}w | CONCURRENT | {:>7.0} ops/s | {} threads x {} ops",
                            $num_workers, ops_per_sec, $num_threads, $ops_per_thread
                        ),
                    );
                }

                duration
            });
        });
    }};
}

/// Macro for INSERT benchmarks with TokenTree
macro_rules! bench_token_insert {
    ($group:expr, $num_workers:expr, $token_size:expr, $workers:expr, $sequences:expr) => {{
        let printed = Arc::new(AtomicBool::new(false));
        let bench_name = format!("token_insert_{}w_{}tok", $num_workers, $token_size);
        let workers_clone = $workers.clone();
        let seqs_clone = $sequences.clone();

        $group.bench_function(&bench_name, |b| {
            let workers = workers_clone.clone();
            let seqs = seqs_clone.clone();
            let printed = printed.clone();

            b.iter_custom(|iters| {
                let tree = TokenTree::new();
                let start = Instant::now();
                for i in 0..iters {
                    let tenant = &workers[i as usize % workers.len()];
                    let tokens = &seqs[i as usize % seqs.len()];
                    tree.insert_tokens(black_box(tokens), tenant);
                }
                let duration = start.elapsed();

                if !printed.swap(true, Ordering::Relaxed) {
                    let ops_per_sec = iters as f64 / duration.as_secs_f64();
                    let latency_us = duration.as_nanos() as f64 / iters as f64 / 1000.0;
                    let throughput_mtok = (ops_per_sec * $token_size as f64) / 1_000_000.0;
                    add_result(
                        "token",
                        format!(
                            "{:>3}w | {:>5} tokens | INSERT | {:>8.0} ops/s | {:>6.1} µs | {:>6.1} Mtok/s",
                            $num_workers, $token_size, ops_per_sec, latency_us, throughput_mtok
                        ),
                    );
                }

                duration
            });
        });
    }};
}

/// Macro for MATCH benchmarks with TokenTree
macro_rules! bench_token_match {
    ($group:expr, $num_workers:expr, $token_size:expr, $tree:expr, $sequences:expr) => {{
        let printed = Arc::new(AtomicBool::new(false));
        let bench_name = format!("token_match_{}w_{}tok", $num_workers, $token_size);
        let tree_clone = $tree.clone();
        let seqs_clone = $sequences.clone();

        $group.bench_function(&bench_name, |b| {
            let tree = tree_clone.clone();
            let seqs = seqs_clone.clone();
            let mut idx = 0;
            let printed = printed.clone();

            b.iter_custom(|iters| {
                let start = Instant::now();
                for _ in 0..iters {
                    let result = tree.prefix_match_legacy(black_box(&seqs[idx % seqs.len()]));
                    black_box(result);
                    idx += 1;
                }
                let duration = start.elapsed();

                if !printed.swap(true, Ordering::Relaxed) {
                    let ops_per_sec = iters as f64 / duration.as_secs_f64();
                    let latency_us = duration.as_nanos() as f64 / iters as f64 / 1000.0;
                    let throughput_mtok = (ops_per_sec * $token_size as f64) / 1_000_000.0;
                    add_result(
                        "token",
                        format!(
                            "{:>3}w | {:>5} tokens | MATCH  | {:>8.0} ops/s | {:>6.1} µs | {:>6.1} Mtok/s",
                            $num_workers, $token_size, ops_per_sec, latency_us, throughput_mtok
                        ),
                    );
                }

                duration
            });
        });
    }};
}

/// Macro for CONCURRENT benchmarks with TokenTree
macro_rules! bench_token_concurrent {
    ($group:expr, $num_workers:expr, $workers:expr, $num_threads:expr, $ops_per_thread:expr) => {{
        let printed = Arc::new(AtomicBool::new(false));
        let bench_name = format!("token_concurrent_{}w", $num_workers);
        let workers_clone = $workers.clone();

        $group.bench_function(&bench_name, |b| {
            let printed = printed.clone();
            let workers = workers_clone.clone();

            b.iter_custom(|iters| {
                let start = Instant::now();
                for _ in 0..iters {
                    let tree = Arc::new(TokenTree::new());
                    let workers_ref = &workers;
                    let handles: Vec<_> = (0..$num_threads)
                        .map(|t| {
                            let tree = Arc::clone(&tree);
                            let worker = workers_ref[t % workers_ref.len()].clone();
                            thread::spawn(move || {
                                for i in 0..$ops_per_thread {
                                    // Generate deterministic token sequence
                                    let tokens: Vec<TokenId> = (0..1024)
                                        .map(|j| (t * 10000 + i * 100 + j) as u32)
                                        .collect();
                                    if i % 3 == 0 {
                                        tree.prefix_match_legacy(&tokens);
                                    } else {
                                        tree.insert_tokens(&tokens, &worker);
                                    }
                                }
                            })
                        })
                        .collect();

                    for h in handles {
                        h.join().unwrap();
                    }
                }
                let duration = start.elapsed();

                if !printed.swap(true, Ordering::Relaxed) {
                    let total_ops = iters * $num_threads as u64 * $ops_per_thread as u64;
                    let ops_per_sec = total_ops as f64 / duration.as_secs_f64();
                    add_result(
                        "token",
                        format!(
                            "{:>3}w | CONCURRENT | {:>7.0} ops/s | {} threads x {} ops",
                            $num_workers, ops_per_sec, $num_threads, $ops_per_thread
                        ),
                    );
                }

                duration
            });
        });
    }};
}

// ============================================================================
// Main Benchmark
// ============================================================================

/// Main benchmark for StringTree and TokenTree
fn bench_summary(c: &mut Criterion) {
    let mut group = c.benchmark_group("benchmark_summary");

    // Reduce warmup and measurement time for faster runs
    group.warm_up_time(std::time::Duration::from_millis(500));
    group.measurement_time(std::time::Duration::from_secs(2));
    group.sample_size(50);

    // Configuration constants
    const TREE_SIZE: usize = 2_000;
    const INSERT_POOL_SIZE: usize = 2_000;
    const NUM_THREADS: usize = 32;
    const OPS_PER_THREAD: usize = 100;

    // Worker counts and sizes to test (reduced for faster runs)
    const WORKER_COUNTS: [usize; 3] = [10, 100, 500];
    const TOKEN_SIZES: [usize; 3] = [1024, 4096, 16384];
    const CHAR_SIZES: [usize; 3] = [4096, 16384, 65536];

    // ========================================================================
    // StringTree Benchmark
    // ========================================================================
    for &num_workers in &WORKER_COUNTS {
        let workers = generate_worker_endpoints(num_workers);

        for &char_size in &CHAR_SIZES {
            let fixed_strings = generate_fixed_char_strings(INSERT_POOL_SIZE, char_size);

            // Pre-populate tree for MATCH
            let string_tree = Arc::new(StringTree::new());
            for (i, s) in fixed_strings.iter().take(TREE_SIZE).enumerate() {
                let tenant = &workers[i % workers.len()];
                string_tree.insert_text(s, tenant);
            }

            bench_string_insert!(
                group,
                num_workers,
                char_size,
                workers,
                fixed_strings,
                "string",
                "string",
                StringTree::new(),
                insert_text
            );
            bench_string_match!(
                group,
                num_workers,
                char_size,
                string_tree,
                fixed_strings,
                "string",
                "string",
                prefix_match_legacy
            );
        }

        bench_string_concurrent!(
            group,
            num_workers,
            workers,
            NUM_THREADS,
            OPS_PER_THREAD,
            "string",
            "string",
            StringTree::new(),
            insert_text,
            prefix_match_legacy
        );
    }

    // ========================================================================
    // TokenTree Benchmark
    // ========================================================================
    for &num_workers in &WORKER_COUNTS {
        let workers = generate_worker_endpoints(num_workers);

        for &token_size in &TOKEN_SIZES {
            let fixed_sequences = generate_fixed_token_sequences(INSERT_POOL_SIZE, token_size);

            // Pre-populate tree for MATCH
            let token_tree = Arc::new(TokenTree::new());
            for (i, seq) in fixed_sequences.iter().take(TREE_SIZE).enumerate() {
                let tenant = &workers[i % workers.len()];
                token_tree.insert_tokens(seq, tenant);
            }

            bench_token_insert!(group, num_workers, token_size, workers, fixed_sequences);
            bench_token_match!(group, num_workers, token_size, token_tree, fixed_sequences);
        }

        bench_token_concurrent!(group, num_workers, workers, NUM_THREADS, OPS_PER_THREAD);
    }

    group.finish();
}

/// Print final summary table
fn print_summary() {
    use std::io::Write;
    let _ = std::io::stdout().flush();
    let _ = std::io::stderr().flush();

    let results = BENCHMARK_RESULTS.lock().unwrap();

    // Collect results by category
    let mut string_results = Vec::new();
    let mut token_results = Vec::new();

    for (key, value) in results.iter() {
        let category = key.split('_').skip(1).collect::<Vec<_>>().join("_");
        match category.as_str() {
            "string" => string_results.push(value.clone()),
            "token" => token_results.push(value.clone()),
            _ => {}
        }
    }

    eprintln!("\n{}", "=".repeat(90));
    eprintln!("STRINGTREE (kv_index::StringTree)");
    eprintln!("{}", "=".repeat(90));
    eprintln!(
        "{:>4} | {:>12} | {:>6} | {:>10} | {:>8} | {:>12}",
        "Work", "Size", "Op", "Throughput", "Latency", "Bandwidth"
    );
    eprintln!("{}", "-".repeat(90));
    for v in &string_results {
        eprintln!("{}", v);
    }

    eprintln!("\n{}", "=".repeat(90));
    eprintln!("TOKENTREE (kv_index::TokenTree)");
    eprintln!("{}", "=".repeat(90));
    eprintln!(
        "{:>4} | {:>12} | {:>6} | {:>10} | {:>8} | {:>12}",
        "Work", "Size", "Op", "Throughput", "Latency", "Bandwidth"
    );
    eprintln!("{}", "-".repeat(90));
    for v in &token_results {
        eprintln!("{}", v);
    }

    eprintln!("\n{}", "=".repeat(90));
}

fn run_benchmarks(c: &mut Criterion) {
    bench_summary(c);
    print_summary();
}

criterion_group!(benches, run_benchmarks);
criterion_main!(benches);
