open Noun

(** Nock interpreter

    Based on the reference implementation from vere/pkg/noun/nock.c

    The Nock spec has 12 opcodes (0-11):
    - 0: slot/fragment lookup
    - 1: constant
    - 2: nock (recursion)
    - 3: is-cell test
    - 4: increment
    - 5: equality test
    - 6: if-then-else
    - 7: composition
    - 8: push
    - 9: call with axis
    - 10: hint (ignored in reference implementation)
    - 11: scry (errors in reference implementation)
*)

(* Trace tracking - matches C logging style
   - call_count: global counter for all nock calls
   - depth: current recursion depth
   - max_calls: how many calls to log (like C's limit)
   - max_mug_depth: only compute mugs at shallow depths (expensive!)
*)
let call_count = ref 0
let depth = ref 0
let max_calls = 100
let max_mug_depth = 3  (* Only log mugs for depth <= 3 to avoid performance hit *)

(* Helper to generate indentation based on depth *)
let indent () =
  match !depth with
  | 0 -> ""
  | 1 -> "  "
  | 2 -> "    "
  | 3 -> "      "
  | 4 -> "        "
  | _ -> "          "

(* Helper to get opcode name/description *)
let opcode_name op =
  match op with
  | 0 -> "0(slot)"
  | 1 -> "1(const)"
  | 2 -> "2(eval)"
  | 3 -> "3(cell?)"
  | 4 -> "4(inc)"
  | 5 -> "5(eq)"
  | 6 -> "6(if)"
  | 7 -> "7(compose)"
  | 8 -> "8(extend)"
  | 9 -> "9(invoke)"
  | 10 -> "10(edit)"
  | 11 -> "11(hint)"
  | n -> Printf.sprintf "%d(?)" n

(** Main nock evaluation function: nock(subject, formula)

    In Nock notation: *[subject formula]

    This is a direct port of _n_nock_on from nock.c:157-396
*)
let rec nock_on bus fol =
  (* Capture current call number in local variable (like C's my_call)
     This ensures ENTER and EXIT logs show the same call number even after recursion *)
  let my_call = !call_count in
  let should_log = my_call < max_calls in

  (* Log entry - shows opcode, depth, subject type, and mug at shallow depths *)
  if should_log then begin
    let opcode_str = match fol with
      | Cell (Atom op, _) when Z.fits_int op -> opcode_name (Z.to_int op)
      | Cell (Cell _, _) -> "CELL(dist)"
      | _ -> "?"
    in
    (* Only compute mugs at shallow depths to avoid performance penalty *)
    if !depth <= max_mug_depth then begin
      let bus_mug = mug bus in
      Printf.eprintf "%s>>> ENTER call #%d depth=%d opcode=%s bus=%s[mug=0x%lx]\n%!"
        (indent ()) my_call !depth opcode_str
        (if is_cell bus then "cell" else "atom")
        bus_mug
    end else begin
      Printf.eprintf "%s>>> ENTER call #%d depth=%d opcode=%s bus=%s\n%!"
        (indent ()) my_call !depth opcode_str
        (if is_cell bus then "cell" else "atom")
    end;
    incr call_count
  end;

  (* Increment depth for recursive calls *)
  incr depth;

  try
    let result = match fol with
    | Cell (hib, gal) when is_cell hib ->
        (* Distribution: [a b] -> compute both sides and cons *)
        let poz = nock_on bus hib in
        let riv = nock_on bus gal in
        cell poz riv

    | Cell (Atom op, gal) ->
        (* Check if opcode fits in int *)
        if Z.compare op (Z.of_int max_int) > 0 then raise Exit;
        let opcode = Z.to_int op in

        (match opcode with
        | 0 ->
            (* Nock 0: /[axis subject] - slot/fragment lookup *)
            if not (is_atom gal) then raise Exit
            else slot (match gal with Atom n -> n | _ -> raise Exit) bus

        | 1 ->
            (* Nock 1: constant - return gal as-is *)
            gal

        | 2 ->
            (* Nock 2: *[subject formula new_subject] - evaluate with new subject *)
            (* C evaluates tail first, then head - must match this order! *)
            if not (is_cell gal) then raise Exit;
            let c_gal = tail gal in
            let b_gal = head gal in
            let nex = nock_on bus c_gal in  (* Tail first like C *)
            let seb = nock_on bus b_gal in  (* Head second like C *)
            nock_on seb nex

        | 3 ->
            (* Nock 3: ?[subject formula] - is-cell test *)
            let gof = nock_on bus gal in
            if is_cell gof then atom 0 else atom 1

        | 4 ->
            (* Nock 4: +[subject formula] - increment *)
            let gof = nock_on bus gal in
            inc gof

        | 5 ->
            (* Nock 5: =[subject formula] - equality test *)
            let wim = nock_on bus gal in
            if not (is_cell wim) then raise Exit;
            let a = head wim in
            let b = tail wim in
            if equal a b then atom 0 else atom 1

        | 6 ->
            (* Nock 6: if-then-else *)
            if not (is_cell gal) then raise Exit;
            let b_gal = head gal in
            let cd_gal = tail gal in
            if not (is_cell cd_gal) then raise Exit;
            let c_gal = head cd_gal in
            let d_gal = tail cd_gal in
            let tys = nock_on bus b_gal in
            let nex = match tys with
              | Atom n when Z.equal n Z.zero -> c_gal
              | Atom n when Z.equal n Z.one -> d_gal
              | _ -> raise Exit
            in
            nock_on bus nex

        | 7 ->
            (* Nock 7: composition - *[*[subject b] c] *)
            if not (is_cell gal) then raise Exit;
            let b_gal = head gal in
            let c_gal = tail gal in
            let bod = nock_on bus b_gal in
            nock_on bod c_gal

        | 8 ->
            (* Nock 8: extend - *[[*[subject b] subject] c] *)
            if not (is_cell gal) then raise Exit;
            let b_gal = head gal in
            let c_gal = tail gal in
            let heb = nock_on bus b_gal in
            let bod = cell heb bus in
            nock_on bod c_gal

        | 9 ->
            (* Nock 9: invoke - *[*[subject c] 2 [0 1] 0 axis] *)
            if not (is_cell gal) then raise Exit;
            let b_gal = head gal in
            let c_gal = tail gal in
            if not (is_atom b_gal) then raise Exit;
            let seb = nock_on bus c_gal in
            let nex = slot (match b_gal with Atom n -> n | _ -> raise Exit) seb in
            nock_on seb nex

        | 10 ->
            (* Nock 10: edit/hint - replace at slot
               Two forms:
               - *[a 10 [b c] d]: edit mode, replace slot b in result of d with result of c
               - *[a 10 b c]: hint mode, just evaluate c (ignore hint b)
            *)
            if not (is_cell gal) then raise Exit;
            let _p_gal = head gal in  (* Unused: hint tag/edit slot *)
            let q_gal = tail gal in
            (* For now, treat both forms as hints - just evaluate the formula *)
            nock_on bus q_gal

        | 11 ->
            (* Nock 11: hint (two forms)
               - *[a 11 [b c] d] → *[[*[a c] *[a d]] 0 3]  (dynamic hint)
               - *[a 11 b c] → *[a c]                       (static hint)

               The key insight: b is HINT DATA, not a formula!
               Don't try to evaluate it as Nock code.
            *)
            if not (is_cell gal) then raise Exit;
            let p_gal = head gal in  (* b or [b c] - hint tag/data *)
            let q_gal = tail gal in  (* c or d - formula to evaluate *)

            if is_cell p_gal then begin
              (* Dynamic hint: *[a 11 [b c] d]
                 Spec: *[[*[a c] *[a d]] 0 3]
                 This evaluates both c and d, conses them, then returns slot 3 (= d's result).
                 Since we just want d's result, we can skip the hint evaluation. *)
              nock_on bus q_gal
            end else begin
              (* Static hint: *[a 11 b c]
                 Spec: *[a c]
                 Just evaluate c, ignore the hint atom b. *)
              nock_on bus q_gal
            end

        | _ ->
            (* Invalid opcode *)
            raise Exit
        )

    | _ ->
        (* Invalid formula structure *)
        raise Exit
    in

    (* Restore depth and log exit before returning *)
    decr depth;
    if should_log then begin
      (* Only compute mugs at shallow depths to avoid performance penalty *)
      if !depth <= max_mug_depth then begin
        let result_mug = mug result in
        Printf.eprintf "%s<<< EXIT  call #%d depth=%d returns=%s[mug=0x%lx]\n%!"
          (indent ()) my_call !depth
          (if is_cell result then "cell" else "atom")
          result_mug
      end else begin
        Printf.eprintf "%s<<< EXIT  call #%d depth=%d returns=%s\n%!"
          (indent ()) my_call !depth
          (if is_cell result then "cell" else "atom")
      end
    end;
    result

  with e ->
    (* Restore depth even on exception *)
    decr depth;
    raise e

(** Convenience function: nock(subject, formula) *)
let nock subject formula =
  nock_on subject formula

(** slam: apply gate to sample
    slam(gate, sample) = *[gate [9 2 [0 1] [0 6] [1 sample] [0 7]]]

    In practice this evaluates the gate (which is a core with a formula at axis 2)
    with a modified sample (at axis 6).
*)
let slam gat sam =
  let cor = cell (head gat) (cell sam (tail (tail gat))) in
  let formula = slot (Z.of_int 2) cor in
  nock_on cor formula

(** kick: fire gate without changing sample
    kick(gate) = *[gate 9 2 0 1]
*)
let kick gat =
  let formula = slot (Z.of_int 2) gat in
  nock_on gat formula