use std::env;
use std::fs;
use std::hint::black_box;
use std::time::Instant;
use tree_sitter::{Node, Parser, Query, QueryCursor};

fn char_to_token_est(chars: usize) -> usize {
    (chars + 3) / 4
}

fn extract_skeleton(root: Node, src: &str, q: &Query) -> (String, usize) {
    let mut out = String::new();
    let mut count = 0usize;
    let mut cursor = QueryCursor::new();
    let matches = cursor.matches(q, root, src.as_bytes());
    for mat in matches {
        for cap in mat.captures {
            let n = cap.node;
            let start_byte = n.start_byte();
            let header_end = {
                let from_start = &src[start_byte..];
                let mut depth: i32 = 0;
                let mut end = from_start.len();
                for (i, b) in from_start.char_indices() {
                    match b {
                        '{' => {
                            if depth == 0 {
                                end = i;
                                break;
                            }
                            depth += 1;
                        }
                        ';' => {
                            end = i;
                            break;
                        }
                        _ => {}
                    }
                    let _ = depth;
                }
                start_byte + end.min(from_start.len())
            };
            let header = &src[start_byte..header_end].trim_end();
            out.push_str(header);
            out.push_str(" { ... }\n");
            count += 1;
        }
    }
    (out, count)
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let path = args.get(1).expect("usage: ts-bench <file.rs>");
    let src = fs::read_to_string(path).expect("read");
    let src_chars = src.chars().count();
    let src_tokens_est = char_to_token_est(src.len());

    let mut parser = Parser::new();
    parser.set_language(&tree_sitter_rust::LANGUAGE.into()).unwrap();

    // warmup
    let _ = parser.parse(&src, None).unwrap();

    // ---- PURE PARSE timing ----
    let n = 200;
    let mut times_us = Vec::with_capacity(n);
    for _ in 0..n {
        let t0 = Instant::now();
        let tree = parser.parse(black_box(&src), None).unwrap();
        let t1 = Instant::now();
        drop(tree);
        times_us.push(t1.duration_since(t0).as_nanos() as f64 / 1000.0);
    }
    times_us.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let min = times_us[0];
    let med = times_us[n / 2];
    let p95 = times_us[(n as f64 * 0.95) as usize];
    let mean = times_us.iter().sum::<f64>() / n as f64;

    // ---- INCREMENTAL PARSE (edit one byte at end) ----
    let mut src_edited = src.clone();
    src_edited.push_str("\n// comment\n");
    let tree0 = parser.parse(&src, None).unwrap();
    let mut times_inc = Vec::with_capacity(n);
    for _ in 0..n {
        let t0 = Instant::now();
        let t2 = parser.parse(&src_edited, Some(&tree0)).unwrap();
        let t1 = Instant::now();
        drop(t2);
        times_inc.push(t1.duration_since(t0).as_nanos() as f64 / 1000.0);
    }
    times_inc.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let inc_med = times_inc[n / 2];

    // ---- Query compilation timing ----
    let query_text = r#"
      [
        (function_item name: (identifier) @name)
        (struct_item name: (type_identifier) @name)
        (enum_item name: (type_identifier) @name)
        (trait_item name: (type_identifier) @name)
        (impl_item trait: (type_identifier)? @trait type: (type_identifier)? @type)
        (const_item name: (identifier) @name)
        (static_item name: (identifier) @name)
        (type_item name: (type_identifier) @name)
        (mod_item name: (identifier) @name)
        (macro_definition name: (identifier) @name)
      ]
    "#;
    let mut qtimes = Vec::with_capacity(n);
    for _ in 0..n {
        let t0 = Instant::now();
        let q = Query::new(&tree_sitter_rust::LANGUAGE.into(), query_text).unwrap();
        let t1 = Instant::now();
        drop(q);
        qtimes.push(t1.duration_since(t0).as_nanos() as f64 / 1000.0);
    }
    qtimes.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let q_med = qtimes[n / 2];

    // ---- Skeleton extraction (compiled query reused) ----
    let tree = parser.parse(&src, None).unwrap();
    let root = tree.root_node();
    let q = Query::new(&tree_sitter_rust::LANGUAGE.into(), query_text).unwrap();

    let m = 200;
    let mut ext_us = Vec::with_capacity(m);
    let mut last_skel = String::new();
    let mut last_count = 0;
    for _ in 0..m {
        let t0 = Instant::now();
        let (skel, cnt) = extract_skeleton(root, &src, &q);
        let t1 = Instant::now();
        ext_us.push(t1.duration_since(t0).as_nanos() as f64 / 1000.0);
        last_skel = skel;
        last_count = cnt;
    }
    ext_us.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let ext_med = ext_us[m / 2];

    let skel_chars = last_skel.chars().count();
    let skel_tokens_est = char_to_token_est(last_skel.len());

    println!("=== FILE ===");
    println!("path              : {}", path);
    println!("lines             : {}", src.lines().count());
    println!("chars             : {}", src_chars);
    println!("bytes             : {}", src.len());
    println!("tokens est (~/4c) : {}", src_tokens_est);
    println!();
    println!("=== PARSE SPEED (pure tree-sitter-rust, release, {} iters) ===", n);
    println!("min={:.2}us  median={:.2}us  p95={:.2}us  mean={:.2}us", min, med, p95, mean);
    println!("median ms         : {:.3} ms", med / 1000.0);
    println!();
    println!("=== INCREMENTAL PARSE (small edit) ===");
    println!("median            : {:.2}us ({:.3}ms)", inc_med, inc_med / 1000.0);
    println!();
    println!("=== QUERY COMPILATION (once, should be cached) ===");
    println!("median            : {:.2}us ({:.3}ms)", q_med, q_med / 1000.0);
    println!();
    println!("=== SKELETON EXTRACTION (compiled query reused) ===");
    println!("symbols extracted : {}", last_count);
    println!("skeleton chars    : {}", skel_chars);
    println!("skeleton tokens   : {}", skel_tokens_est);
    println!("extract median    : {:.2}us ({:.3}ms)", ext_med, ext_med / 1000.0);
    println!();
    println!("=== DENSITY (full vs skeleton) ===");
    println!(
        "chars   : {:.1}x  ({:.1}% reduction)",
        src_chars as f64 / skel_chars as f64,
        (1.0 - skel_chars as f64 / src_chars as f64) * 100.0
    );
    println!(
        "tokens  : {:.1}x  ({:.1}% reduction)",
        src_tokens_est as f64 / skel_tokens_est as f64,
        (1.0 - skel_tokens_est as f64 / src_tokens_est as f64) * 100.0
    );
    println!();
    println!("=== SKELETON PREVIEW (first 30 lines) ===");
    for line in last_skel.lines().take(30) {
        println!("{}", line);
    }
}
