(** Noun type and basic operations *)

(** Atom record with inline mug cache *)
type atom_record = {
  z: Z.t;              (** The integer value *)
  mutable mug: int32;  (** Cached 31-bit hash (0l = not computed yet) *)
}

(** Cell record with inline mug cache *)
type cell_record = {
  h: noun;             (** Head *)
  t: noun;             (** Tail *)
  mutable mug: int32;  (** Cached 31-bit hash (0l = not computed yet) *)
}

(** A noun is either an atom (arbitrary-precision integer) or a cell (pair of nouns) *)
and noun =
  | Atom of atom_record  (** Arbitrary-precision integer with mug cache *)
  | Cell of cell_record  (** Pair of nouns with mug cache *)

(** Exception raised on nock evaluation errors *)
exception Exit

(** Create an atom from an int *)
let atom n = Atom { z = Z.of_int n; mug = 0l }

(** Create a cell *)
let cell a b = Cell { h = a; t = b; mug = 0l }

(** Test if a noun is a cell *)
let is_cell = function
  | Cell _ -> true
  | Atom _ -> false

(** Test if a noun is an atom *)
let is_atom = function
  | Atom _ -> true
  | Cell _ -> false

(** Get head of a cell *)
let head = function
  | Cell { h; _ } -> h
  | Atom _ -> raise Exit

(** Get tail of a cell *)
let tail = function
  | Cell { t; _ } -> t
  | Atom _ -> raise Exit

(** Fragment/axis lookup: slot(n, noun)
    This implements the tree-addressing scheme:
    - 1 is the root
    - 2 is head, 3 is tail
    - For n > 1: if even, go left; if odd, go right
*)
let slot_debug = ref false

let rec slot n noun =
  let debug = !slot_debug && (Z.equal n (Z.of_int 2) || Z.equal n (Z.of_int 3)) in
  if debug then
    Printf.eprintf "[SLOT-DEBUG] slot(%s, noun)\n%!" (Z.to_string n);

  if Z.equal n Z.one then begin
    if debug then Printf.eprintf "[SLOT-DEBUG] slot(%s) = identity\n%!" (Z.to_string n);
    noun
  end
  else if Z.equal n Z.zero then
    raise Exit
  else begin
    let bit = Z.testbit n 0 in  (* Check if odd *)
    let parent = Z.shift_right n 1 in
    if debug then
      Printf.eprintf "[SLOT-DEBUG] slot(%s): bit=%b parent=%s\n%!"
        (Z.to_string n) bit (Z.to_string parent);
    let sub = slot parent noun in
    let result = if bit then tail sub else head sub in
    if debug then begin
      let result_str =
        if is_cell result then "cell"
        else "atom"
      in
      Printf.eprintf "[SLOT-DEBUG] slot(%s): taking %s, result is %s\n%!"
        (Z.to_string n) (if bit then "tail" else "head") result_str
    end;
    result
  end

(** Equality test for nouns *)
let rec equal a b =
  match a, b with
  | Atom { z = x; _ }, Atom { z = y; _ } -> Z.equal x y
  | Cell { h = ah; t = at; _ }, Cell { h = bh; t = bt; _ } ->
      equal ah bh && equal at bt
  | _, _ -> false

(** Increment an atom *)
let inc = function
  | Atom { z; _ } -> Atom { z = Z.succ z; mug = 0l }
  | Cell _ -> raise Exit

(** Compute mug (31-bit hash) of a noun with inline caching

    This implements Urbit's mug hash function using FNV-1a.
    Mugs are stored inline in the noun structure (like C's u3 system).

    The mug field is 0l when not yet computed, and computed lazily on first access.
    This matches C's approach where u3r_mug() is just a memory read after first computation.

    For atoms: hash the bytes of the integer representation
    For cells: mix the mugs of head and tail
*)
let rec mug noun =
  match noun with
  | Atom { z; mug = cached_mug } ->
      if cached_mug <> 0l then cached_mug
      else begin
        let computed = compute_mug_atom z in
        (* Update inline cache *)
        (match noun with
         | Atom r -> r.mug <- computed
         | _ -> ());
        computed
      end
  | Cell { h; t; mug = cached_mug } ->
      if cached_mug <> 0l then cached_mug
      else begin
        let h_mug = mug h in
        let t_mug = mug t in
        let computed = mix_mugs h_mug t_mug in
        (* Update inline cache *)
        (match noun with
         | Cell r -> r.mug <- computed
         | _ -> ());
        computed
      end

and compute_mug_atom z =
  (* FNV-1a constants - using hex to avoid signed int32 overflow *)
  let fnv_prime = 16777619l in
  let fnv_basis = 0x811c9dc5l in  (* 2166136261 in decimal *)

  (* Mask to 31 bits (Urbit uses 31-bit mugs) *)
  let mask31 x = Int32.logand x 0x7fffffffl in

  (* Convert atom to bytes and hash using FNV-1a *)
  let bytes = Z.to_bits z in
  let len = String.length bytes in
  let rec loop i hash =
    if i >= len then hash
    else
      let byte = Int32.of_int (Char.code bytes.[i]) in
      let hash' = Int32.mul (Int32.logxor hash byte) fnv_prime in
      loop (i + 1) hash'
  in
  mask31 (loop 0 fnv_basis)

and mix_mugs a_mug b_mug =
  (* Mix two mugs together (for cells) *)
  let fnv_prime = 16777619l in
  let mask31 x = Int32.logand x 0x7fffffffl in
  let mixed = Int32.mul (Int32.logxor a_mug b_mug) fnv_prime in
  mask31 mixed

(** Pretty-print a noun *)
let rec pp_noun fmt = function
  | Atom { z; _ } -> Format.fprintf fmt "%s" (Z.to_string z)
  | Cell { h; t; _ } -> Format.fprintf fmt "[%a %a]" pp_noun h pp_noun t