use arc_swap::ArcSwap;
use cagire_forth::{Dictionary, ExecutionTrace, Forth, Rng, StepContext, Variables};
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use parking_lot::Mutex;
use rand::rngs::StdRng;
use rand::SeedableRng;
use std::collections::HashMap;
use std::sync::Arc;

fn new_vars() -> Variables {
    Arc::new(ArcSwap::from_pointee(HashMap::new()))
}

fn new_dict() -> Dictionary {
    Arc::new(Mutex::new(HashMap::new()))
}

fn new_rng() -> Rng {
    Arc::new(Mutex::new(StdRng::seed_from_u64(42)))
}

fn new_forth() -> Forth {
    Forth::new(new_vars(), new_dict(), new_rng())
}

fn ctx() -> StepContext<'static> {
    StepContext {
        step: 3,
        beat: 7.5,
        bank: 0,
        pattern: 1,
        tempo: 120.0,
        phase: 0.75,
        slot: 0,
        runs: 5,
        iter: 5,
        speed: 1.0,
        fill: false,
        nudge_secs: 0.0,
        cc_access: None,
        speed_key: "__speed_0_1__",
        mouse_x: 0.5,
        mouse_y: 0.5,
        mouse_down: 0.0,
    }
}

const SCRIPTS: &[(&str, &str)] = &[
    ("trivial", "60"),
    ("drum_hit", "\"kick\" s ."),
    ("synth_params", "\"sine\" s 60 note 0.5 amp 2000 lpf ."),
    ("cycle_notes", "{ 60 } { 64 } { 67 } 3 cycle note"),
    (
        "conditional",
        "step 2 mod 0 = if \"kick\" s . else \"hat\" s . then",
    ),
    ("variables", "beat 4 mod \"myvar\" !"),
    (
        "multi_emit",
        "\"kick\" s . \"hat\" s 0.7 amp . \"snare\" s 0.5 amp .",
    ),
    (
        "complex",
        "{ 60 } { 64 } { 67 } { 72 } 4 cycle note \"sine\" s 0.01 0.1 0.7 0.3 adsr 2000 lpf .",
    ),
];

fn bench_evaluate(c: &mut Criterion) {
    let mut group = c.benchmark_group("evaluate");
    let step_ctx = ctx();

    for &(name, script) in SCRIPTS {
        group.bench_function(name, |b| {
            let mut forth = new_forth();
            b.iter(|| {
                let _ = black_box(forth.evaluate(black_box(script), &step_ctx));
                forth.clear_stack();
            });
        });
    }
    group.finish();
}

fn bench_evaluate_with_trace(c: &mut Criterion) {
    let mut group = c.benchmark_group("evaluate_with_trace");
    let step_ctx = ctx();

    for &(name, script) in SCRIPTS {
        group.bench_function(name, |b| {
            let mut forth = new_forth();
            let mut trace = ExecutionTrace::default();
            b.iter(|| {
                trace.executed_spans.clear();
                trace.selected_spans.clear();
                trace.resolved.clear();
                let _ = black_box(forth.evaluate_with_trace(black_box(script), &step_ctx, &mut trace));
                forth.clear_stack();
            });
        });
    }
    group.finish();
}

fn bench_trace_overhead(c: &mut Criterion) {
    let mut group = c.benchmark_group("trace_overhead");
    let step_ctx = ctx();
    let script = "{ 60 } { 64 } { 67 } { 72 } 4 cycle note \"sine\" s 0.01 0.1 0.7 0.3 adsr 2000 lpf .";

    group.bench_function("no_trace", |b| {
        let mut forth = new_forth();
        b.iter(|| {
            let _ = black_box(forth.evaluate(black_box(script), &step_ctx));
            forth.clear_stack();
        });
    });

    group.bench_function("with_trace", |b| {
        let mut forth = new_forth();
        let mut trace = ExecutionTrace::default();
        b.iter(|| {
            trace.executed_spans.clear();
            trace.selected_spans.clear();
            trace.resolved.clear();
            let _ = black_box(forth.evaluate_with_trace(black_box(script), &step_ctx, &mut trace));
            forth.clear_stack();
        });
    });
    group.finish();
}

fn bench_tick_simulation(c: &mut Criterion) {
    let mut group = c.benchmark_group("tick_simulation");

    let patterns: &[(&str, &str)] = &[
        ("drum_pattern", "\"kick\" s ."),
        ("melodic_pattern", "{ 60 } { 64 } { 67 } { 72 } 4 cycle note \"sine\" s 0.5 amp ."),
        ("generative_pattern", "1 8 rand note \"pluck\" s 0.3 0.8 rand amp ."),
        ("conditional_pattern", "step 4 mod 0 = if \"kick\" s . else step 2 mod 0 = if \"hat\" s 0.5 amp . then then"),
    ];

    for &(name, script) in patterns {
        group.bench_function(name, |b| {
            let mut forth = new_forth();
            let step_ctx = ctx();
            let mut trace = ExecutionTrace::default();
            b.iter(|| {
                trace.executed_spans.clear();
                trace.selected_spans.clear();
                trace.resolved.clear();
                let _ = black_box(forth.evaluate_with_trace(black_box(script), &step_ctx, &mut trace));
                forth.clear_stack();
            });
        });
    }
    group.finish();
}

fn bench_throughput(c: &mut Criterion) {
    let mut group = c.benchmark_group("throughput");

    let medium_scripts: &[&str] = &[
        "\"kick\" s .",
        "\"hat\" s 0.5 amp .",
        "\"snare\" s 0.8 amp 1000 lpf .",
        "{ 60 } { 64 } { 67 } 3 cycle note \"sine\" s .",
        "\"bass\" s 36 note 0.7 amp .",
        "{ 48 } { 55 } { 60 } { 64 } 4 cycle note \"pad\" s 0.3 amp 800 lpf .",
        "\"pluck\" s 72 note 0.4 amp 3000 lpf .",
        "step 4 mod 0 = if \"kick\" s . then",
    ];

    group.bench_function("8_patterns_one_step", |b| {
        let mut forths: Vec<Forth> = (0..8).map(|_| new_forth()).collect();
        let step_ctx = ctx();
        let mut trace = ExecutionTrace::default();
        b.iter(|| {
            for (i, forth) in forths.iter_mut().enumerate() {
                trace.executed_spans.clear();
                trace.selected_spans.clear();
                trace.resolved.clear();
                let _ = black_box(forth.evaluate_with_trace(
                    black_box(medium_scripts[i]),
                    &step_ctx,
                    &mut trace,
                ));
                forth.clear_stack();
            }
        });
    });
    group.finish();
}

criterion_group!(
    benches,
    bench_evaluate,
    bench_evaluate_with_trace,
    bench_trace_overhead,
    bench_tick_simulation,
    bench_throughput,
);
criterion_main!(benches);