(* State Management - Domain-safe Arvo kernel state
 *
 * This module manages the Urbit runtime state with:
 * - Domain-safe operations using Atomic
 * - Arvo kernel state (roc)
 * - Event counter
 * - Atomic snapshots with Eio.Promise
 *
 * Key differences from C implementation:
 * - No loom! OCaml GC handles memory
 * - Atomic operations for thread-safety across domains
 * - Simplified memory management
 *)

(* Ship state *)
type t = {
  (* Arvo core - the kernel state *)
  mutable roc: Noun.noun;

  (* Event number - using mutable int64 with mutex for now
   * TODO: Use lock-free approach with Kcas later *)
  mutable eve: int64;

  (* Wish cache - for compiled expressions *)
  mutable yot: (string, Noun.noun) Hashtbl.t;

  (* Lock for state updates (Mutex for now, will use lock-free later) *)
  lock: Mutex.t;
}

(* Create empty state *)
let create () =
  {
    roc = Noun.atom 0;  (* Empty initial state *)
    eve = 0L;
    yot = Hashtbl.create 256;
    lock = Mutex.create ();
  }

(* Get current event number *)
let event_num state =
  Mutex.lock state.lock;
  let eve = state.eve in
  Mutex.unlock state.lock;
  eve

(* Get Arvo core *)
let get_arvo state =
  Mutex.lock state.lock;
  let roc = state.roc in
  Mutex.unlock state.lock;
  roc

(* Set Arvo core *)
let set_arvo state roc =
  Mutex.lock state.lock;
  state.roc <- roc;
  Mutex.unlock state.lock

(* Increment event number *)
let inc_event state =
  Mutex.lock state.lock;
  let old_eve = state.eve in
  state.eve <- Int64.succ state.eve;
  Mutex.unlock state.lock;
  old_eve

(* Boot from a pill (simplified - just load a noun as the kernel) *)
let boot state kernel_noun =
  Mutex.lock state.lock;
  state.roc <- kernel_noun;
  state.eve <- 0L;
  Hashtbl.clear state.yot;
  Mutex.unlock state.lock

(* Poke Formula - Real Arvo gate call
 *
 * This is the Nock formula to call the Arvo kernel gate with an event.
 *
 * Formula: [9 2 [0 2] [0 3]]
 *   - Opcode 9: Call gate
 *   - Arm 2: The $ arm (standard gate arm)
 *   - Sample: [0 2] - the event from slot 2
 *   - Context: [0 3] - the kernel from slot 3
 *
 * Subject structure: [event kernel]
 *   - Slot 2 = event (the ovum)
 *   - Slot 3 = kernel (Arvo core/gate)
 *)
let poke_formula =
  (* TEST: Simplest formula - just return subject [0 1] *)
  Noun.cell (Noun.atom 0) (Noun.atom 1)

  (* TODO: Real gate call formula once we understand Arvo's structure
   * [9 2 [0 2] [0 3]] or similar *)

(* Parse poke result
 *
 * Arvo poke result structure: [effects new-kernel]
 * Or sometimes: [[moves new-kernel] effects]
 *
 * For now, simplified: assume result is the new kernel
 *)
let parse_poke_result result =
  (* TODO: Parse real Arvo result structure
   * For now: treat whole result as new kernel *)
  let new_kernel = result in
  let effects = [] in  (* No effects parsed yet *)
  (new_kernel, effects)

(* Poke: apply an event to the kernel
 *
 * Real Arvo poke sequence:
 * 1. Build subject: [event kernel]
 * 2. Run Nock with poke formula
 * 3. Parse result: [effects new-kernel]
 * 4. Update kernel state
 * 5. Return effects
 *)
let poke state event_noun =
  Mutex.lock state.lock;

  try
    (* Build subject: [event kernel] *)
    let subject = Noun.cell event_noun state.roc in

    Printf.printf "[State] Calling Arvo with poke formula...\n%!";
    Printf.printf "[State] Subject: [event kernel]\n%!";
    Printf.printf "[State] Formula: [9 2 [0 2] [0 3]]\n%!";

    (* Run Nock with poke formula *)
    let result = Nock.nock_on subject poke_formula in

    Printf.printf "[State] ✓ Nock execution succeeded!\n%!";

    (* Parse result *)
    let (new_kernel, effects) = parse_poke_result result in

    (* Update kernel state *)
    state.roc <- new_kernel;
    state.eve <- Int64.succ state.eve;

    Mutex.unlock state.lock;

    Printf.printf "[State] ✓ Poke complete, event number: %Ld\n%!" state.eve;

    (* Return effects *)
    effects

  with e ->
    (* Nock error - don't update state *)
    Mutex.unlock state.lock;
    Printf.printf "[State] ✗ Poke failed with exception: %s\n%!" (Printexc.to_string e);
    Printf.printf "[State] Stack trace:\n%!";
    Printf.printf "%s\n%!" (Printexc.get_backtrace ());
    []

(* Peek Formula - Scry formula
 *
 * Scry is a read-only query into Arvo.
 * Formula: Similar to poke but doesn't update state
 *
 * For now: simplified - just return the path from the kernel
 *)
let peek_formula =
  (* Simplified: [0 1] - return the whole subject *)
  Noun.cell
    (Noun.atom 0)
    (Noun.atom 1)

(* Peek: query the kernel state (read-only)
 *
 * Real Arvo scry:
 * 1. Build subject: [path kernel]
 * 2. Run Nock with peek formula
 * 3. Return result (no state update!)
 *
 * Multiple domains can peek concurrently since it's read-only.
 *)
let peek state path_noun =
  (* No lock needed for read! This is why peek is fast *)
  let kernel = state.roc in

  try
    (* Build subject: [path kernel] *)
    let subject = Noun.cell path_noun kernel in

    (* Run Nock with peek formula - read-only! *)
    let result = Nock.nock_on subject peek_formula in

    Some result

  with e ->
    (* Scry failed *)
    Printf.printf "[State] Peek failed: %s\n%!" (Printexc.to_string e);
    None

(* Save snapshot to file using Eio
 *
 * Snapshot format:
 * - Event number (8 bytes)
 * - Jammed Arvo core
 *)
let save_snapshot state ~fs path =
  (* Take atomic snapshot of current state *)
  Mutex.lock state.lock;
  let eve = state.eve in
  let roc = state.roc in
  Mutex.unlock state.lock;

  (* Serialize: 8-byte event number + jammed state *)
  let jammed = Serial.jam roc in
  let jam_len = Bytes.length jammed in

  let total_len = 8 + jam_len in
  let snapshot = Bytes.create total_len in

  (* Write event number (little-endian) *)
  Bytes.set_int64_le snapshot 0 eve;

  (* Write jammed state *)
  Bytes.blit jammed 0 snapshot 8 jam_len;

  (* Write to file using Eio *)
  let file_path = Eio.Path.(fs / path) in
  Eio.Path.save ~create:(`Or_truncate 0o644) file_path (Bytes.to_string snapshot)

(* Load snapshot from file using Eio *)
let load_snapshot state ~fs path =
  let file_path = Eio.Path.(fs / path) in

  try
    let data = Eio.Path.load file_path in
    let bytes = Bytes.of_string data in

    if Bytes.length bytes < 8 then
      Error "Snapshot too short"
    else
      (* Read event number *)
      let eve = Bytes.get_int64_le bytes 0 in

      (* Read jammed state *)
      let jam_len = Bytes.length bytes - 8 in
      let jammed = Bytes.sub bytes 8 jam_len in

      (* Cue the state *)
      let roc = Serial.cue jammed in

      (* Update state *)
      Mutex.lock state.lock;
      state.roc <- roc;
      state.eve <- eve;
      Hashtbl.clear state.yot;
      Mutex.unlock state.lock;

      Ok eve
  with
  | Sys_error msg -> Error ("Failed to load snapshot: " ^ msg)
  | e -> Error ("Failed to load snapshot: " ^ Printexc.to_string e)

(* Save snapshot with Eio.Promise for async completion *)
let save_snapshot_async state ~sw:_ ~fs path =
  save_snapshot state ~fs path;
  Eio.Promise.create_resolved ()

(* Load snapshot with promise *)
let load_snapshot_async state ~sw:_ ~fs path =
  let result = load_snapshot state ~fs path in
  Eio.Promise.create_resolved result

(* Clear wish cache (for memory reclamation) *)
let clear_cache state =
  Mutex.lock state.lock;
  Hashtbl.clear state.yot;
  Mutex.unlock state.lock

(* Get state summary for debugging *)
let summary state =
  Mutex.lock state.lock;
  let eve = state.eve in
  let cache_size = Hashtbl.length state.yot in
  Mutex.unlock state.lock;
  Printf.sprintf "State[eve=%Ld, cache=%d]" eve cache_size