open Noun

(** Nock interpreter with trampoline to match C's tail-call optimization

    This is a more direct port of the C version which uses continue for tail calls
*)

(* Trace tracking *)
let call_count = ref 0
let depth = ref 0
let max_calls = 100
let max_mug_depth = ref (-1)
let show_mugs = ref false

let indent () =
  match !depth with
  | 0 -> ""
  | 1 -> "  "
  | 2 -> "    "
  | 3 -> "      "
  | 4 -> "        "
  | _ -> "          "

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 using trampoline pattern like C *)
let rec nock_on init_bus init_fol =
  let bus = ref init_bus in
  let fol = ref init_fol in

  let rec loop () =
    let my_call = !call_count in

    (* Check if this is opcode 0 (slot lookup) - C doesn't log these *)
    let is_slot = match !fol with
      | Cell { h = Atom { z = op; _ }; _ } when Z.fits_int op && Z.to_int op = 0 -> true
      | _ -> false
    in

    let should_log = my_call < max_calls && not is_slot in

    (* Log entry *)
    if should_log then begin
      let opcode_str = match !fol with
        | Cell { h = Atom { z = op; _ }; _ } when Z.fits_int op -> opcode_name (Z.to_int op)
        | Cell { h = Cell _; _ } -> "CELL(dist)"
        | _ -> "?"
      in
      if !show_mugs || !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;

    incr depth;

    try
      let result = match !fol with
      | Cell { h = hib; t = gal; _ } when is_cell hib ->
          (* Distribution *)
          let poz = nock_on !bus hib in
          let riv = nock_on !bus gal in
          cell poz riv

      | Cell { h = Atom { z = op; _ }; t = gal; _ } ->
          if Z.compare op (Z.of_int max_int) > 0 then raise Exit;
          let opcode = Z.to_int op in

          (match opcode with
          | 0 ->
              if not (is_atom gal) then raise Exit
              else slot (match gal with Atom { z = n; _ } -> n | _ -> raise Exit) !bus

          | 1 ->
              gal

          | 2 ->
              (* Tail-call optimization: set bus/fol and loop *)
              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
              let seb = nock_on !bus b_gal in
              bus := seb;
              fol := nex;
              decr depth;
              if should_log then begin
                if !show_mugs || !depth <= !max_mug_depth then begin
                  let result_mug = mug !bus in
                  Printf.eprintf "%s<<< CONTINUE call #%d depth=%d with bus[mug=0x%lx]\n%!"
                    (indent ()) my_call !depth result_mug
                end
              end;
              loop ()

          | 3 ->
              let gof = nock_on !bus gal in
              if is_cell gof then atom 0 else atom 1

          | 4 ->
              let gof = nock_on !bus gal in
              inc gof

          | 5 ->
              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 ->
              (* Tail-call optimization *)
              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 { z = n; _ } when Z.equal n Z.zero -> c_gal
                | Atom { z = n; _ } when Z.equal n Z.one -> d_gal
                | _ -> raise Exit
              in
              fol := nex;
              decr depth;
              if should_log then begin
                Printf.eprintf "%s<<< CONTINUE call #%d depth=%d (if-branch)\n%!"
                  (indent ()) my_call !depth
              end;
              loop ()

          | 7 ->
              (* Tail-call optimization *)
              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
              bus := bod;
              fol := c_gal;
              decr depth;
              if should_log then begin
                if !show_mugs || !depth <= !max_mug_depth then begin
                  let result_mug = mug !bus in
                  Printf.eprintf "%s<<< CONTINUE call #%d depth=%d with bus[mug=0x%lx]\n%!"
                    (indent ()) my_call !depth result_mug
                end
              end;
              loop ()

          | 8 ->
              (* Tail-call optimization *)
              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
              bus := bod;
              fol := c_gal;
              decr depth;
              if should_log then begin
                if !show_mugs || !depth <= !max_mug_depth then begin
                  let result_mug = mug !bus in
                  Printf.eprintf "%s<<< CONTINUE call #%d depth=%d with bus[mug=0x%lx]\n%!"
                    (indent ()) my_call !depth result_mug
                end
              end;
              loop ()

          | 9 ->
              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 { z; _ } -> z | _ -> raise Exit) seb in
              nock_on seb nex

          | 10 ->
              if not (is_cell gal) then raise Exit;
              let _p_gal = head gal in
              let q_gal = tail gal in
              nock_on !bus q_gal

          | 11 ->
              if not (is_cell gal) then raise Exit;
              let p_gal = head gal in
              let q_gal = tail gal in
              if is_cell p_gal then
                nock_on !bus q_gal
              else
                nock_on !bus q_gal

          | _ ->
              raise Exit
          )

      | _ ->
          raise Exit
      in

      (* Log exit and return *)
      decr depth;
      if should_log then begin
        if !show_mugs || !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 ->
      decr depth;
      raise e
  in

  loop ()

let nock subject formula =
  nock_on subject formula