use std::collections::HashMap;

use criterion::{criterion_group, criterion_main, Criterion};
use tree_sitter_highlight::{HighlightConfiguration, Highlighter};

use kak_highlight::daemon::worker::{highlight_content, make_highlighter_config};

fn make_tokens() -> HashMap<String, String> {
    [
        "constructor",
        "function",
        "function.macro",
        "function.method",
        "keyword",
        "module",
        "operator",
        "punctuation",
        "string",
        "type",
        "variable",
        "variable.parameter",
    ]
    .into_iter()
    .fold(HashMap::new(), |mut map, val| {
        map.insert(String::from(val), String::from(val));
        map
    })
}

fn run_benchmark(c: &mut Criterion, name: &str, content: &str) {
    let timestamp = 0;
    let tokens = make_tokens();
    let names: Vec<_> = tokens.keys().collect();

    let mut highlighter = Highlighter::new();
    let highlight_config = make_highlighter_config(&names).unwrap();

    c.bench_function(name, |b| {
        b.iter(|| {
            let _ = highlight_content(
                &mut highlighter,
                &highlight_config,
                &tokens,
                content.as_bytes(),
                timestamp,
            );
        })
    });
}

fn bench_main_rs(c: &mut Criterion) {
    let content = r#"fn main() {
    println!("Hello, world!");
}
"#;

    run_benchmark(c, "main.rs", content);
}

fn bench_ast_rs(c: &mut Criterion) {
    let content = include_str!("../tree-sitter-rust/examples/ast.rs");
    run_benchmark(c, "ast.rs", content);
}

criterion_group!(benches, bench_main_rs, bench_ast_rs);
criterion_main!(benches);