(** {1 VM extension registry}

    Holds the live registry of extension opcodes and installs the
    [dispatch] function into [Sx_vm.extension_dispatch_ref] at module
    init time, replacing Phase A's stub.

    See [plans/sx-vm-opcode-extension.md] and [Sx_vm_extension] for the
    extension interface. *)

open Sx_vm_extension

(** The opcode range an extension is allowed to claim.
    Mirrors the partition comment in [Sx_vm]. *)
let extension_min = 200
let extension_max = 247

(** opcode_id → handler *)
let by_id : (int, handler) Hashtbl.t = Hashtbl.create 64

(** opcode_name → opcode_id *)
let by_name : (string, int) Hashtbl.t = Hashtbl.create 64

(** opcode_id → opcode_name (reverse of [by_name]; used by
    [Sx_vm.opcode_name] for disassembly). *)
let name_of_id_table : (int, string) Hashtbl.t = Hashtbl.create 64

(** extension_name → state *)
let states : (string, extension_state) Hashtbl.t = Hashtbl.create 8

(** Registered extension names, newest first. *)
let extensions : string list ref = ref []

(** Dispatch an extension opcode to its registered handler. Raises
    [Sx_vm.Invalid_opcode] if no handler is registered for [op]. *)
let dispatch op vm frame =
  match Hashtbl.find_opt by_id op with
  | Some handler -> handler vm frame
  | None -> raise (Sx_vm.Invalid_opcode op)

(** Register an extension. Fails if the extension name is already
    registered, or if any opcode_id is outside the extension range or
    collides with an already-registered opcode. *)
let register (m : (module EXTENSION)) =
  let module M = (val m) in
  if Hashtbl.mem states M.name then
    failwith (Printf.sprintf
      "Sx_vm_extensions: extension %S already registered" M.name);
  let st = M.init () in
  let ops = M.opcodes st in
  List.iter (fun (id, opname, _h) ->
    if id < extension_min || id > extension_max then
      failwith (Printf.sprintf
        "Sx_vm_extensions: opcode %d (%s) outside extension range %d-%d"
        id opname extension_min extension_max);
    if Hashtbl.mem by_id id then
      failwith (Printf.sprintf
        "Sx_vm_extensions: opcode %d (%s) already registered" id opname);
    if Hashtbl.mem by_name opname then
      failwith (Printf.sprintf
        "Sx_vm_extensions: opcode name %S already registered" opname)
  ) ops;
  Hashtbl.add states M.name st;
  List.iter (fun (id, opname, h) ->
    Hashtbl.add by_id id h;
    Hashtbl.add by_name opname id;
    Hashtbl.add name_of_id_table id opname
  ) ops;
  extensions := M.name :: !extensions

(** Look up the opcode_id for an opcode_name. Returns [None] if no
    extension provides that opcode. *)
let id_of_name name = Hashtbl.find_opt by_name name

(** Look up the opcode_name for an opcode_id. Returns [None] if no
    extension provides that opcode. Used by disassembly. *)
let name_of_id id = Hashtbl.find_opt name_of_id_table id

(** Look up the state of an extension by name. Returns [None] if the
    extension is not registered. *)
let state_of_extension name = Hashtbl.find_opt states name

(** Names of all registered extensions, newest first. *)
let registered_extensions () = !extensions

(** Test-only: clear the registry. Used by unit tests to isolate
    extensions between test cases. The dispatch_ref is left in place. *)
let _reset_for_tests () =
  Hashtbl.clear by_id;
  Hashtbl.clear by_name;
  Hashtbl.clear name_of_id_table;
  Hashtbl.clear states;
  extensions := []

(** Install our [dispatch] into [Sx_vm.extension_dispatch_ref] and our
    [name_of_id] into [Sx_vm.extension_opcode_name_ref], replacing
    the Phase A stubs. Idempotent. Called automatically at module init. *)
let install_dispatch () =
  Sx_vm.extension_dispatch_ref := dispatch;
  Sx_vm.extension_opcode_name_ref := name_of_id

let () = install_dispatch ()

(** Compiler-side opcode lookup: register the [extension-opcode-id]
    primitive. Compilers ([lib/compiler.sx]) call this to emit
    extension opcodes by name. Returns [Integer id] when registered,
    [Nil] otherwise — so missing extensions degrade to a fallback
    rather than failure. *)
let () =
  Sx_primitives.register "extension-opcode-id" (fun args ->
    match args with
    | [Sx_types.String name] ->
      (match id_of_name name with
       | Some id -> Sx_types.Integer id
       | None -> Sx_types.Nil)
    | [Sx_types.Symbol name] ->
      (match id_of_name name with
       | Some id -> Sx_types.Integer id
       | None -> Sx_types.Nil)
    | _ -> raise (Sx_types.Eval_error
        "extension-opcode-id: expected one string or symbol"))