let root () = try Sys.getenv "DELTA_ROOT" with Not_found -> "."

let fail name message =
  prerr_endline (Printf.sprintf "bench: %s: %s" name message);
  exit 1

let deltac () = Filename.concat (root ()) "deltac"

let build_dir () = try Sys.getenv "DELTA_BUILD_DIR" with Not_found -> "_build"

let run command =
  let log = Filename.temp_file "delta_bench" ".log" in
  let fd = Unix.openfile log [ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC ] 0o600 in
  let argv = Array.of_list command in
  let pid = Unix.create_process argv.(0) argv Unix.stdin fd fd in
  let status = snd (Unix.waitpid [] pid) in
  Unix.close fd;
  let text = Native.read_file log in
  (try Sys.remove log with _ -> ());
  (status, text)

let write_file path text =
  let channel = open_out path in
  output_string channel text;
  close_out channel

type row = { customer : string; total : int }

let make_row index = { customer = Printf.sprintf "customer%d" index; total = (index * 37) mod 4001 }

let row_text row = Printf.sprintf "(record (customer %S) (total %d))" row.customer row.total

let round_of size seed = ((seed * 1103515245) + size) mod size

let input_text size =
  let buffer = Buffer.create (size * 48) in
  for index = 0 to size - 1 do
    Buffer.add_string buffer (Printf.sprintf "(%d %s)\n" (index + 1) (row_text (make_row index)))
  done;
  Buffer.contents buffer

let updates_text size batches =
  let buffer = Buffer.create (batches * 64) in
  for batch = 0 to batches - 1 do
    let key = (round_of size batch) + 1 in
    if batch mod 2 = 0 then
      Buffer.add_string buffer
        (Printf.sprintf "(batch (replace %d %s))\n" key
           (row_text { customer = Printf.sprintf "customer%d" key; total = 2000 + (batch mod 500) }))
    else
      Buffer.add_string buffer
        (Printf.sprintf "(batch (insert %d %s))\n" (size + batch + 1)
           (row_text { customer = Printf.sprintf "extra%d" batch; total = 1500 + (batch mod 250) }))
  done;
  Buffer.contents buffer

type stats = {
  st_init_seconds : float;
  st_update_seconds : float;
  st_minor_words : float;
  st_init_counters : string;
  st_update_counters : string;
}

let field prefix line =
  let length = String.length prefix in
  if String.length line > length && String.sub line 0 length = prefix then
    Some (String.trim (String.sub line length (String.length line - length)))
  else None

let counter_value name text =
  let parts = String.split_on_char ' ' text in
  let rec search = function
    | [] -> 0
    | part :: rest -> (
        match (try Some (String.index part '=') with Not_found -> None) with
        | Some index
          when String.sub part 0 index = name ->
            int_of_string (String.sub part (index + 1) (String.length part - index - 1))
        | _ -> search rest)
  in
  search parts

let parse_stats text =
  let lines = String.split_on_char '\n' text in
  let value prefix default =
    let rec search = function
      | [] -> default
      | line :: rest -> (
          match field prefix line with
          | Some text -> float_of_string text
          | None -> search rest)
    in
    search lines
  in
  let counter prefix default =
    let rec search = function
      | [] -> default
      | line :: rest -> (
          match field prefix line with Some text -> text | None -> search rest)
    in
    search lines
  in
  {
    st_init_seconds = value "init_seconds: " 0.0;
    st_update_seconds = value "update_seconds: " 0.0;
    st_minor_words = value "minor_words: " 0.0;
    st_init_counters = counter "init_counters: " "";
    st_update_counters = counter "update_counters: " "";
  }

let query_source name =
  let path = Filename.concat (Filename.concat (root ()) "example") (name ^ ".delta") in
  Native.read_file path

let compile_query name directory =
  let source = query_source name in
  let program = Filename.concat directory (name ^ ".ml") in
  write_file program source;
  let plan = Simplify.simplify (Graph.build (Anf.program (Specialize.program (Infer.program (Resolve.program (Parse.program source)))))) in
  let emitted = Filename.concat directory (name ^ "_generated.ml") in
  write_file emitted (Emit.program_to_string plan);
  let executable = Filename.concat directory name in
  let command =
    [ (try Sys.getenv "DELTA_OCAMLOPT" with Not_found -> "ocamlopt");
      "-w"; "-26"; "-I"; build_dir (); "-I"; "+unix"; "-o"; executable;
      Filename.concat (build_dir ()) "delta_runtime.cmxa"; "unix.cmxa"; emitted ]
  in
  match run command with
  | Unix.WEXITED 0, _ -> executable
  | _, output -> fail "compile" output

let sizes = [ 1000; 100000; 1000000 ]

let queries = [ "expensive_order"; "revenue"; "count_large" ]

let () =
  let directory = Filename.temp_file "delta_bench" "" in
  Sys.remove directory;
  Unix.mkdir directory 0o700;
  print_endline "delta benchmarks";
  print_endline "";
  print_endline
    "query             rows      init_ms  update_ms  batches  update_us_per_batch  minor_words  predicate_evals  mapping_evals  changed_keys  full_traversals";
  let regression_failures = ref 0 in
  List.iter
    (fun name ->
      let executable = compile_query name directory in
      List.iter
        (fun size ->
          let input_path = Filename.concat directory (Printf.sprintf "input_%s_%d.sexp" name size) in
          let updates_path = Filename.concat directory (Printf.sprintf "update_%s_%d.sexp" name size) in
          write_file input_path (input_text size);
          write_file updates_path (updates_text size 200);
          let status, output =
            run [ executable; "--input"; input_path; "--updates"; updates_path; "--print-result"; "--stats" ]
          in
          if status <> Unix.WEXITED 0 then fail "benchmark run" output;
          let stats = parse_stats output in
          let batches = 200 in
          Printf.printf "%-17s %-9d %-8.3f %-9.3f %-8d %-20.3f %-12.0f %-16d %-14d %-13d %d\n" name size
            (stats.st_init_seconds *. 1000.0) (stats.st_update_seconds *. 1000.0) batches
            (stats.st_update_seconds *. 1000000.0 /. float_of_int batches)
            stats.st_minor_words
            (counter_value "predicate_evaluations" stats.st_update_counters)
            (counter_value "mapping_evaluations" stats.st_update_counters)
            (counter_value "changed_key_visits" stats.st_update_counters)
            (counter_value "full_traversals" stats.st_update_counters);
          if counter_value "full_traversals" stats.st_update_counters <> 0 then (
            incr regression_failures;
            Printf.printf "  regression: updates performed %d full traversals for %s at %d rows\n"
              (counter_value "full_traversals" stats.st_update_counters) name size);
          if counter_value "predicate_evaluations" stats.st_update_counters > 200 * 4 then (
            incr regression_failures;
            Printf.printf "  regression: too many predicate evaluations (%d) for %s at %d rows\n"
              (counter_value "predicate_evaluations" stats.st_update_counters) name size);
          if counter_value "mapping_evaluations" stats.st_update_counters > 200 * 4 then (
            incr regression_failures;
            Printf.printf "  regression: too many mapping evaluations (%d) for %s at %d rows\n"
              (counter_value "mapping_evaluations" stats.st_update_counters) name size);
          Sys.remove input_path;
          Sys.remove updates_path)
        sizes)
    queries;
  print_endline "";
  if !regression_failures = 0 then
    print_endline
      "regression check: single-key updates performed no full traversals and touched only the changed keys"
  else Printf.printf "%d regression checks failed\n" !regression_failures;
  Native.remove_dir directory;
  exit (if !regression_failures = 0 then 0 else 1)