use criterion::{
    BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
    measurement::WallTime,
};
use pyo3::types::PyAnyMethods;
use rustpython_compiler::Mode;
use rustpython_vm::{AsObject, Interpreter, PyResult, Settings, scope::Scope};
use std::{fs, io, path::Path};

// List of microbenchmarks to skip.
//
// These result in excessive memory usage, some more so than others. For example, while
// exception_context.py consumes a lot of memory, it still finishes. On the other hand,
// call_kwargs.py seems like it performs an excessive amount of allocations and results in
// a system freeze.
// In addition, the fact that we don't yet have a GC means that benchmarks which might consume
// a bearable amount of memory accumulate. As such, best to skip them for now.
const SKIP_MICROBENCHMARKS: [&str; 8] = [
    "call_simple.py",
    "call_kwargs.py",
    "construct_object.py",
    "define_function.py",
    "define_class.py",
    "exception_nested.py",
    "exception_simple.py",
    "exception_context.py",
];

pub struct MicroBenchmark {
    name: String,
    setup: String,
    code: String,
    iterate: bool,
}

fn bench_cpython_code(group: &mut BenchmarkGroup<'_, WallTime>, bench: &MicroBenchmark) {
    pyo3::Python::attach(|py| {
        let setup_name = format!("{}_setup", bench.name);
        let setup_code = cpy_compile_code(py, &bench.setup, &setup_name).unwrap();

        let code = cpy_compile_code(py, &bench.code, &bench.name).unwrap();

        // Grab the exec function in advance so we don't have lookups in the hot code
        let builtins =
            pyo3::types::PyModule::import(py, "builtins").expect("Failed to import builtins");
        let exec = builtins.getattr("exec").expect("no exec in builtins");

        let bench_func = |globals: &mut pyo3::Bound<'_, pyo3::types::PyDict>| {
            let res = exec.call((&code, &*globals), None);
            if let Err(e) = res {
                e.print(py);
                panic!("Error running microbenchmark")
            }
        };

        let bench_setup = |iterations| {
            let globals = pyo3::types::PyDict::new(py);
            globals.set_item("__name__", "__main__").unwrap();
            if let Some(idx) = iterations {
                globals.set_item("ITERATIONS", idx).unwrap();
            }

            let res = exec.call((&setup_code, &globals), None);
            if let Err(e) = res {
                e.print(py);
                panic!("Error running microbenchmark setup code")
            }
            globals
        };

        if bench.iterate {
            for idx in iteration_counts() {
                group.throughput(Throughput::Elements(idx as u64));
                let id = BenchmarkId::new("cpython", format!("{}[{idx}]", bench.name));
                group.bench_with_input(id, &idx, |b, idx| {
                    b.iter_batched_ref(
                        || bench_setup(Some(*idx)),
                        bench_func,
                        BatchSize::LargeInput,
                    );
                });
            }
        } else {
            group.bench_function(BenchmarkId::new("cpython", &bench.name), move |b| {
                b.iter_batched_ref(|| bench_setup(None), bench_func, BatchSize::LargeInput);
            });
        }
    })
}

fn cpy_compile_code<'a>(
    py: pyo3::Python<'a>,
    code: &str,
    name: &str,
) -> pyo3::PyResult<pyo3::Bound<'a, pyo3::types::PyCode>> {
    let builtins =
        pyo3::types::PyModule::import(py, "builtins").expect("Failed to import builtins");
    let compile = builtins.getattr("compile").expect("no compile in builtins");
    Ok(compile
        .call1((code, name, "exec"))?
        .cast_into()
        .expect("compile() should return a code object"))
}

fn bench_rustpython_code(group: &mut BenchmarkGroup<'_, WallTime>, bench: &MicroBenchmark) {
    let mut settings = Settings::default();
    settings.path_list.push("Lib/".to_string());
    settings.write_bytecode = false;
    settings.user_site_directory = false;
    // Pin the str/bytes hash secret so dict/set probe sequences (and thus
    // timings) are reproducible across benchmark runs.
    settings.hash_seed = Some(0);

    let builder = Interpreter::builder(settings);
    let defs = rustpython_stdlib::stdlib_module_defs(&builder.ctx);
    let interp = builder.add_native_modules(&defs).build();
    interp.enter(|vm| {
        let setup_code = vm
            .compile(&bench.setup, Mode::Exec, &bench.name)
            .expect("Error compiling setup code");
        let bench_code = vm
            .compile(&bench.code, Mode::Exec, &bench.name)
            .expect("Error compiling bench code");

        let bench_func = |scope: &mut Scope| {
            // Keep setup objects alive until after timing, as in the CPython benchmark.
            let res: PyResult = vm.run_code_obj(bench_code.clone(), scope.clone());
            vm.unwrap_pyresult(res);
        };

        let bench_setup = |iterations| {
            let scope = vm.new_scope_with_builtins();
            scope
                .globals
                .set_item("__name__", vm.ctx.new_str("__main__").into(), vm)
                .unwrap();
            if let Some(idx) = iterations {
                scope
                    .locals
                    .as_ref()
                    .expect("new_scope_with_builtins always provides locals")
                    .as_object()
                    .set_item("ITERATIONS", vm.new_pyobj(idx), vm)
                    .expect("Error adding ITERATIONS local variable");
            }
            let setup_result = vm.run_code_obj(setup_code.clone(), scope.clone());
            vm.unwrap_pyresult(setup_result);
            scope
        };

        if bench.iterate {
            for idx in iteration_counts() {
                group.throughput(Throughput::Elements(idx as u64));
                let name = if is_codspeed() {
                    bench.name.clone()
                } else {
                    format!("{}[{idx}]", bench.name)
                };
                group.bench_with_input(BenchmarkId::new("rustpython", name), &idx, |b, idx| {
                    b.iter_batched_ref(
                        || bench_setup(Some(*idx)),
                        bench_func,
                        BatchSize::LargeInput,
                    );
                });
            }
        } else {
            group.bench_function(BenchmarkId::new("rustpython", &bench.name), move |b| {
                b.iter_batched_ref(|| bench_setup(None), bench_func, BatchSize::LargeInput);
            });
        }
    })
}

/// `true` when the benchmarks are executed by the CodSpeed runner.
///
/// CodSpeed tracks the performance of RustPython itself, so the CPython
/// reference benchmarks are skipped there: they double the (already slow)
/// instrumented run without ever reporting a change of RustPython.
fn is_codspeed() -> bool {
    std::env::var_os("CODSPEED_ENV").is_some()
}

/// The `ITERATIONS` values the benchmarks referencing them are run with.
///
/// A single size is used under CodSpeed: the criterion benchmark id does not
/// include the iteration count, so every size would be reported under the same
/// name.
fn iteration_counts() -> Vec<i32> {
    if is_codspeed() {
        vec![1_000]
    } else {
        (100..=1_000).step_by(200).collect()
    }
}

pub fn run_micro_benchmark(c: &mut Criterion, benchmark: MicroBenchmark) {
    let mut group = c.benchmark_group("microbenchmarks");

    if !is_codspeed() {
        bench_cpython_code(&mut group, &benchmark);
    }
    bench_rustpython_code(&mut group, &benchmark);

    group.finish();
}

pub fn criterion_benchmark(c: &mut Criterion) {
    let benchmark_dir = Path::new("./benches/microbenchmarks/");
    let mut paths: Vec<_> = benchmark_dir
        .read_dir()
        .unwrap()
        .collect::<io::Result<Vec<_>>>()
        .unwrap()
        .into_iter()
        .map(|entry| entry.path())
        .filter(|path| path.is_file() && path.extension().is_some_and(|ext| ext == "py"))
        .collect();
    paths.sort();

    let benchmarks: Vec<MicroBenchmark> = paths
        .iter()
        .map(|p| {
            let name = p.file_name().unwrap().to_os_string();
            let contents = fs::read_to_string(p).unwrap();
            let iterate = contents.contains("ITERATIONS");

            let (setup, code) = if contents.contains("# ---") {
                let split: Vec<&str> = contents.splitn(2, "# ---").collect();
                (split[0].to_string(), split[1].to_string())
            } else {
                ("".to_string(), contents)
            };
            let name = name.into_string().unwrap();
            MicroBenchmark {
                name,
                setup,
                code,
                iterate,
            }
        })
        .collect();

    for benchmark in benchmarks {
        if SKIP_MICROBENCHMARKS.contains(&benchmark.name.as_str()) {
            continue;
        }
        run_micro_benchmark(c, benchmark);
    }
}

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