open Nock_lib

(* module Murmur3 = struct *)
  (* let rotl32 x r = *)
    (* Int32.logor (Int32.shift_left x r) (Int32.shift_right_logical x (32 - r)) *)

  (* let fmix32 h = *)
    (* let open Int32 in *)
    (* let h = logxor h (shift_right_logical h 16) in *)
    (* let h = mul h 0x85ebca6bl in *)
    (* let h = logxor h (shift_right_logical h 13) in *)
    (* let h = mul h 0xc2b2ae35l in *)
    (* logxor h (shift_right_logical h 16) *)

  (* let hash32 ?(seed = 0l) bytes ~length = *)
    (* let c1 = 0xcc9e2d51l in *)
    (* let c2 = 0x1b873593l in *)
    (* let nblocks = length lsr 2 in *)
    (* let h1 = ref seed in *)

    (* for block = 0 to nblocks - 1 do *)
      (* let i = block lsl 2 in *)
      (* let k1 = *)
        (* let open Int32 in *)
        (* let b0 = of_int (Bytes.get_uint8 bytes i) in *)
        (* let b1 = shift_left (of_int (Bytes.get_uint8 bytes (i + 1))) 8 in *)
        (* let b2 = shift_left (of_int (Bytes.get_uint8 bytes (i + 2))) 16 in *)
        (* let b3 = shift_left (of_int (Bytes.get_uint8 bytes (i + 3))) 24 in *)
        (* logor b0 (logor b1 (logor b2 b3)) *)
      (* in *)
      (* let k1 = *)
        (* let open Int32 in *)
        (* let k1 = mul k1 c1 in *)
        (* let k1 = rotl32 k1 15 in *)
        (* mul k1 c2 *)
      (* in *)
      (* let open Int32 in *)
      (* let h = !h1 in *)
      (* let h = logxor h k1 in *)
      (* let h = rotl32 h 13 in *)
      (* let h = add (mul h 5l) 0xe6546b64l in *)
      (* h1 := h *)
    (* done; *)

    (* let tail_index = nblocks lsl 2 in *)
    (* let tail_len = length land 3 in *)
    (* let k1 = *)
      (* let open Int32 in *)
      (* let k = ref 0l in *)
      (* if tail_len >= 3 then *)
        (* k := logor !k (shift_left (of_int (Bytes.get_uint8 bytes (tail_index + 2))) 16); *)
      (* if tail_len >= 2 then *)
        (* k := logor !k (shift_left (of_int (Bytes.get_uint8 bytes (tail_index + 1))) 8); *)
      (* if tail_len >= 1 then begin *)
        (* k := logor !k (of_int (Bytes.get_uint8 bytes tail_index)); *)
        (* let kx = mul !k c1 in *)
        (* let kx = rotl32 kx 15 in *)
        (* Some (mul kx c2) *)
      (* end else *)
        (* None *)
    (* in *)
    (* let h1 = *)
      (* match k1 with *)
      (* | None -> !h1 *)
      (* | Some k1 -> Int32.logxor !h1 k1 *)
    (* in *)
    (* let h1 = Int32.logxor h1 (Int32.of_int length) in *)
    (* fmix32 h1 *)
(* end *)

external murmur3_hash32_seed : string -> int32 -> int32 = "caml_murmur3_hash32_seed"

module Mug = struct
  open Noun

  module Tbl = Hashtbl.Make(struct
    type t = noun
    let equal = (==)
    let hash = Hashtbl.hash
  end)

  let memo = Tbl.create 1024

  let adjust hash =
    let mask = Int32.of_int 0x7fffffff in
    let ham =
      Int32.logxor
        (Int32.shift_right_logical hash 31)
        (Int32.logand hash mask)
    in
    if Int32.equal ham Int32.zero then 0x7fff else Int32.to_int ham

  let mug_bytes bytes ~length ~seed =
    let rec loop seed attempts =
      if attempts >= 8 then Int32.of_int 0x7fff
      else
        (* Convert bytes to string, taking only 'length' bytes *)
        let data = Bytes.sub_string bytes 0 length in
        let hash = murmur3_hash32_seed data seed in
        let ham = adjust hash in
        if ham = 0 then
          loop (Int32.add seed (Int32.of_int 1)) (attempts + 1)
        else
          Int32.of_int ham
    in
    loop seed 0

  let mug_both left right =
    let len =
      let bits =
        let rec loop count value =
          if Int32.equal value Int32.zero then count
          else loop (count + 1) (Int32.shift_right_logical value 1)
        in
        loop 0 right
      in
      4 + ((bits + 7) lsr 3)
    in
    let buf = Bytes.make 8 '\000' in
    let store value offset =
      let mask = Int32.of_int 0xff in
      Bytes.set buf offset (Char.chr (Int32.to_int (Int32.logand value mask)));
      Bytes.set buf (offset + 1)
        (Char.chr (Int32.to_int (Int32.logand (Int32.shift_right_logical value 8) mask)));
      Bytes.set buf (offset + 2)
        (Char.chr (Int32.to_int (Int32.logand (Int32.shift_right_logical value 16) mask)));
      Bytes.set buf (offset + 3)
        (Char.chr (Int32.to_int (Int32.logand (Int32.shift_right_logical value 24) mask)));
    in
    store left 0;
    store right 4;
    mug_bytes buf ~length:len ~seed:(Int32.of_int 0xdeadbeef)

  let trim_trailing_zeros str =
    let len = String.length str in
    let rec find idx =
      if idx < 0 then -1
      else if Char.equal str.[idx] '\000' then find (idx - 1)
      else idx
    in
    match find (len - 1) with
    | -1 -> Bytes.create 0
    | last -> Bytes.sub (Bytes.of_string str) 0 (last + 1)

  let rec mug noun =
    match Tbl.find_opt memo noun with
    | Some value -> value
    | None ->
        let value =
          match noun with
          | Atom z ->
              let bytes = trim_trailing_zeros (Z.to_bits z) in
              let len = Bytes.length bytes in
              if len = 0 then Int32.of_int 0x79ff04e8
              else mug_bytes bytes ~length:len ~seed:(Int32.of_int 0xcafebabe)
          | Cell (h, t) ->
              let left = mug h in
              let right = mug t in
              mug_both left right
        in
        Tbl.add memo noun value;
        value
end

let timed f =
  let start = Unix.gettimeofday () in
  let res = f () in
  let elapsed_ms = (Unix.gettimeofday () -. start) *. 1000.0 in
  res, elapsed_ms

let rec find_project_root dir =
  let pills_dir = Filename.concat dir "pills" in
  if Sys.file_exists pills_dir && Sys.is_directory pills_dir then dir
  else
    let parent = Filename.dirname dir in
    if String.equal parent dir then invalid_arg "unable to locate project root containing pills/"
    else find_project_root parent

let project_root =
  match Sys.getenv_opt "NEOVERE_ROOT" with
  | Some root -> root
  | None ->
      let exe_dir = Filename.dirname Sys.executable_name in
      find_project_root exe_dir

let read_file path =
  let ic = open_in_bin path in
  let len = in_channel_length ic in
  let data = really_input_string ic len in
  close_in ic;
  Bytes.of_string data

let count_list noun =
  let rec loop acc = function
    | Noun.Atom z when Z.equal z Z.zero -> acc
    | Noun.Cell (_, t) -> loop (acc + 1) t
    | _ -> raise Noun.Exit
  in
  loop 0 noun

let lifecycle_formula =
  let open Noun in
  let axis03 = cell (atom_of_int 0) (atom_of_int 3) in
  let axis02 = cell (atom_of_int 0) (atom_of_int 2) in
  cell (atom_of_int 2) (cell axis03 axis02)

let run_lifecycle events =
  let gate = Nock.nock_on events lifecycle_formula in
  Noun.slot (Z.of_int 7) gate

let hex32 x =
  Printf.sprintf "0x%08x" (Int32.to_int x)

let () =
  let pill =
    match Array.to_list Sys.argv with
    | _ :: path :: _ -> path
    | _ ->
        Printf.eprintf "usage: compare_ivory PATH/ivory.pill\n%!";
        exit 1
  in
  let pill_path =
    let raw =
      if Filename.is_relative pill then Filename.concat project_root pill else pill
    in
    Unix.realpath raw
  in
  Printf.printf "Loading ivory pill from %s...\n%!" pill_path;
  let pill_bytes = read_file pill_path in
  let pill, cue_ms = timed (fun () -> Serial.cue ~verbose:true pill_bytes) in
  Printf.printf "perf: ivory cue %.3f ms\n%!" cue_ms;
  let ivory_mug = Mug.mug pill in
  Printf.printf "ivory_pil mug: %s\n%!" (hex32 ivory_mug);
  let arvo_core =
    match pill with
    | Noun.Cell (_, tail) -> tail
    | _ -> failwith "ivory pill must be a cell"
  in
  Printf.printf "arvo_core mug: %s\n%!" (hex32 (Mug.mug arvo_core));
  Printf.printf "ivory event count=%d\n%!" (count_list arvo_core);
  (* Check pill head *)
  let pill_head = match pill with
    | Noun.Cell (h, _) -> h
    | _ -> failwith "not a cell"
  in
  Printf.printf "pill head (should be %%ivory) mug: %s\n%!" (hex32 (Mug.mug pill_head));
  let kernel, boot_ms = timed (fun () -> run_lifecycle arvo_core) in
  Printf.printf "perf: ivory boot %.3f ms\n%!" boot_ms;
  let kernel_mug = Mug.mug kernel in
  Printf.printf "lite: core %s\n%!" (hex32 kernel_mug);
  Printf.printf "lite: final state %s\n%!" (hex32 kernel_mug);
  let slot axis =
    try
      let noun = Noun.slot (Z.of_int axis) kernel in
      Some (hex32 (Mug.mug noun))
    with Noun.Exit -> None
  in
  let print_slot axis label =
    match slot axis with
    | Some value -> Printf.printf "%s mug: %s\n%!" label value
    | None -> Printf.printf "%s unavailable\n%!" label
  in
  print_slot 2 "kernel slot 2";
  print_slot 3 "kernel slot 3";
  print_slot 23 "kernel slot 23";
  let jammed, jam_ms = timed (fun () -> Serial.jam ~verbose:true kernel) in
  Printf.printf "jam kernel %.3f ms\n%!" jam_ms;
  Printf.printf "jam kernel bytes=%d\n%!" (Bytes.length jammed);
  let digest = Digest.string (Bytes.unsafe_to_string jammed) |> Digest.to_hex in
  Printf.printf "kernel jam digest=%s\n%!" digest