open Nock_lib.Noun
open Nock_lib.Serial

(** Benchmark utilities *)

let time_it f =
  let start = Unix.gettimeofday () in
  let result = f () in
  let elapsed = Unix.gettimeofday () -. start in
  (result, elapsed)

let benchmark name iterations f =
  (* Warmup *)
  for _i = 1 to min 100 (iterations / 10) do
    let _ = f () in
    ()
  done;

  (* Actual benchmark *)
  let times = ref [] in
  for _i = 1 to iterations do
    let (_, elapsed) = time_it f in
    times := elapsed :: !times
  done;

  let total = List.fold_left (+.) 0.0 !times in
  let avg = total /. float_of_int iterations in
  let sorted = List.sort compare !times in
  let median = List.nth sorted (iterations / 2) in

  Printf.printf "%-40s %d iters: avg=%.6f median=%.6f total=%.6f\n"
    name iterations avg median total

(** Benchmark cases *)

let bench_atom_small () =
  benchmark "jam/cue small atom (42)" 100000 (fun () ->
    let n = atom 42 in
    let j = jam n in
    let c = cue j in
    c
  )

let bench_atom_large () =
  benchmark "jam/cue large atom (2^64)" 10000 (fun () ->
    let n = Atom (Z.shift_left Z.one 64) in
    let j = jam n in
    let c = cue j in
    c
  )

let bench_cell_simple () =
  benchmark "jam/cue simple cell [1 2]" 100000 (fun () ->
    let n = cell (atom 1) (atom 2) in
    let j = jam n in
    let c = cue j in
    c
  )

let bench_tree_balanced () =
  let tree = 
    cell 
      (cell (cell (atom 1) (atom 2)) (cell (atom 3) (atom 4)))
      (cell (cell (atom 5) (atom 6)) (cell (atom 7) (atom 8)))
  in
  benchmark "jam/cue balanced tree (depth 3)" 50000 (fun () ->
    let j = jam tree in
    let c = cue j in
    c
  )

let bench_list_structure () =
  let rec make_list n =
    if n = 0 then atom 0
    else cell (atom n) (make_list (n - 1))
  in
  let list = make_list 20 in
  benchmark "jam/cue list structure (20 elements)" 10000 (fun () ->
    let j = jam list in
    let c = cue j in
    c
  )

let bench_deep_nesting () =
  let rec make_deep n =
    if n = 0 then atom 0
    else cell (atom n) (make_deep (n - 1))
  in
  let deep = make_deep 100 in
  benchmark "jam/cue deep nesting (100 levels)" 1000 (fun () ->
    let j = jam deep in
    let c = cue j in
    c
  )

let bench_jam_only_small () =
  let n = atom 42 in
  benchmark "jam only (small atom)" 100000 (fun () ->
    let j = jam n in
    j
  )

let bench_cue_only_small () =
  let n = atom 42 in
  let j = jam n in
  (* Copy the bytes to avoid any mutation issues *)
  let j_copy = Bytes.copy j in
  benchmark "cue only (small atom)" 100000 (fun () ->
    let c = cue j_copy in
    c
  )

let bench_jam_only_tree () =
  let tree =
    cell
      (cell (cell (atom 1) (atom 2)) (cell (atom 3) (atom 4)))
      (cell (cell (atom 5) (atom 6)) (cell (atom 7) (atom 8)))
  in
  benchmark "jam only (balanced tree)" 50000 (fun () ->
    let j = jam tree in
    j
  )

let bench_cue_only_tree () =
  let tree =
    cell
      (cell (cell (atom 1) (atom 2)) (cell (atom 3) (atom 4)))
      (cell (cell (atom 5) (atom 6)) (cell (atom 7) (atom 8)))
  in
  let j = jam tree in
  (* Copy the bytes to avoid any mutation issues *)
  let j_copy = Bytes.copy j in
  benchmark "cue only (balanced tree)" 50000 (fun () ->
    let c = cue j_copy in
    c
  )

(** Run all benchmarks *)
let () =
  Printf.printf "========================================\n";
  Printf.printf "Jam/Cue Serialization Benchmarks\n";
  Printf.printf "========================================\n\n";
  
  Printf.printf "Round-trip benchmarks:\n";
  bench_atom_small ();
  bench_atom_large ();
  bench_cell_simple ();
  bench_tree_balanced ();
  bench_list_structure ();
  bench_deep_nesting ();
  
  Printf.printf "\nJam-only benchmarks:\n";
  bench_jam_only_small ();
  bench_jam_only_tree ();
  
  Printf.printf "\nCue-only benchmarks:\n";
  bench_cue_only_small ();
  bench_cue_only_tree ();
  
  Printf.printf "\n========================================\n"