//! Criterion benchmarks — `ScopedSymbolTable` (axhash-backed) vs a naïve
//! `Vec<std::collections::HashMap<String, SymbolData>>` baseline.
//!
//! Scenarios (simulating a compiler's name-resolution pass)
//! ─────────────────────────────────────────────────────────
//! 1. **define_only**        — Fill N symbols across a realistic scope stack.
//! 2. **resolve_hit**        — Resolve N times; every name exists (hot path).
//! 3. **resolve_miss**       — Resolve N times; name is never found (cold path).
//! 4. **resolve_deep_outer** — Name lives in the global scope; must walk all scopes.
//! 5. **mixed_workload**     — Interleaved define + resolve (type-checker loop).
//!
//! Run:
//!   cargo bench --bench symbol_table
//!
//! HTML reports → target/criterion/

use std::collections::HashMap as StdHashMap;
use std::hint::black_box;
use std::time::Duration;

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};

use ax_sym::{ScopedSymbolTable, SourceLocation, SymbolData, SymbolKind};

// ── Configuration ─────────────────────────────────────────────────────────────

const N: usize = 10_000;
const SCOPE_DEPTH: usize = 8; // realistic function + block nesting
const SYMS_PER_SCOPE: usize = 50; // ~50 locals per scope

fn loc() -> SourceLocation {
    SourceLocation::new(0, 1, 1)
}

fn configured_criterion() -> Criterion {
    Criterion::default()
        .measurement_time(Duration::from_secs(5))
        .warm_up_time(Duration::from_secs(2))
        .sample_size(100)
}

// ── Baseline (std HashMap stack) ──────────────────────────────────────────────

struct NaiveSymbolTable {
    scopes: Vec<StdHashMap<String, SymbolData>>,
}

impl NaiveSymbolTable {
    fn new() -> Self {
        Self {
            scopes: vec![StdHashMap::new()],
        }
    }
    fn push_scope(&mut self) {
        self.scopes.push(StdHashMap::new());
    }
    fn pop_scope(&mut self) {
        if self.scopes.len() > 1 {
            self.scopes.pop();
        }
    }
    fn define(&mut self, name: String, data: SymbolData) {
        self.scopes.last_mut().unwrap().insert(name, data);
    }
    fn resolve(&self, name: &str) -> Option<&SymbolData> {
        for scope in self.scopes.iter().rev() {
            if let Some(s) = scope.get(name) {
                return Some(s);
            }
        }
        None
    }
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Build the symbol names that will be used across all benchmarks.
fn build_names(n: usize) -> Vec<String> {
    (0..n).map(|i| format!("identifier_{i:06}")).collect()
}

/// Build an ax-sym table pre-loaded with `total` symbols spread across
/// `SCOPE_DEPTH` scopes.
fn build_ax_table(names: &[String]) -> ScopedSymbolTable {
    let mut t = ScopedSymbolTable::with_capacity(SCOPE_DEPTH + 1, SYMS_PER_SCOPE);
    let chunk = names.len() / SCOPE_DEPTH;
    for (i, name) in names.iter().enumerate() {
        if i % chunk == 0 && i > 0 {
            t.push_scope();
        }
        t.define(
            name.clone(),
            SymbolData::new(i as u32, SymbolKind::Variable, loc()),
        );
    }
    t
}

/// Build the naïve baseline with the same layout.
fn build_naive_table(names: &[String]) -> NaiveSymbolTable {
    let mut t = NaiveSymbolTable::new();
    let chunk = names.len() / SCOPE_DEPTH;
    for (i, name) in names.iter().enumerate() {
        if i % chunk == 0 && i > 0 {
            t.push_scope();
        }
        t.define(
            name.clone(),
            SymbolData::new(i as u32, SymbolKind::Variable, loc()),
        );
    }
    t
}

// ── 1. Define ─────────────────────────────────────────────────────────────────

fn bench_define(c: &mut Criterion) {
    let names = build_names(N);
    let chunk = N / SCOPE_DEPTH;

    let mut group = c.benchmark_group("define");
    group.throughput(Throughput::Elements(N as u64));

    group.bench_function(BenchmarkId::new("AxSymbolTable", N), |b| {
        b.iter(|| {
            let mut t = ScopedSymbolTable::with_capacity(SCOPE_DEPTH + 1, SYMS_PER_SCOPE);
            for (i, name) in names.iter().enumerate() {
                if i % chunk == 0 && i > 0 {
                    t.push_scope();
                }
                t.define(
                    name.clone(),
                    SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                );
            }
            black_box(t.total_symbols())
        });
    });

    group.bench_function(BenchmarkId::new("NaiveHashMap", N), |b| {
        b.iter(|| {
            let mut t = NaiveSymbolTable::new();
            for (i, name) in names.iter().enumerate() {
                if i % chunk == 0 && i > 0 {
                    t.push_scope();
                }
                t.define(
                    name.clone(),
                    SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                );
            }
            black_box(t.scopes.len())
        });
    });

    group.finish();
}

// ── 2. Resolve hit ────────────────────────────────────────────────────────────

fn bench_resolve_hit(c: &mut Criterion) {
    let names = build_names(N);
    let ax_table = build_ax_table(&names);
    let std_table = build_naive_table(&names);

    let mut group = c.benchmark_group("resolve_hit");
    group.throughput(Throughput::Elements(N as u64));

    group.bench_function(BenchmarkId::new("AxSymbolTable", N), |b| {
        b.iter(|| {
            let mut found = 0usize;
            for name in &names {
                if ax_table.resolve(name.as_str()).is_some() {
                    found += 1;
                }
            }
            black_box(found)
        });
    });

    group.bench_function(BenchmarkId::new("NaiveHashMap", N), |b| {
        b.iter(|| {
            let mut found = 0usize;
            for name in &names {
                if std_table.resolve(name.as_str()).is_some() {
                    found += 1;
                }
            }
            black_box(found)
        });
    });

    group.finish();
}

// ── 3. Resolve miss ───────────────────────────────────────────────────────────

fn bench_resolve_miss(c: &mut Criterion) {
    let names = build_names(N);
    // Query names that are NOT in the table.
    let missing: Vec<String> = (0..N).map(|i| format!("missing_{i:06}")).collect();

    let ax_table = build_ax_table(&names);
    let std_table = build_naive_table(&names);

    let mut group = c.benchmark_group("resolve_miss");
    group.throughput(Throughput::Elements(N as u64));

    group.bench_function(BenchmarkId::new("AxSymbolTable", N), |b| {
        b.iter(|| {
            let mut nones = 0usize;
            for name in &missing {
                if ax_table.resolve(name.as_str()).is_none() {
                    nones += 1;
                }
            }
            black_box(nones)
        });
    });

    group.bench_function(BenchmarkId::new("NaiveHashMap", N), |b| {
        b.iter(|| {
            let mut nones = 0usize;
            for name in &missing {
                if std_table.resolve(name.as_str()).is_none() {
                    nones += 1;
                }
            }
            black_box(nones)
        });
    });

    group.finish();
}

// ── 4. Resolve deep outer (worst-case scope walk) ─────────────────────────────

fn bench_resolve_deep_outer(c: &mut Criterion) {
    const DEEP: usize = 20;
    let global_sym = "global_type".to_string();

    // Pre-build the ax-sym table with the target symbol in the global scope,
    // then push DEEP additional scopes (none of which contain the symbol).
    let ax_table = {
        let mut t = ScopedSymbolTable::with_capacity(DEEP + 1, 4);
        t.define(
            global_sym.clone(),
            SymbolData::new(0, SymbolKind::Type, loc()),
        );
        for _ in 0..DEEP {
            t.push_scope();
        }
        t
    };

    let std_table = {
        let mut t = NaiveSymbolTable::new();
        t.define(
            global_sym.clone(),
            SymbolData::new(0, SymbolKind::Type, loc()),
        );
        for _ in 0..DEEP {
            t.push_scope();
        }
        t
    };

    let mut group = c.benchmark_group("resolve_deep_outer");
    group.throughput(Throughput::Elements(N as u64));

    group.bench_function(BenchmarkId::new("AxSymbolTable", N), |b| {
        b.iter(|| {
            let mut found = 0usize;
            for _ in 0..N {
                if ax_table.resolve(global_sym.as_str()).is_some() {
                    found += 1;
                }
            }
            black_box(found)
        });
    });

    group.bench_function(BenchmarkId::new("NaiveHashMap", N), |b| {
        b.iter(|| {
            let mut found = 0usize;
            for _ in 0..N {
                if std_table.resolve(global_sym.as_str()).is_some() {
                    found += 1;
                }
            }
            black_box(found)
        });
    });

    group.finish();
}

// ── 5. Mixed workload (define + resolve interleaved) ──────────────────────────

fn bench_mixed(c: &mut Criterion) {
    let names = build_names(N);

    let mut group = c.benchmark_group("mixed_workload");
    group.throughput(Throughput::Elements(N as u64));

    group.bench_function(BenchmarkId::new("AxSymbolTable", N), |b| {
        b.iter(|| {
            let mut t = ScopedSymbolTable::with_capacity(SCOPE_DEPTH + 1, SYMS_PER_SCOPE);
            let mut hits = 0usize;
            for (i, name) in names.iter().enumerate() {
                // Push a new scope every SYMS_PER_SCOPE symbols (simulate block enter).
                if i % SYMS_PER_SCOPE == 0 && i > 0 {
                    t.push_scope();
                }
                t.define(
                    name.clone(),
                    SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                );
                // Resolve the previous symbol (simulating reference right after definition).
                if i > 0 {
                    if t.resolve(names[i - 1].as_str()).is_some() {
                        hits += 1;
                    }
                }
            }
            black_box(hits)
        });
    });

    group.bench_function(BenchmarkId::new("NaiveHashMap", N), |b| {
        b.iter(|| {
            let mut t = NaiveSymbolTable::new();
            let mut hits = 0usize;
            for (i, name) in names.iter().enumerate() {
                if i % SYMS_PER_SCOPE == 0 && i > 0 {
                    t.push_scope();
                }
                t.define(
                    name.clone(),
                    SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                );
                if i > 0 {
                    if t.resolve(names[i - 1].as_str()).is_some() {
                        hits += 1;
                    }
                }
            }
            black_box(hits)
        });
    });

    group.finish();
}

// ── 6. Scope lifecycle (push → define → pop full teardown) ───────────────────
//
// Measures the full round-trip cost of entering and leaving scopes, including
// the `pop_scope` overhead (dropping all `SymbolData` entries in a scope).
// This is the critical path when a compiler exits a block: it must reclaim all
// locals before continuing with the enclosing scope.
//
// Two sub-scenarios:
//   a) **shallow** — SCOPE_DEPTH scopes, SYMS_PER_SCOPE symbols each.
//      Mirrors a typical function with a few nested blocks.
//   b) **deep**   — DEEP_SCOPE_DEPTH scopes, fewer symbols each.
//      Mirrors deeply-recursive descent or template instantiation.

fn bench_scope_lifecycle(c: &mut Criterion) {
    const DEEP_SCOPE_DEPTH: usize = 32;
    const DEEP_SYMS_PER_SCOPE: usize = 10;

    let mut group = c.benchmark_group("scope_lifecycle");

    // ── a) Shallow: realistic function body ───────────────────────────────────

    let shallow_total = SCOPE_DEPTH * SYMS_PER_SCOPE;
    let shallow_names = build_names(shallow_total);

    group.throughput(Throughput::Elements(shallow_total as u64));

    group.bench_function(
        BenchmarkId::new("AxSymbolTable/shallow", shallow_total),
        |b| {
            b.iter(|| {
                let mut t = ScopedSymbolTable::with_capacity(SCOPE_DEPTH + 1, SYMS_PER_SCOPE);
                // Phase 1: build — push scopes and fill with symbols.
                for (i, name) in shallow_names.iter().enumerate() {
                    if i % SYMS_PER_SCOPE == 0 && i > 0 {
                        t.push_scope();
                    }
                    t.define(
                        name.clone(),
                        SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                    );
                }
                // Phase 2: teardown — pop every scope until only global remains.
                while t.pop_scope() {}
                black_box(t.scope_count())
            });
        },
    );

    group.bench_function(
        BenchmarkId::new("NaiveHashMap/shallow", shallow_total),
        |b| {
            b.iter(|| {
                let mut t = NaiveSymbolTable::new();
                for (i, name) in shallow_names.iter().enumerate() {
                    if i % SYMS_PER_SCOPE == 0 && i > 0 {
                        t.push_scope();
                    }
                    t.define(
                        name.clone(),
                        SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                    );
                }
                while t.scopes.len() > 1 {
                    t.pop_scope();
                }
                black_box(t.scopes.len())
            });
        },
    );

    // ── b) Deep: many scopes, few symbols each ────────────────────────────────

    let deep_total = DEEP_SCOPE_DEPTH * DEEP_SYMS_PER_SCOPE;
    let deep_names = build_names(deep_total);

    group.throughput(Throughput::Elements(deep_total as u64));

    group.bench_function(BenchmarkId::new("AxSymbolTable/deep", deep_total), |b| {
        b.iter(|| {
            let mut t = ScopedSymbolTable::with_capacity(DEEP_SCOPE_DEPTH + 1, DEEP_SYMS_PER_SCOPE);
            for (i, name) in deep_names.iter().enumerate() {
                if i % DEEP_SYMS_PER_SCOPE == 0 && i > 0 {
                    t.push_scope();
                }
                t.define(
                    name.clone(),
                    SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                );
            }
            while t.pop_scope() {}
            black_box(t.scope_count())
        });
    });

    group.bench_function(BenchmarkId::new("NaiveHashMap/deep", deep_total), |b| {
        b.iter(|| {
            let mut t = NaiveSymbolTable::new();
            for (i, name) in deep_names.iter().enumerate() {
                if i % DEEP_SYMS_PER_SCOPE == 0 && i > 0 {
                    t.push_scope();
                }
                t.define(
                    name.clone(),
                    SymbolData::new(i as u32, SymbolKind::Variable, loc()),
                );
            }
            while t.scopes.len() > 1 {
                t.pop_scope();
            }
            black_box(t.scopes.len())
        });
    });

    group.finish();
}

// ── Entry point ───────────────────────────────────────────────────────────────

criterion_group!(
    name    = benches;
    config  = configured_criterion();
    targets = bench_define, bench_resolve_hit, bench_resolve_miss,
              bench_resolve_deep_outer, bench_mixed, bench_scope_lifecycle
);
criterion_main!(benches);
