//! Benchmarks for the radix tree implementation used in cache-aware routing.
//!
//! Run with: cargo bench --bench tree_benchmark

use std::{sync::Arc, thread};

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use rand::{
    distr::{Alphanumeric, SampleString},
    rng as thread_rng,
};
// Import the tree module
use sgl_model_gateway::policies::tree::Tree;

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

/// Generate random strings with common prefixes (simulates real request patterns)
fn random_prefixed_strings(prefix: &str, suffix_len: usize, count: usize) -> Vec<String> {
    (0..count)
        .map(|_| format!("{}{}", prefix, random_ascii_string(suffix_len)))
        .collect()
}

/// Benchmark single-threaded insert throughput
fn bench_insert_throughput(c: &mut Criterion) {
    let mut group = c.benchmark_group("insert_throughput");

    for text_len in [10, 50, 100, 500].iter() {
        group.throughput(Throughput::Elements(1));
        group.bench_with_input(
            BenchmarkId::new("random_text", text_len),
            text_len,
            |b, &len| {
                let tree = Tree::new();
                let strings: Vec<String> = (0..1000).map(|_| random_ascii_string(len)).collect();
                let mut idx = 0;

                b.iter(|| {
                    tree.insert(black_box(&strings[idx % strings.len()]), "tenant1");
                    idx += 1;
                });
            },
        );
    }

    // Benchmark with shared prefixes (common cache scenario)
    group.bench_function("shared_prefix_100", |b| {
        let tree = Tree::new();
        let prefixes = ["system:", "user:", "assistant:", "tool:"];
        let strings: Vec<String> = prefixes
            .iter()
            .flat_map(|p| random_prefixed_strings(p, 50, 250))
            .collect();
        let mut idx = 0;

        b.iter(|| {
            tree.insert(black_box(&strings[idx % strings.len()]), "tenant1");
            idx += 1;
        });
    });

    group.finish();
}

/// Benchmark prefix_match latency
fn bench_prefix_match_latency(c: &mut Criterion) {
    let mut group = c.benchmark_group("prefix_match_latency");

    // Setup: pre-populate tree with data
    let tree = Tree::new();
    let prefixes = ["system:", "user:", "assistant:", "tool:"];
    let strings: Vec<String> = prefixes
        .iter()
        .flat_map(|p| random_prefixed_strings(p, 50, 1000))
        .collect();

    for s in &strings {
        tree.insert(s, "tenant1");
    }

    // Benchmark cache hit (exact match)
    group.bench_function("cache_hit", |b| {
        let mut idx = 0;
        b.iter(|| {
            let result = tree.prefix_match(black_box(&strings[idx % strings.len()]));
            idx += 1;
            result
        });
    });

    // Benchmark cache miss (no match)
    let miss_strings: Vec<String> = (0..1000).map(|_| random_ascii_string(50)).collect();
    group.bench_function("cache_miss", |b| {
        let mut idx = 0;
        b.iter(|| {
            let result = tree.prefix_match(black_box(&miss_strings[idx % miss_strings.len()]));
            idx += 1;
            result
        });
    });

    // Benchmark partial match
    group.bench_function("partial_match", |b| {
        let partial_strings: Vec<String> = prefixes
            .iter()
            .map(|p| format!("{}partial_query", p))
            .collect();
        let mut idx = 0;
        b.iter(|| {
            let result =
                tree.prefix_match(black_box(&partial_strings[idx % partial_strings.len()]));
            idx += 1;
            result
        });
    });

    group.finish();
}

/// Benchmark concurrent operations
fn bench_concurrent_operations(c: &mut Criterion) {
    let mut group = c.benchmark_group("concurrent");
    group.sample_size(50); // Reduce sample size for concurrent tests

    // Mixed read/write workload
    for num_threads in [2, 4, 8].iter() {
        group.bench_with_input(
            BenchmarkId::new("mixed_workload", num_threads),
            num_threads,
            |b, &threads| {
                b.iter(|| {
                    let tree = Arc::new(Tree::new());
                    let handles: Vec<_> = (0..threads)
                        .map(|t| {
                            let tree = Arc::clone(&tree);
                            thread::spawn(move || {
                                let tenant = format!("tenant{}", t);
                                for i in 0..100 {
                                    let text = format!("thread{}_request{}", t, i);
                                    if i % 3 == 0 {
                                        tree.prefix_match(&text);
                                    } else {
                                        tree.insert(&text, &tenant);
                                    }
                                }
                            })
                        })
                        .collect();

                    for h in handles {
                        h.join().unwrap();
                    }
                });
            },
        );
    }

    group.finish();
}

/// Benchmark eviction performance
fn bench_eviction(c: &mut Criterion) {
    let mut group = c.benchmark_group("eviction");
    group.sample_size(20); // Eviction is expensive

    for tree_size in [1000, 5000, 10000].iter() {
        group.bench_with_input(
            BenchmarkId::new("evict_to_half", tree_size),
            tree_size,
            |b, &size| {
                b.iter_with_setup(
                    || {
                        // Setup: create tree with many entries
                        let tree = Tree::new();
                        for i in 0..size {
                            tree.insert(&format!("entry_{:05}", i), "tenant1");
                        }
                        tree
                    },
                    |tree| {
                        // Evict to half size
                        tree.evict_tenant_by_size(size / 2);
                    },
                );
            },
        );
    }

    group.finish();
}

/// Benchmark UTF-8 handling vs ASCII
fn bench_utf8_vs_ascii(c: &mut Criterion) {
    let mut group = c.benchmark_group("encoding");

    let tree_ascii = Tree::new();
    let tree_utf8 = Tree::new();

    // Pre-populate
    let ascii_strings: Vec<String> = (0..1000).map(|_| random_ascii_string(50)).collect();
    let utf8_strings: Vec<String> = (0..1000).map(|i| format!("你好世界_{}", i)).collect();

    for s in &ascii_strings {
        tree_ascii.insert(s, "tenant1");
    }
    for s in &utf8_strings {
        tree_utf8.insert(s, "tenant1");
    }

    group.bench_function("ascii_match", |b| {
        let mut idx = 0;
        b.iter(|| {
            let result =
                tree_ascii.prefix_match(black_box(&ascii_strings[idx % ascii_strings.len()]));
            idx += 1;
            result
        });
    });

    group.bench_function("utf8_match", |b| {
        let mut idx = 0;
        b.iter(|| {
            let result = tree_utf8.prefix_match(black_box(&utf8_strings[idx % utf8_strings.len()]));
            idx += 1;
            result
        });
    });

    group.finish();
}

/// Benchmark multi-tenant scenarios
fn bench_multi_tenant(c: &mut Criterion) {
    let mut group = c.benchmark_group("multi_tenant");

    let tree = Tree::new();

    // Setup: multiple tenants with overlapping data
    let tenants = ["worker1", "worker2", "worker3", "worker4"];
    let prefixes = ["prompt:", "completion:", "context:"];

    for tenant in &tenants {
        for prefix in &prefixes {
            for i in 0..100 {
                tree.insert(&format!("{}data_{}", prefix, i), tenant);
            }
        }
    }

    group.bench_function("shared_prefix_lookup", |b| {
        let queries: Vec<String> = prefixes
            .iter()
            .flat_map(|p| (0..10).map(move |i| format!("{}data_{}", p, i)))
            .collect();
        let mut idx = 0;

        b.iter(|| {
            let result = tree.prefix_match(black_box(&queries[idx % queries.len()]));
            idx += 1;
            result
        });
    });

    group.bench_function("tenant_specific_match", |b| {
        let queries: Vec<(String, &str)> = tenants
            .iter()
            .flat_map(|&t| (0..10).map(move |i| (format!("prompt:data_{}", i), t)))
            .collect();
        let mut idx = 0;

        b.iter(|| {
            let (query, tenant) = &queries[idx % queries.len()];
            let result = tree.prefix_match_tenant(black_box(query), black_box(tenant));
            idx += 1;
            result
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_insert_throughput,
    bench_prefix_match_latency,
    bench_concurrent_operations,
    bench_eviction,
    bench_utf8_vs_ascii,
    bench_multi_tenant,
);
criterion_main!(benches);