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 depth counter - only log first 20 levels *)
let trace_depth = ref 0
let max_trace_depth = 20

(** 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 =
  let should_trace = !trace_depth < max_trace_depth in
  if should_trace then incr trace_depth;
  try
    let result = match fol with
    | Cell (hib, gal) when is_cell hib ->
        (* [a b] -> compute both sides and cons *)
        if should_trace then Printf.eprintf "[Nock:%d] Cell-cell formula\n%!" !trace_depth;
        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 (
          if should_trace then Printf.eprintf "[Nock:%d] Opcode too large: %s\n%!" !trace_depth (Z.to_string op);
          raise Exit
        );
        let opcode = Z.to_int op in
        if should_trace then Printf.eprintf "[Nock:%d] Opcode %d\n%!" !trace_depth opcode;
        (match opcode with
        | 0 ->
            (* /[axis subject] - slot/fragment lookup *)
            if not (is_atom gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op0: gal not atom\n%!" !trace_depth;
              raise Exit
            )
            else slot (match gal with Atom n -> n | _ -> raise Exit) bus

        | 1 ->
            (* =[constant subject] - return constant *)
            gal

        | 2 ->
            (* *[subject formula new_subject] - evaluate with new subject *)
            if not (is_cell gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op2: gal not cell\n%!" !trace_depth;
              raise Exit
            );
            let b_gal = head gal in
            let c_gal = tail gal in
            let seb = nock_on bus b_gal in
            let nex = nock_on bus c_gal in
            nock_on seb nex

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

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

        | 5 ->
            (* =[subject formula] - equality test *)
            let wim = nock_on bus gal in
            if not (is_cell wim) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op5: wim not cell\n%!" !trace_depth;
              raise Exit
            );
            let a = head wim in
            let b = tail wim in
            if equal a b then atom 0 else atom 1

        | 6 ->
            (* if-then-else *)
            if not (is_cell gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op6: gal not cell\n%!" !trace_depth;
              raise Exit
            );
            let b_gal = head gal in
            let cd_gal = tail gal in
            if not (is_cell cd_gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op6: cd_gal not cell\n%!" !trace_depth;
              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
              | _ ->
                  if should_trace then Printf.eprintf "[Nock:%d] Op6: tys not 0 or 1\n%!" !trace_depth;
                  raise Exit
            in
            nock_on bus nex

        | 7 ->
            (* composition: *[*[subject b] c] *)
            if not (is_cell gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op7: gal not cell\n%!" !trace_depth;
              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 ->
            (* push: *[[*[subject b] subject] c] *)
            if not (is_cell gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op8: gal not cell\n%!" !trace_depth;
              raise Exit
            );
            let b_gal = head gal in
            let c_gal = tail gal in
            if should_trace then Printf.eprintf "[Nock:%d] Op8: computing b_gal...\n%!" !trace_depth;
            let heb = nock_on bus b_gal in
            if should_trace then Printf.eprintf "[Nock:%d] Op8: creating new subject [heb bus]...\n%!" !trace_depth;
            let bod = cell heb bus in
            if should_trace then Printf.eprintf "[Nock:%d] Op8: computing c_gal on new subject...\n%!" !trace_depth;
            nock_on bod c_gal

        | 9 ->
            (* call: *[*[subject c] axis] *)
            if not (is_cell gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op9: gal not cell\n%!" !trace_depth;
              raise Exit
            );
            let b_gal = head gal in
            let c_gal = tail gal in
            if not (is_atom b_gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op9: b_gal not atom\n%!" !trace_depth;
              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 ->
            (* hint - in reference implementation, hints are mostly ignored *)
            if should_trace then Printf.eprintf "[Nock:%d] Op10: hint (gal is %s)\n%!" !trace_depth
              (if is_cell gal then "cell" else "atom");
            let nex =
              if is_cell gal then begin
                (* [[hint-tag hint-value] formula] *)
                let hint_part = head gal in
                let formula = tail gal in
                if should_trace then Printf.eprintf "[Nock:%d] Op10: hint_part is %s, formula is %s\n%!" !trace_depth
                  (if is_cell hint_part then "cell" else "atom")
                  (if is_cell formula then "cell" else "atom");
                formula
              end else begin
                (* [hint-tag formula] where hint-value is implicit *)
                if should_trace then Printf.eprintf "[Nock:%d] Op10: implicit hint\n%!" !trace_depth;
                gal
              end
            in
            nock_on bus nex

        | 11 ->
            (* scry - static scry *)
            if should_trace then Printf.eprintf "[Nock:%d] Op11: scry (gal is %s)\n%!" !trace_depth
              (if is_cell gal then "cell" else "atom");
            if not (is_cell gal) then (
              if should_trace then Printf.eprintf "[Nock:%d] Op11: gal not cell\n%!" !trace_depth;
              raise Exit
            );
            let ref_formula = head gal in
            let gof_formula = tail gal in

            if should_trace then Printf.eprintf "[Nock:%d] Op11: ref_formula is %s, gof_formula is %s\n%!" !trace_depth
              (if is_cell ref_formula then "cell" else "atom")
              (if is_cell gof_formula then "cell" else "atom");

            (* Check if ref_formula looks valid *)
            (match ref_formula with
            | Cell (Atom op, _) when Z.to_int op > 11 ->
                if should_trace then Printf.eprintf "[Nock:%d] Op11: WARNING ref_formula has invalid opcode %d\n%!" !trace_depth (Z.to_int op)
            | _ -> ());

            (* Evaluate both formulas *)
            if should_trace then Printf.eprintf "[Nock:%d] Op11: evaluating ref...\n%!" !trace_depth;
            let _ref = nock_on bus ref_formula in
            if should_trace then Printf.eprintf "[Nock:%d] Op11: evaluating gof...\n%!" !trace_depth;
            let _gof = nock_on bus gof_formula in

            (* For now, scry always fails (returns block)
             * In real Urbit, this would call into the scry handler
             * C Vere calls u3m_soft_esc which can fail
             * We'll return a crash for now *)
            if should_trace then Printf.eprintf "[Nock:%d] Op11: scry not supported, crashing\n%!" !trace_depth;
            raise Exit

        | n ->
            (* Invalid opcode *)
            if should_trace then Printf.eprintf "[Nock:%d] Invalid opcode: %d\n%!" !trace_depth n;
            raise Exit
        )

    | _ ->
        (* Invalid formula structure *)
        if should_trace then Printf.eprintf "[Nock:%d] Invalid formula (not [atom cell] or [cell cell])\n%!" !trace_depth;
        raise Exit
    in
    if should_trace then decr trace_depth;
    result
  with e ->
    if should_trace then decr trace_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