[@@@ocaml.warning "-32-39"]

open Nock_lib
open Noun
open Serial

let atom_int n = atom (Z.of_int n)

let rec mass = function
  | Atom z -> 1 + Z.numbits z
  | Cell (h, t) -> 1 + mass h + mass t

let noun_to_string noun =
  let rec loop = function
    | Atom z -> Z.to_string z
    | Cell (h, t) -> Printf.sprintf "[%s %s]" (loop h) (loop t)
  in
  loop noun

let bytes_to_hex bytes =
  let buf = Buffer.create (Bytes.length bytes * 2) in
  Bytes.iter
    (fun byte -> Buffer.add_string buf (Printf.sprintf "%02X" (int_of_char byte)))
    bytes;
  Buffer.contents buf

let log_roundtrip name noun jammed =
  Printf.printf "[%s] noun=%s jam=%s (%d bytes)\n%!"
    name (noun_to_string noun) (bytes_to_hex jammed) (Bytes.length jammed)

let roundtrip name noun =
  let jammed = jam noun in
  log_roundtrip name noun jammed;
  let cued = cue jammed in
  Alcotest.(check bool) (name ^ " roundtrip") true (equal noun cued);
  jammed

let setup_stack () = ()

let test_jam_cue_atom () =
  ignore (roundtrip "test_jam_cue_atom" (atom_int 42))

let test_jam_cue_cell () =
  ignore (roundtrip "test_jam_cue_cell" (cell (atom_int 1) (atom_int 2)))

let test_jam_cue_nested_cell () =
  let inner = cell (atom_int 3) (atom_int 4) in
  ignore (roundtrip "test_jam_cue_nested_cell" (cell (atom_int 1) inner))

let test_jam_cue_shared_structure () =
  let shared = atom_int 99 in
  ignore (roundtrip "test_jam_cue_shared_structure" (cell shared shared))

let test_jam_cue_large_atom () =
  let noun = atom (Z.of_string "18446744073709551615") in
  ignore (roundtrip "test_jam_cue_large_atom" noun)

let test_jam_cue_empty_atom () =
  ignore (roundtrip "test_jam_cue_empty_atom" (atom Z.zero))

let test_jam_cue_complex_structure () =
  let atom1 = atom_int 1 in
  let atom2 = atom_int 2 in
  let cell1 = cell atom1 atom2 in
  let cell2 = cell cell1 atom2 in
  let cell3 = cell cell2 cell1 in
  ignore (roundtrip "test_jam_cue_complex_structure" cell3)

let test_cue_invalid_input () =
  let invalid = Bytes.of_string "\x03" in
  Alcotest.check_raises "test_cue_invalid_input"
    (Invalid_argument "read_bit: end of stream")
    (fun () -> ignore (cue invalid))

let rec space_needed_noun noun = mass noun

let rec random_noun rng depth acc =
  if depth <= 0 || Random.State.bool rng then
    let value = Random.State.int64 rng Int64.max_int in
    let noun = atom (Z.of_int64 value) in
    (noun, acc + mass noun)
  else
    let left, left_size = random_noun rng (depth - 1) acc in
    let right, total = random_noun rng (depth - 1) left_size in
    let noun = cell left right in
    (noun, total + mass noun)

let generate_random_noun _stack bits rng =
  let depth = max 1 (bits / 16) in
  random_noun rng depth 0

let rec generate_deeply_nested_noun stack depth rng =
  if depth = 0 then
    generate_random_noun stack 64 rng
  else
    let left, left_size = generate_deeply_nested_noun stack (depth - 1) rng in
    let right, right_size = generate_deeply_nested_noun stack (depth - 1) rng in
    let noun = cell left right in
    (noun, left_size + right_size + mass noun)

let test_jam_cue_roundtrip_property () =
  let rng = Random.State.make [| 0x1A2B3C4D |] in
  for depth = 1 to 6 do
    let noun, _ = generate_deeply_nested_noun () depth rng in
    ignore (roundtrip (Printf.sprintf "test_jam_cue_roundtrip_property depth=%d" depth) noun)
  done

let test_cue_invalid_backreference () =
  let invalid = Bytes.of_string "\x03" in
  Alcotest.check_raises "test_cue_invalid_backreference"
    (Invalid_argument "read_bit: end of stream")
    (fun () -> ignore (cue invalid))

let test_cue_nondeterministic_error () =
  Alcotest.(check bool) "test_cue_nondeterministic_error" true true

let test_cell_construction () =
  match cell (atom_int 10) (atom_int 11) with
  | Cell (h, t) ->
      Alcotest.(check bool) "head" true (equal h (atom_int 10));
      Alcotest.(check bool) "tail" true (equal t (atom_int 11))
  | Atom _ -> Alcotest.fail "expected cell"

let tests = [
  "test_jam_cue_atom", `Quick, test_jam_cue_atom;
  "test_jam_cue_cell", `Quick, test_jam_cue_cell;
  "test_jam_cue_nested_cell", `Quick, test_jam_cue_nested_cell;
  "test_jam_cue_shared_structure", `Quick, test_jam_cue_shared_structure;
  "test_jam_cue_large_atom", `Quick, test_jam_cue_large_atom;
  "test_jam_cue_empty_atom", `Quick, test_jam_cue_empty_atom;
  "test_jam_cue_complex_structure", `Quick, test_jam_cue_complex_structure;
  "test_cue_invalid_input", `Quick, test_cue_invalid_input;
  "test_jam_cue_roundtrip_property", `Quick, test_jam_cue_roundtrip_property;
  "test_cue_invalid_backreference", `Quick, test_cue_invalid_backreference;
  "test_cue_nondeterministic_error", `Quick, test_cue_nondeterministic_error;
  "test_cell_construction", `Quick, test_cell_construction;
]

let () = Alcotest.run "serialization" [ "tests", tests ]