open Nock_lib
module State = Nock_lib.State
open Noun

let atom_int n = atom (Z.of_int n)

let make_test_kernel ~effects ~new_core =
  let gate_result = cell effects new_core in
  let gate_battery = cell (atom_int 1) gate_result in
  let gate_payload = cell (atom_int 0) (atom_int 0) in
  let gate = cell gate_battery gate_payload in
  let formula = cell (atom_int 0) (atom (Z.of_int 6)) in
  let battery =
    cell (atom_int 0)
      (cell (atom_int 0)
        (cell (atom_int 0) formula))
  in
  let payload = cell gate (atom_int 0) in
  cell battery payload

let test_boot_and_event_counter () =
  let state = State.create () in
  Alcotest.(check int64) "initial eve" 0L (State.event_number state);
  let kernel = atom_int 99 in
  State.boot state kernel;
  Alcotest.(check bool) "kernel after boot" true (equal kernel (State.arvo_core state));
  Alcotest.(check int64) "eve reset" 0L (State.event_number state)

let test_boot_with_explicit_eve () =
  let state = State.create () in
  let kernel = atom_int 7 in
  State.boot ~events_played:3L state kernel;
  Alcotest.(check bool) "kernel after boot" true (equal kernel (State.arvo_core state));
  Alcotest.(check int64) "eve preset" 3L (State.event_number state)

let test_poke_updates_core_and_counter () =
  let effects = atom_int 0 in
  let new_core = atom_int 123 in
  let kernel = make_test_kernel ~effects ~new_core in
  Alcotest.(check bool) "axis 23" true
    (equal (Noun.slot (Z.of_int 23) kernel) (cell (atom_int 0) (atom (Z.of_int 6))));
  let state = State.create ~initial:kernel () in
  let event = atom_int 42 in
  let returned = State.poke state event in
  Alcotest.(check bool) "effects" true (equal returned effects);
  Alcotest.(check bool) "new core" true (equal new_core (State.arvo_core state));
  Alcotest.(check int64) "eve increment" 1L (State.event_number state)

let test_peek () =
  let kernel = cell (atom_int 1) (atom_int 2) in
  let state = State.create ~initial:kernel () in
  let path = atom_int 0 in
  match State.peek state path with
  | Some result ->
      Alcotest.(check bool) "peek subject" true (equal result (cell path kernel))
  | None -> Alcotest.fail "peek should succeed"

let test_snapshot_roundtrip () =
  let kernel = cell (atom_int 5) (atom_int 6) in
  let state = State.create ~initial:kernel () in
  let jammed, eve = State.snapshot state in
  Alcotest.(check int64) "snapshot eve" 0L eve;
  let state' = State.create () in
  State.load_snapshot state' jammed eve;
  Alcotest.(check bool) "core restored" true (equal (State.arvo_core state') kernel)

let tests = [
  "boot", `Quick, test_boot_and_event_counter;
  "boot preset eve", `Quick, test_boot_with_explicit_eve;
  "poke", `Quick, test_poke_updates_core_and_counter;
  "peek", `Quick, test_peek;
  "snapshot", `Quick, test_snapshot_roundtrip;
]

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