(** Noun type and basic operations *)

(** A noun is either an atom (arbitrary-precision integer) or a cell (pair of nouns) *)
type noun =
  | Atom of Z.t  (** Arbitrary-precision integer using Zarith *)
  | Cell of noun * noun  (** Pair of nouns *)

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

(** Create an atom from an int *)
let atom n = Atom (Z.of_int n)

(** Create a cell *)
let cell a b = Cell (a, b)

(** 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
      Printf.eprintf "[SLOT-DEBUG] slot(%s): taking %s, result is %s\n%!"
        (Z.to_string n) (if bit then "tail" else "head")
        (if is_cell result then
          let h = head result in
          if is_atom h then
            (match h with Atom z -> "cell[" ^ Z.to_string z ^ " ...]" | _ -> "cell[? ...]")
          else "cell[cell ...]"
        else "atom")
    end;
    result
  end

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

(** Increment an atom *)
let inc = function
  | Atom n -> Atom (Z.succ n)
  | Cell _ -> raise Exit

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

    This implements Urbit's mug hash function using FNV-1a.
    The mug is cached in the C implementation but we compute it fresh each time.

    For atoms: hash the bytes of the integer representation
    For cells: mix the mugs of head and tail
*)
let rec mug noun =
  (* 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

  (* Hash bytes using FNV-1a *)
  let hash_bytes bytes =
    let len = Bytes.length bytes in
    let rec loop i hash =
      if i >= len then hash
      else
        let byte = Int32.of_int (Bytes.get_uint8 bytes i) in
        let hash' = Int32.mul (Int32.logxor hash byte) fnv_prime in
        loop (i + 1) hash'
    in
    mask31 (loop 0 fnv_basis)
  in

  (* Mix two mugs together (for cells) *)
  let mix_mugs a_mug b_mug =
    (* Mix by XOR and multiply, then mask *)
    let mixed = Int32.mul (Int32.logxor a_mug b_mug) fnv_prime in
    mask31 mixed
  in

  match noun with
  | Atom z ->
      (* Convert atom to bytes and hash *)
      let bytes = Z.to_bits z in
      hash_bytes (Bytes.of_string bytes)

  | Cell (h, t) ->
      (* Mix the mugs of head and tail *)
      let h_mug = mug h in
      let t_mug = mug t in
      mix_mugs h_mug t_mug

(** Pretty-print a noun *)
let rec pp_noun fmt = function
  | Atom n -> Format.fprintf fmt "%s" (Z.to_string n)
  | Cell (a, b) -> Format.fprintf fmt "[%a %a]" pp_noun a pp_noun b