(** sx_browser.ml — OCaml SX kernel compiled to WASM/JS for browser use.

    Exposes the CEK machine, bytecode VM, parser, and primitives as a
    global [SxKernel] object that the JS platform layer binds to.

    Fresh implementation on the ocaml-vm branch — builds on the bytecode
    VM + lazy JIT infrastructure. *)

open Js_of_ocaml
open Sx_types

(* ================================================================== *)
(* Opaque value handle table                                           *)
(*                                                                     *)
(* Non-primitive SX values (lambdas, components, signals, etc.) are    *)
(* stored here and represented on the JS side as objects with an       *)
(* __sx_handle integer key.  Preserves identity across JS↔OCaml.      *)
(* ================================================================== *)

let _next_handle = ref 0
let _handle_table : (int, value) Hashtbl.t = Hashtbl.create 256

let alloc_handle (v : value) : int =
  let id = !_next_handle in
  incr _next_handle;
  Hashtbl.replace _handle_table id v;
  id

let get_handle (id : int) : value =
  match Hashtbl.find_opt _handle_table id with
  | Some v -> v
  | None -> raise (Eval_error (Printf.sprintf "Invalid SX handle: %d" id))

(* JS-side opaque host object table.
   Host objects (DOM elements, console, etc.) are stored here to preserve
   identity across the OCaml↔JS boundary. Represented as Dict with
   __host_handle key on the OCaml side. *)
let _next_host_handle = ref 0
let _alloc_host_handle = Js.Unsafe.pure_js_expr
  "(function() { var t = {}; var n = 0; return { put: function(obj) { var id = n++; t[id] = obj; return id; }, get: function(id) { return t[id]; } }; })()"
let host_put (obj : Js.Unsafe.any) : int =
  let id = !_next_host_handle in
  incr _next_host_handle;
  ignore (Js.Unsafe.meth_call _alloc_host_handle "put" [| obj |]);
  id
let host_get_js (id : int) : Js.Unsafe.any =
  Js.Unsafe.meth_call _alloc_host_handle "get" [| Js.Unsafe.inject id |]

(* ================================================================== *)
(* Global environment                                                  *)
(* ================================================================== *)

(* Force module initialization — these modules register primitives
   in their let () = ... blocks but aren't referenced directly. *)
let () = Sx_scope.clear_all ()

let global_env = make_env ()
let _sx_render_mode = ref false

let call_sx_fn (fn : value) (args : value list) : value =
  let result = Sx_runtime.sx_call fn args in
  !Sx_primitives._sx_trampoline_fn result

(* ================================================================== *)
(* Value conversion: OCaml <-> JS                                      *)
(* ================================================================== *)

(** Tag a JS function with __sx_handle and _type properties. *)
let _tag_fn = Js.Unsafe.pure_js_expr
  "(function(fn, handle, type) { fn.__sx_handle = handle; fn._type = type; return fn; })"

let rec value_to_js (v : value) : Js.Unsafe.any =
  match v with
  | Nil -> Js.Unsafe.inject Js.null
  | Bool b -> Js.Unsafe.inject (Js.bool b)
  | Number n -> Js.Unsafe.inject (Js.number_of_float n)
  | String s -> Js.Unsafe.inject (Js.string s)
  | RawHTML s -> Js.Unsafe.inject (Js.string s)
  | Symbol s ->
    Js.Unsafe.inject (Js.Unsafe.obj [|
      ("_type", Js.Unsafe.inject (Js.string "symbol"));
      ("name", Js.Unsafe.inject (Js.string s)) |])
  | Keyword k ->
    Js.Unsafe.inject (Js.Unsafe.obj [|
      ("_type", Js.Unsafe.inject (Js.string "keyword"));
      ("name", Js.Unsafe.inject (Js.string k)) |])
  | List items | ListRef { contents = items } ->
    let arr = items |> List.map value_to_js |> Array.of_list in
    Js.Unsafe.inject (Js.Unsafe.obj [|
      ("_type", Js.Unsafe.inject (Js.string "list"));
      ("items", Js.Unsafe.inject (Js.array arr)) |])
  | Dict d ->
    (* Check for __host_handle — return original JS object *)
    (match Hashtbl.find_opt d "__host_handle" with
     | Some (Number n) -> host_get_js (int_of_float n)
     | _ ->
       let obj = Js.Unsafe.obj [||] in
       Js.Unsafe.set obj (Js.string "_type") (Js.string "dict");
       Hashtbl.iter (fun k v ->
         Js.Unsafe.set obj (Js.string k) (value_to_js v)) d;
       Js.Unsafe.inject obj)
  (* Callable values: wrap as JS functions with __sx_handle *)
  | Lambda _ | NativeFn _ | Continuation _ | VmClosure _ ->
    let handle = alloc_handle v in
    let inner = Js.wrap_callback (fun args_js ->
      try
        let arg = js_to_value args_js in
        let args = match arg with Nil -> [] | _ -> [arg] in
        let result = call_sx_fn v args in
        value_to_js result
      with
      | Eval_error msg ->
        let fn_info = Printf.sprintf " [callback %s handle=%d]" (type_of v) handle in
        ignore (Js.Unsafe.meth_call
          (Js.Unsafe.get Js.Unsafe.global (Js.string "console"))
          "error" [| Js.Unsafe.inject (Js.string ("[sx] " ^ msg ^ fn_info)) |]);
        Js.Unsafe.inject Js.null
      | exn ->
        let fn_info = Printf.sprintf " [callback %s handle=%d]" (type_of v) handle in
        ignore (Js.Unsafe.meth_call
          (Js.Unsafe.get Js.Unsafe.global (Js.string "console"))
          "error" [| Js.Unsafe.inject (Js.string ("[sx] UNCAUGHT: " ^ Printexc.to_string exn ^ fn_info)) |]);
        Js.Unsafe.inject Js.null) in
    Js.Unsafe.fun_call _tag_fn [|
      Js.Unsafe.inject inner;
      Js.Unsafe.inject handle;
      Js.Unsafe.inject (Js.string (type_of v)) |]
  (* Non-callable compound: tagged object with handle *)
  | _ ->
    let handle = alloc_handle v in
    Js.Unsafe.inject (Js.Unsafe.obj [|
      ("_type", Js.Unsafe.inject (Js.string (type_of v)));
      ("__sx_handle", Js.Unsafe.inject handle) |])

and js_to_value (js : Js.Unsafe.any) : value =
  if Js.Unsafe.equals js Js.null || Js.Unsafe.equals js Js.undefined then Nil
  else
    let ty = Js.to_string (Js.typeof js) in
    match ty with
    | "number" -> Number (Js.float_of_number (Js.Unsafe.coerce js))
    | "boolean" -> Bool (Js.to_bool (Js.Unsafe.coerce js))
    | "string" -> String (Js.to_string (Js.Unsafe.coerce js))
    | "function" ->
      let h = Js.Unsafe.get js (Js.string "__sx_handle") in
      if not (Js.Unsafe.equals h Js.undefined) then
        get_handle (Js.float_of_number (Js.Unsafe.coerce h) |> int_of_float)
      else
        (* Plain JS function — wrap as NativeFn *)
        NativeFn ("js-callback", fun args ->
          let js_args = args |> List.map value_to_js |> Array.of_list in
          js_to_value (Js.Unsafe.fun_call js (Array.map Fun.id js_args)))
    | "object" ->
      let h = Js.Unsafe.get js (Js.string "__sx_handle") in
      if not (Js.Unsafe.equals h Js.undefined) then
        get_handle (Js.float_of_number (Js.Unsafe.coerce h) |> int_of_float)
      else
        let type_field = Js.Unsafe.get js (Js.string "_type") in
        if Js.Unsafe.equals type_field Js.undefined then begin
          if Js.to_bool (Js.Unsafe.global##._Array##isArray js) then begin
            let n = Js.float_of_number (Js.Unsafe.coerce (Js.Unsafe.get js (Js.string "length"))) |> int_of_float in
            List (List.init n (fun i ->
              js_to_value (Js.array_get (Js.Unsafe.coerce js) i |> Js.Optdef.to_option |> Option.get)))
          end else begin
            (* Opaque host object — store in JS-side table, return Dict with __host_handle *)
            let id = host_put js in
            let d = Hashtbl.create 2 in
            Hashtbl.replace d "__host_handle" (Number (float_of_int id));
            Dict d
          end
        end else begin
          let tag = Js.to_string (Js.Unsafe.coerce type_field) in
          match tag with
          | "symbol" -> Symbol (Js.to_string (Js.Unsafe.get js (Js.string "name")))
          | "keyword" -> Keyword (Js.to_string (Js.Unsafe.get js (Js.string "name")))
          | "list" ->
            let items_js = Js.Unsafe.get js (Js.string "items") in
            let n = Js.float_of_number (Js.Unsafe.coerce (Js.Unsafe.get items_js (Js.string "length"))) |> int_of_float in
            List (List.init n (fun i ->
              js_to_value (Js.array_get (Js.Unsafe.coerce items_js) i |> Js.Optdef.to_option |> Option.get)))
          | "dict" ->
            let d = Hashtbl.create 8 in
            let keys = Js.Unsafe.global##._Object##keys js in
            let len = keys##.length in
            for i = 0 to len - 1 do
              let k = Js.to_string (Js.array_get keys i |> Js.Optdef.to_option |> Option.get) in
              if k <> "_type" then
                Hashtbl.replace d k (js_to_value (Js.Unsafe.get js (Js.string k)))
            done;
            Dict d
          | _ -> Nil
        end
    | _ -> Nil

(* ================================================================== *)
(* Side-channel return (bypasses js_of_ocaml stripping properties)     *)
(* ================================================================== *)

let return_via_side_channel (v : Js.Unsafe.any) : Js.Unsafe.any =
  Js.Unsafe.set Js.Unsafe.global (Js.string "__sxR") v; v

(* ================================================================== *)
(* Persistent VM globals — synced with global_env                      *)
(* ================================================================== *)

(* String-keyed mirror of global_env.bindings for VmClosures.
   VmClosures from bytecode modules hold vm_env_ref pointing here.
   Must stay in sync so VmClosures see post-boot definitions. *)
let _vm_globals : (string, value) Hashtbl.t = Hashtbl.create 512
let _in_batch = ref false

(* Sync env→VM: copy all bindings from global_env.bindings to _vm_globals.
   Called after CEK eval/load so VmClosures can see new definitions. *)
let sync_env_to_vm () =
  Hashtbl.iter (fun id v ->
    Hashtbl.replace _vm_globals (unintern id) v
  ) global_env.bindings

(* Hook: intercept env_bind on global_env to also update _vm_globals.
   This ensures VmClosures see new definitions immediately, even during
   a single boot-init call that loads page scripts and components. *)
let () =
  Sx_types._env_bind_hook := Some (fun env name v ->
    if env == global_env then
      Hashtbl.replace _vm_globals name v)

(* ================================================================== *)
(* Core API                                                            *)
(* ================================================================== *)

let api_parse src_js =
  let src = Js.to_string src_js in
  try
    let values = Sx_parser.parse_all src in
    Js.Unsafe.inject (Js.array (values |> List.map value_to_js |> Array.of_list))
  with Parse_error msg ->
    Js.Unsafe.inject (Js.string ("Parse error: " ^ msg))

let api_eval src_js =
  let src = Js.to_string src_js in
  try
    let exprs = Sx_parser.parse_all src in
    let env = Env global_env in
    let result = List.fold_left (fun _acc expr -> Sx_ref.eval_expr expr env) Nil exprs in
    sync_env_to_vm ();
    return_via_side_channel (value_to_js result)
  with
  | Eval_error msg -> Js.Unsafe.inject (Js.string ("Error: " ^ msg))
  | Parse_error msg -> Js.Unsafe.inject (Js.string ("Parse error: " ^ msg))

(** evalVM: compile SX source to bytecode and run through the VM.
    Globals defined with `define` are visible to subsequent evalVM/eval calls.
    This tests the exact same code path as island hydration and click handlers. *)
let api_eval_vm src_js =
  let src = Js.to_string src_js in
  try
    let exprs = Sx_parser.parse_all src in
    let compile_fn = match Hashtbl.find_opt _vm_globals "compile-module" with
      | Some v -> v
      | None -> env_get global_env "compile-module" in
    let code_val = Sx_runtime.trampoline (Sx_runtime.sx_call compile_fn [List exprs]) in
    let code = Sx_vm.code_from_value code_val in
    let result = Sx_vm.execute_module code _vm_globals in
    (* Sync VM globals → CEK env so subsequent eval() calls see defines *)
    Hashtbl.iter (fun name v ->
      let id = intern name in
      if not (Hashtbl.mem global_env.bindings id) then
        Hashtbl.replace global_env.bindings id v
      else (match Hashtbl.find global_env.bindings id, v with
        | VmClosure _, VmClosure _ | _, VmClosure _ -> Hashtbl.replace global_env.bindings id v
        | _ -> ())
    ) _vm_globals;
    return_via_side_channel (value_to_js result)
  with
  | Eval_error msg -> Js.Unsafe.inject (Js.string ("Error: " ^ msg))
  | Parse_error msg -> Js.Unsafe.inject (Js.string ("Parse error: " ^ msg))
  | Not_found -> Js.Unsafe.inject (Js.string "Error: compile-module not loaded")

let api_eval_expr expr_js _env_js =
  let expr = js_to_value expr_js in
  try
    let result = Sx_ref.eval_expr expr (Env global_env) in
    sync_env_to_vm ();
    return_via_side_channel (value_to_js result)
  with Eval_error msg ->
    Js.Unsafe.inject (Js.string ("Error: " ^ msg))

let api_load src_js =
  let src = Js.to_string src_js in
  try
    let exprs = Sx_parser.parse_all src in
    let env = Env global_env in
    let count = ref 0 in
    List.iter (fun expr -> ignore (Sx_ref.eval_expr expr env); incr count) exprs;
    sync_env_to_vm ();
    Js.Unsafe.inject !count
  with
  | Eval_error msg -> Js.Unsafe.inject (Js.string ("Error: " ^ msg))
  | Parse_error msg -> Js.Unsafe.inject (Js.string ("Parse error: " ^ msg))

let api_begin_module_load () =
  (* Snapshot current env into the persistent VM globals table *)
  Hashtbl.clear _vm_globals;
  Hashtbl.iter (fun id v -> Hashtbl.replace _vm_globals (unintern id) v) global_env.bindings;
  _in_batch := true;
  Js.Unsafe.inject true

let api_end_module_load () =
  if !_in_batch then begin
    (* Copy VM globals back to env (bytecode modules defined new symbols) *)
    Hashtbl.iter (fun k v ->
      Hashtbl.replace global_env.bindings (intern k) v
    ) _vm_globals;
    _in_batch := false
  end;
  Js.Unsafe.inject true

let sync_vm_to_env () =
  Hashtbl.iter (fun name v ->
    let id = intern name in
    if not (Hashtbl.mem global_env.bindings id) then
      Hashtbl.replace global_env.bindings id v
    else begin
      (* Update existing binding if the VM has a newer value *)
      let existing = Hashtbl.find global_env.bindings id in
      match existing, v with
      | VmClosure _, VmClosure _ -> Hashtbl.replace global_env.bindings id v
      | _, VmClosure _ -> Hashtbl.replace global_env.bindings id v
      | _ -> ()
    end
  ) _vm_globals

let api_load_module module_js =
  try
    let code_val = js_to_value module_js in
    let code = Sx_vm.code_from_value code_val in
    let _result = Sx_vm.execute_module code _vm_globals in
    sync_vm_to_env ();
    Js.Unsafe.inject (Hashtbl.length _vm_globals)
  with
  | Eval_error msg -> Js.Unsafe.inject (Js.string ("Error: " ^ msg))
  | exn -> Js.Unsafe.inject (Js.string ("Error: " ^ Printexc.to_string exn))

let api_debug_env name_js =
  let name = Js.to_string name_js in
  let id = intern name in
  let found_env = Hashtbl.find_opt global_env.bindings id in
  let found_vm = Hashtbl.find_opt _vm_globals name in
  let total_env = Hashtbl.length global_env.bindings in
  let total_vm = Hashtbl.length _vm_globals in
  let env_s = match found_env with Some v -> "env:" ^ type_of v | None -> "env:MISSING" in
  let vm_s = match found_vm with Some v -> "vm:" ^ type_of v | None -> "vm:MISSING" in
  Js.Unsafe.inject (Js.string (Printf.sprintf "%s %s (env=%d vm=%d)" env_s vm_s total_env total_vm))

let api_compile_module src_js =
  let src = Js.to_string src_js in
  try
    let exprs = Sx_parser.parse_all src in
    let compile_fn = env_get global_env "compile-module" in
    let code = Sx_ref.eval_expr (List [compile_fn; List exprs]) (Env global_env) in
    return_via_side_channel (value_to_js code)
  with
  | Eval_error msg -> Js.Unsafe.inject (Js.string ("Error: " ^ msg))
  | Parse_error msg -> Js.Unsafe.inject (Js.string ("Parse error: " ^ msg))
  | Not_found -> Js.Unsafe.inject (Js.string "Error: compile-module not loaded")

let api_render_to_html expr_js =
  let expr = js_to_value expr_js in
  let prev = !_sx_render_mode in
  _sx_render_mode := true;
  (try
    let html = Sx_render.render_to_html expr global_env in
    _sx_render_mode := prev;
    Js.Unsafe.inject (Js.string html)
  with Eval_error msg ->
    _sx_render_mode := prev;
    Js.Unsafe.inject (Js.string ("Error: " ^ msg)))

let api_stringify v_js =
  Js.Unsafe.inject (Js.string (inspect (js_to_value v_js)))

let api_type_of v_js =
  Js.Unsafe.inject (Js.string (type_of (js_to_value v_js)))

let api_inspect v_js =
  Js.Unsafe.inject (Js.string (inspect (js_to_value v_js)))

let api_engine () =
  Js.Unsafe.inject (Js.string "ocaml-vm-wasm")

let api_register_native name_js callback_js =
  let name = Js.to_string name_js in
  let native_fn args =
    let js_args = args |> List.map value_to_js |> Array.of_list in
    js_to_value (Js.Unsafe.fun_call callback_js [| Js.Unsafe.inject (Js.array js_args) |])
  in
  let v = NativeFn (name, native_fn) in
  ignore (env_bind global_env name v);
  Hashtbl.replace _vm_globals name v;
  Js.Unsafe.inject Js.null

let api_call_fn fn_js args_js =
  try
    let fn = js_to_value fn_js in
    let args = Array.to_list (Array.map js_to_value (Js.to_array (Js.Unsafe.coerce args_js))) in
    return_via_side_channel (value_to_js (call_sx_fn fn args))
  with
  | Eval_error msg ->
    ignore (Js.Unsafe.meth_call
      (Js.Unsafe.get Js.Unsafe.global (Js.string "console"))
      "error" [| Js.Unsafe.inject (Js.string ("[sx] callFn: " ^ msg)) |]);
    Js.Unsafe.inject Js.null
  | exn ->
    ignore (Js.Unsafe.meth_call
      (Js.Unsafe.get Js.Unsafe.global (Js.string "console"))
      "error" [| Js.Unsafe.inject (Js.string ("[sx] callFn: " ^ Printexc.to_string exn)) |]);
    Js.Unsafe.inject Js.null

let api_is_callable fn_js =
  if Js.Unsafe.equals fn_js Js.null || Js.Unsafe.equals fn_js Js.undefined then
    Js.Unsafe.inject (Js.bool false)
  else
    let h = Js.Unsafe.get fn_js (Js.string "__sx_handle") in
    if Js.Unsafe.equals h Js.undefined then Js.Unsafe.inject (Js.bool false)
    else Js.Unsafe.inject (Js.bool (is_callable (get_handle (Js.float_of_number (Js.Unsafe.coerce h) |> int_of_float))))

let api_fn_arity fn_js =
  let h = Js.Unsafe.get fn_js (Js.string "__sx_handle") in
  if Js.Unsafe.equals h Js.undefined then Js.Unsafe.inject (Js.number_of_float (-1.0))
  else
    let v = get_handle (Js.float_of_number (Js.Unsafe.coerce h) |> int_of_float) in
    match v with
    | Lambda l -> Js.Unsafe.inject (Js.number_of_float (float_of_int (List.length l.l_params)))
    | _ -> Js.Unsafe.inject (Js.number_of_float (-1.0))

(* ================================================================== *)
(* Platform bindings (registered in global env)                        *)
(* ================================================================== *)

let () =
  let bind name fn = ignore (env_bind global_env name (NativeFn (name, fn))) in

  (* --- Evaluation --- *)
  bind "cek-eval" (fun args ->
    match args with
    | [String s] -> let e = Sx_parser.parse_all s in (match e with h :: _ -> Sx_ref.eval_expr h (Env global_env) | [] -> Nil)
    | [expr] -> Sx_ref.eval_expr expr (Env global_env)
    | [expr; env_val] -> Sx_ref.eval_expr expr env_val
    | _ -> raise (Eval_error "cek-eval: expected 1-2 args"));

  bind "eval-expr-cek" (fun args ->
    match args with
    | [expr; e] -> Sx_ref.eval_expr expr e
    | [expr] -> Sx_ref.eval_expr expr (Env global_env)
    | _ -> raise (Eval_error "eval-expr-cek: expected 1-2 args"));

  bind "cek-call" (fun args ->
    match args with
    | [f; a] when is_callable f ->
      let arg_list = match a with List l -> l | Nil -> [] | v -> [v] in
      Sx_runtime.trampoline (Sx_runtime.sx_call f arg_list)
    | [f; _] -> raise (Eval_error ("cek-call: not callable: " ^ type_of f))
    | _ -> raise (Eval_error "cek-call: expected (fn args)"));

  bind "sx-parse" (fun args ->
    match args with
    | [String src] -> List (Sx_parser.parse_all src)
    | _ -> raise (Eval_error "sx-parse: expected string"));

  bind "sx-serialize" (fun args ->
    match args with
    | [v] -> String (inspect v)
    | _ -> raise (Eval_error "sx-serialize: expected 1 arg"));

  (* --- Assertions & equality --- *)
  let rec deep_equal a b =
    match a, b with
    | Nil, Nil -> true | Bool a, Bool b -> a = b
    | Number a, Number b -> a = b | String a, String b -> a = b
    | Symbol a, Symbol b -> a = b | Keyword a, Keyword b -> a = b
    | (List a | ListRef { contents = a }), (List b | ListRef { contents = b }) ->
      List.length a = List.length b && List.for_all2 deep_equal a b
    | Dict a, Dict b ->
      let ka = Hashtbl.fold (fun k _ acc -> k :: acc) a [] in
      let kb = Hashtbl.fold (fun k _ acc -> k :: acc) b [] in
      List.length ka = List.length kb &&
      List.for_all (fun k -> Hashtbl.mem b k &&
        deep_equal (Hashtbl.find a k) (Hashtbl.find b k)) ka
    | _ -> false
  in
  bind "equal?" (fun args -> match args with [a; b] -> Bool (deep_equal a b) | _ -> raise (Eval_error "equal?: 2 args"));
  bind "assert" (fun args ->
    match args with
    | [cond] -> if not (sx_truthy cond) then raise (Eval_error "Assertion failed"); Bool true
    | [cond; msg] -> if not (sx_truthy cond) then raise (Eval_error ("Assertion: " ^ value_to_string msg)); Bool true
    | _ -> raise (Eval_error "assert: 1-2 args"));

  bind "try-call" (fun args ->
    match args with
    | [thunk] ->
      (try ignore (Sx_ref.eval_expr (List [thunk]) (Env global_env));
        let d = Hashtbl.create 2 in Hashtbl.replace d "ok" (Bool true); Dict d
      with Eval_error msg ->
        let d = Hashtbl.create 2 in Hashtbl.replace d "ok" (Bool false); Hashtbl.replace d "error" (String msg); Dict d)
    | _ -> raise (Eval_error "try-call: 1 arg"));

  (* --- Bytecode loading from s-expression format ---
     (sxbc version hash (code :arity N :upvalue-count N :bytecode (...) :constants (...)))
     Recursively converts the SX tree into the dict format that loadModule expects. *)
  bind "load-sxbc" (fun args ->
    match args with
    | [List (_ :: _ :: _ :: code_form :: _)] | [List (_ :: _ :: code_form :: _)] ->
      let rec convert_code form =
        match form with
        | List (Symbol "code" :: rest) ->
          let d = Hashtbl.create 8 in
          let rec parse_kv = function
            | Keyword "arity" :: Number n :: rest -> Hashtbl.replace d "arity" (Number n); parse_kv rest
            | Keyword "upvalue-count" :: Number n :: rest -> Hashtbl.replace d "upvalue-count" (Number n); parse_kv rest
            | Keyword "bytecode" :: List nums :: rest ->
              Hashtbl.replace d "bytecode" (List nums); parse_kv rest
            | Keyword "constants" :: List consts :: rest ->
              Hashtbl.replace d "constants" (List (List.map convert_const consts)); parse_kv rest
            | _ :: rest -> parse_kv rest  (* skip unknown keywords *)
            | [] -> ()
          in
          parse_kv rest;
          Dict d
        | _ -> raise (Eval_error ("load-sxbc: expected (code ...), got " ^ type_of form))
      and convert_const = function
        | List (Symbol "code" :: _) as form -> convert_code form
        | List (Symbol "list" :: items) -> List (List.map convert_const items)
        | v -> v  (* strings, numbers, booleans, nil, symbols, keywords pass through *)
      in
      let module_val = convert_code code_form in
      let code = Sx_vm.code_from_value module_val in
      let _result = Sx_vm.execute_module code _vm_globals in
      sync_vm_to_env ();
      Number (float_of_int (Hashtbl.length _vm_globals))
    | _ -> raise (Eval_error "load-sxbc: expected (sxbc version hash (code ...))"));

  (* --- List mutation --- *)
  bind "append!" (fun args ->
    match args with
    | [ListRef r; v] -> r := !r @ [v]; ListRef r
    | [List items; v] -> List (items @ [v])
    | _ -> raise (Eval_error "append!: expected list and value"));

  (* remove! — mutate ListRef in-place, removing by identity (==) *)
  bind "remove!" (fun args ->
    match args with
    | [ListRef r; target] ->
      r := List.filter (fun x -> x != target) !r; ListRef r
    | [List items; target] ->
      List (List.filter (fun x -> x != target) items)
    | _ -> raise (Eval_error "append!: list and value"));

  (* --- Environment ops --- *)
  (* Use unwrap_env for nil/dict tolerance, matching the server kernel *)
  let uw = Sx_runtime.unwrap_env in
  bind "make-env" (fun _ -> Env (make_env ()));
  bind "global-env" (fun _ -> Env global_env);
  bind "env-has?" (fun args -> match args with [e; String k] | [e; Keyword k] -> Bool (env_has (uw e) k) | _ -> raise (Eval_error "env-has?"));
  bind "env-get" (fun args -> match args with [e; String k] | [e; Keyword k] -> env_get (uw e) k | _ -> raise (Eval_error "env-get"));
  bind "env-bind!" (fun args -> match args with [e; String k; v] | [e; Keyword k; v] -> env_bind (uw e) k v | _ -> raise (Eval_error "env-bind!"));
  bind "env-set!" (fun args -> match args with [e; String k; v] | [e; Keyword k; v] -> env_set (uw e) k v | _ -> raise (Eval_error "env-set!"));
  bind "env-extend" (fun args -> match args with [e] -> Env (env_extend (uw e)) | _ -> raise (Eval_error "env-extend"));
  bind "env-merge" (fun args -> match args with [a; b] -> Sx_runtime.env_merge a b | _ -> raise (Eval_error "env-merge"));

  (* --- Type constructors --- *)
  bind "make-symbol" (fun args -> match args with [String s] -> Symbol s | [v] -> Symbol (value_to_string v) | _ -> raise (Eval_error "make-symbol"));
  bind "make-keyword" (fun args -> match args with [String s] -> Keyword s | _ -> raise (Eval_error "make-keyword"));
  bind "keyword-name" (fun args -> match args with [Keyword k] -> String k | _ -> raise (Eval_error "keyword-name"));
  bind "symbol-name" (fun args -> match args with [Symbol s] -> String s | _ -> raise (Eval_error "symbol-name"));

  (* --- Component/Island accessors (must handle both types) --- *)
  bind "component-name" (fun args ->
    match args with [Component c] -> String c.c_name | [Island i] -> String i.i_name | _ -> String "");
  bind "component-closure" (fun args ->
    match args with [Component c] -> Env c.c_closure | [Island i] -> Env i.i_closure | _ -> Dict (Hashtbl.create 0));
  bind "component-params" (fun args ->
    match args with
    | [Component c] -> List (List.map (fun s -> String s) c.c_params)
    | [Island i] -> List (List.map (fun s -> String s) i.i_params)
    | _ -> Nil);
  bind "component-body" (fun args ->
    match args with [Component c] -> c.c_body | [Island i] -> i.i_body | _ -> Nil);
  let has_children_impl = NativeFn ("component-has-children?", fun args ->
    match args with [Component c] -> Bool c.c_has_children | [Island i] -> Bool i.i_has_children | _ -> Bool false) in
  ignore (env_bind global_env "component-has-children" has_children_impl);
  ignore (env_bind global_env "component-has-children?" has_children_impl);
  bind "component-affinity" (fun args ->
    match args with [Component c] -> String c.c_affinity | [Island _] -> String "client" | _ -> String "auto");
  bind "component-param-types" (fun _ -> Nil);
  bind "component-set-param-types!" (fun _ -> Nil);

  (* --- CEK stepping --- *)
  bind "make-cek-state" (fun args -> match args with [c; e; k] -> Sx_ref.make_cek_state c e k | _ -> raise (Eval_error "make-cek-state"));
  bind "cek-step" (fun args -> match args with [s] -> Sx_ref.cek_step s | _ -> raise (Eval_error "cek-step"));
  bind "cek-phase" (fun args -> match args with [s] -> Sx_ref.cek_phase s | _ -> raise (Eval_error "cek-phase"));
  bind "cek-value" (fun args -> match args with [s] -> Sx_ref.cek_value s | _ -> raise (Eval_error "cek-value"));
  bind "cek-terminal?" (fun args -> match args with [s] -> Sx_ref.cek_terminal_p s | _ -> raise (Eval_error "cek-terminal?"));
  bind "cek-kont" (fun args -> match args with [s] -> Sx_ref.cek_kont s | _ -> raise (Eval_error "cek-kont"));
  bind "frame-type" (fun args -> match args with [f] -> Sx_ref.frame_type f | _ -> raise (Eval_error "frame-type"));

  (* --- Strict mode --- *)
  ignore (env_bind global_env "*strict*" (Bool false));
  ignore (env_bind global_env "*prim-param-types*" Nil);
  bind "set-strict!" (fun args -> match args with [v] -> Sx_ref._strict_ref := v; ignore (env_set global_env "*strict*" v); Nil | _ -> Nil);
  bind "set-prim-param-types!" (fun args -> match args with [v] -> Sx_ref._prim_param_types_ref := v; ignore (env_set global_env "*prim-param-types*" v); Nil | _ -> Nil);
  bind "value-matches-type?" (fun args -> match args with [v; t] -> Sx_ref.value_matches_type_p v t | _ -> Nil);

  (* --- Apply --- *)
  bind "apply" (fun args ->
    match args with
    | f :: rest ->
      let all_args = match List.rev rest with List last :: prefix -> List.rev prefix @ last | _ -> rest in
      Sx_runtime.sx_call f all_args
    | _ -> raise (Eval_error "apply"));

  (* --- Scope stack --- *)
  (* Scope primitives (scope-push!, scope-pop!, context, collect!, collected,
     emit!, emitted, scope-emit!, scope-emitted, etc.) are registered by
     Sx_scope module initialization in the primitives table.
     The CEK evaluator falls through to the primitives table when a symbol
     isn't in the env, so these work automatically.
     Only provide-push!/provide-pop! need explicit env bindings as aliases. *)
  bind "provide-push!" (fun args -> match args with [n; v] -> Sx_runtime.provide_push n v | _ -> raise (Eval_error "provide-push!"));
  bind "provide-pop!" (fun args -> match args with [n] -> Sx_runtime.provide_pop n | _ -> raise (Eval_error "provide-pop!"));

  (* --- Fragment / raw HTML --- *)
  bind "<>" (fun args ->
    RawHTML (String.concat "" (List.map (fun a ->
      match a with String s | RawHTML s -> s | Nil -> ""
      | List _ -> Sx_render.render_to_html a global_env
      | _ -> value_to_string a) args)));
  bind "raw!" (fun args ->
    RawHTML (String.concat "" (List.map (fun a ->
      match a with String s | RawHTML s -> s | _ -> value_to_string a) args)));

  bind "define-page-helper" (fun _ -> Nil);

  (* --- Render --- *)
  Sx_render.setup_render_env global_env;
  bind "set-render-active!" (fun _ -> Nil);
  bind "render-active?" (fun _ -> Bool true);
  bind "is-html-tag?" (fun args -> match args with [String s] -> Bool (Sx_render.is_html_tag s) | _ -> Bool false);

  (* --- Render constants needed by web adapters --- *)
  let html_tags = List (List.map (fun s -> String s) Sx_render.html_tags) in
  let void_elements = List (List.map (fun s -> String s) Sx_render.void_elements) in
  let boolean_attrs = List (List.map (fun s -> String s) Sx_render.boolean_attrs) in
  ignore (env_bind global_env "HTML_TAGS" html_tags);
  ignore (env_bind global_env "VOID_ELEMENTS" void_elements);
  ignore (env_bind global_env "BOOLEAN_ATTRS" boolean_attrs);

  (* --- Error handling --- *)
  bind "cek-try" (fun args ->
    match args with
    | [thunk; handler] ->
      (try Sx_ref.cek_call thunk Nil
       with Eval_error msg -> Sx_ref.cek_call handler (List [String msg]))
    | [thunk] ->
      (try let r = Sx_ref.cek_call thunk Nil in
        List [Symbol "ok"; r]
       with Eval_error msg -> List [Symbol "error"; String msg])
    | _ -> Nil);

  (* --- Evaluator bridge functions needed by spec .sx files --- *)
  bind "eval-expr" (fun args ->
    match args with [expr; e] -> Sx_ref.eval_expr expr e | [expr] -> Sx_ref.eval_expr expr (Env global_env) | _ -> Nil);
  bind "trampoline" (fun args -> match args with [v] -> !Sx_primitives._sx_trampoline_fn v | _ -> Nil);
  bind "expand-macro" (fun args ->
    match args with [mac; raw; Env e] -> Sx_ref.expand_macro mac raw (Env e) | [mac; raw] -> Sx_ref.expand_macro mac raw (Env global_env) | _ -> Nil);
  bind "call-lambda" (fun args ->
    match args with
    | [f; a; _] | [f; a] when is_callable f ->
      let arg_list = match a with List l -> l | Nil -> [] | v -> [v] in
      Sx_runtime.trampoline (Sx_runtime.sx_call f arg_list)
    | _ -> raise (Eval_error "call-lambda: expected (fn args env?)"));
  bind "cek-call" (fun args ->
    match args with
    | [f; a] when is_callable f ->
      let arg_list = match a with List l -> l | Nil -> [] | v -> [v] in
      Sx_runtime.trampoline (Sx_runtime.sx_call f arg_list)
    | [f; _] -> raise (Eval_error ("cek-call: not callable: " ^ type_of f))
    | _ -> raise (Eval_error "cek-call: expected (fn args)"));
  bind "cek-eval" (fun args ->
    match args with [expr] -> Sx_ref.eval_expr expr (Env global_env) | [expr; e] -> Sx_ref.eval_expr expr e | _ -> Nil);
  bind "qq-expand-runtime" (fun args ->
    match args with [template] -> Sx_ref.qq_expand template (Env global_env) | [template; Env e] -> Sx_ref.qq_expand template (Env e) | _ -> Nil);

  (* --- Type predicates needed by adapters --- *)
  bind "thunk?" (fun args -> match args with [Thunk _] -> Bool true | _ -> Bool false);
  bind "thunk-expr" (fun args -> match args with [v] -> thunk_expr v | _ -> Nil);
  bind "thunk-env" (fun args -> match args with [v] -> thunk_env v | _ -> Nil);
  bind "lambda?" (fun args -> match args with [Lambda _] -> Bool true | _ -> Bool false);
  bind "macro?" (fun args -> match args with [Macro _] -> Bool true | _ -> Bool false);
  bind "island?" (fun args -> match args with [Island _] -> Bool true | _ -> Bool false);
  bind "component?" (fun args -> match args with [Component _] | [Island _] -> Bool true | _ -> Bool false);
  bind "callable?" (fun args -> match args with [NativeFn _] | [Lambda _] | [Component _] | [Island _] -> Bool true | _ -> Bool false);
  bind "spread?" (fun args -> match args with [Spread _] -> Bool true | _ -> Bool false);
  bind "continuation?" (fun args -> match args with [Continuation _] -> Bool true | _ -> Bool false);
  bind "lambda-params" (fun args -> match args with [Lambda l] -> List (List.map (fun s -> String s) l.l_params) | _ -> List []);
  bind "lambda-body" (fun args -> match args with [Lambda l] -> l.l_body | _ -> Nil);
  bind "lambda-closure" (fun args -> match args with [Lambda l] -> Env l.l_closure | _ -> Dict (Hashtbl.create 0));

  (* --- Core operations needed by adapters --- *)
  bind "spread-attrs" (fun args ->
    match args with [Spread pairs] -> let d = Hashtbl.create 4 in List.iter (fun (k, v) -> Hashtbl.replace d k v) pairs; Dict d | _ -> Dict (Hashtbl.create 0));
  bind "make-spread" (fun args ->
    match args with [Dict d] -> Spread (Hashtbl.fold (fun k v acc -> (k, v) :: acc) d []) | _ -> Nil);
  bind "make-raw-html" (fun args -> match args with [String s] -> RawHTML s | [v] -> RawHTML (value_to_string v) | _ -> Nil);
  bind "raw-html-content" (fun args -> match args with [RawHTML s] -> String s | [String s] -> String s | _ -> String "");
  bind "empty-dict?" (fun args -> match args with [Dict d] -> Bool (Hashtbl.length d = 0) | _ -> Bool true);
  bind "identical?" (fun args -> match args with [a; b] -> Bool (a == b) | _ -> raise (Eval_error "identical?"));
  bind "for-each-indexed" (fun args ->
    match args with
    | [fn_val; List items] | [fn_val; ListRef { contents = items }] ->
      List.iteri (fun i item ->
        ignore (Sx_ref.eval_expr (List [fn_val; Number (float_of_int i); item]) (Env global_env))
      ) items; Nil
    | _ -> Nil);

  (* --- String/number helpers used by orchestration/browser --- *)
  bind "make-sx-expr" (fun args -> match args with [String s] -> SxExpr s | _ -> raise (Eval_error "make-sx-expr"));
  bind "sx-expr-source" (fun args -> match args with [SxExpr s] -> String s | [String s] -> String s | _ -> raise (Eval_error "sx-expr-source"));
  bind "parse-int" (fun args ->
    match args with
    | [String s] -> (try Number (float_of_int (int_of_string s)) with _ -> Nil)
    | [String s; default_val] -> (try Number (float_of_int (int_of_string s)) with _ -> default_val)
    | [Number n] | [Number n; _] -> Number (Float.round n)
    | [_; default_val] -> default_val | _ -> Nil);
  bind "parse-number" (fun args -> match args with [String s] -> (try Number (float_of_string s) with _ -> Nil) | _ -> Nil);

  (* --- Server-only stubs (no-ops in browser) --- *)
  bind "query" (fun _ -> Nil);
  bind "action" (fun _ -> Nil);
  bind "request-arg" (fun args -> match args with [_; d] -> d | _ -> Nil);
  bind "request-method" (fun _ -> String "GET");
  bind "ctx" (fun _ -> Nil);
  bind "helper" (fun _ -> Nil);
  ()

(* ================================================================== *)
(* JIT compilation hook                                                *)
(*                                                                     *)
(* On first call to a named lambda, try to compile it to bytecode via  *)
(* compiler.sx (loaded as an .sx platform file). Compiled closures run *)
(* on the bytecode VM; failures fall back to the CEK interpreter.      *)
(* ================================================================== *)

let _jit_compiling = ref false
let _jit_enabled = ref false

let () =
  Sx_ref.jit_call_hook := Some (fun f args ->
    match f with
    | Lambda l when !_jit_enabled ->
      (match l.l_compiled with
       | Some cl when not (Sx_vm.is_jit_failed cl) ->
         (try Some (Sx_vm.call_closure cl args _vm_globals)
          with Eval_error msg ->
            let fn_name = match l.l_name with Some n -> n | None -> "?" in
            Printf.eprintf "[jit] DISABLED %s — %s\n%!" fn_name msg;
            l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
            None)
       | Some _ -> None
       | None ->
         if !_jit_compiling then None
         else begin
           _jit_compiling := true;
           let compiled = Sx_vm.jit_compile_lambda l _vm_globals in
           _jit_compiling := false;
           (match compiled with
           | Some cl ->
             l.l_compiled <- Some cl;
             (try Some (Sx_vm.call_closure cl args _vm_globals)
              with _ ->
                l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
                None)
           | None -> None)
         end)
    | _ -> None)

let () = ignore (env_bind global_env "enable-jit!" (NativeFn ("enable-jit!", fun _ -> _jit_enabled := true; Nil)))

(* ================================================================== *)
(* Register global SxKernel object                                     *)
(* ================================================================== *)

let () =
  let sx = Js.Unsafe.obj [||] in
  let wrap fn = Js.Unsafe.fun_call
    (Js.Unsafe.pure_js_expr "(function(fn) { return function() { globalThis.__sxR = undefined; var r = fn.apply(null, arguments); return globalThis.__sxR !== undefined ? globalThis.__sxR : r; }; })")
    [| Js.Unsafe.inject (Js.wrap_callback fn) |] in

  Js.Unsafe.set sx (Js.string "parse") (Js.wrap_callback api_parse);
  Js.Unsafe.set sx (Js.string "stringify") (Js.wrap_callback api_stringify);
  Js.Unsafe.set sx (Js.string "eval") (wrap api_eval);
  Js.Unsafe.set sx (Js.string "evalVM") (wrap api_eval_vm);
  Js.Unsafe.set sx (Js.string "evalExpr") (wrap api_eval_expr);
  Js.Unsafe.set sx (Js.string "renderToHtml") (Js.wrap_callback api_render_to_html);
  Js.Unsafe.set sx (Js.string "load") (Js.wrap_callback api_load);
  Js.Unsafe.set sx (Js.string "loadModule") (Js.wrap_callback api_load_module);
  Js.Unsafe.set sx (Js.string "beginModuleLoad") (Js.wrap_callback (fun () -> api_begin_module_load ()));
  Js.Unsafe.set sx (Js.string "endModuleLoad") (Js.wrap_callback (fun () -> api_end_module_load ()));
  Js.Unsafe.set sx (Js.string "compileModule") (wrap api_compile_module);
  Js.Unsafe.set sx (Js.string "typeOf") (Js.wrap_callback api_type_of);
  Js.Unsafe.set sx (Js.string "inspect") (Js.wrap_callback api_inspect);
  Js.Unsafe.set sx (Js.string "engine") (Js.wrap_callback api_engine);
  Js.Unsafe.set sx (Js.string "registerNative") (Js.wrap_callback api_register_native);
  Js.Unsafe.set sx (Js.string "loadSource") (Js.wrap_callback api_load);
  Js.Unsafe.set sx (Js.string "callFn") (wrap api_call_fn);
  Js.Unsafe.set sx (Js.string "isCallable") (Js.wrap_callback api_is_callable);
  Js.Unsafe.set sx (Js.string "fnArity") (Js.wrap_callback api_fn_arity);
  Js.Unsafe.set sx (Js.string "debugEnv") (Js.wrap_callback api_debug_env);

  (* Scope tracing API *)
  Js.Unsafe.set sx (Js.string "scopeTraceOn") (Js.wrap_callback (fun () ->
    Sx_scope.scope_trace_enable (); Js.Unsafe.inject Js.null));
  Js.Unsafe.set sx (Js.string "scopeTraceOff") (Js.wrap_callback (fun () ->
    Sx_scope.scope_trace_disable (); Js.Unsafe.inject Js.null));
  Js.Unsafe.set sx (Js.string "scopeTraceDrain") (Js.wrap_callback (fun () ->
    let log = Sx_scope.scope_trace_drain () in
    Js.Unsafe.inject (Js.array (Array.of_list (List.map (fun s -> Js.Unsafe.inject (Js.string s)) log)))));

  Js.Unsafe.set Js.Unsafe.global (Js.string "SxKernel") sx