rose-ash/hosts/ocaml/bin/sx_server.ml

(** SX coroutine subprocess server.

    Persistent process that accepts commands on stdin and writes
    responses on stdout.  All messages are single-line SX expressions,
    newline-delimited.

    Protocol:
      Python → OCaml: (ping), (load path), (load-source src),
                       (eval src), (render src), (reset),
                       (io-response value)
      OCaml → Python: (ready), (ok), (ok value), (error msg),
                       (io-request name args...)

    IO primitives (query, action, request-arg, request-method, ctx)
    yield (io-request ...) and block on stdin for (io-response ...). *)

module Sx_types = Sx.Sx_types
module Sx_parser = Sx.Sx_parser
module Sx_primitives = Sx.Sx_primitives
module Sx_runtime = Sx.Sx_runtime
module Sx_ref = Sx.Sx_ref
module Sx_render = Sx.Sx_render
module Sx_vm = Sx.Sx_vm

open Sx_types


(* ====================================================================== *)
(* Output helpers                                                          *)
(* ====================================================================== *)

(** Escape a string for embedding in an SX string literal. *)
let escape_sx_string s =
  let buf = Buffer.create (String.length s + 16) in
  String.iter (function
    | '"'  -> Buffer.add_string buf "\\\""
    | '\\' -> Buffer.add_string buf "\\\\"
    | '\n' -> Buffer.add_string buf "\\n"
    | '\r' -> Buffer.add_string buf "\\r"
    | '\t' -> Buffer.add_string buf "\\t"
    | c    -> Buffer.add_char buf c) s;
  Buffer.contents buf

(** Serialize a value to SX text (for io-request args). *)
let rec serialize_value = function
  | Nil -> "nil"
  | Bool true -> "true"
  | Bool false -> "false"
  | Number n ->
    if Float.is_integer n then string_of_int (int_of_float n)
    else Printf.sprintf "%g" n
  | String s -> "\"" ^ escape_sx_string s ^ "\""
  | Symbol s -> s
  | Keyword k -> ":" ^ k
  | List items | ListRef { contents = items } ->
    "(list " ^ String.concat " " (List.map serialize_value items) ^ ")"
  | Dict d ->
    let pairs = Hashtbl.fold (fun k v acc ->
      (Printf.sprintf ":%s %s" k (serialize_value v)) :: acc) d [] in
    "{" ^ String.concat " " pairs ^ "}"
  | RawHTML s -> "\"" ^ escape_sx_string s ^ "\""
  | _ -> "nil"

let send line =
  print_string line;
  print_char '\n';
  flush stdout

let send_ok () = send "(ok)"
let send_ok_value v = send (Printf.sprintf "(ok %s)" (serialize_value v))
let send_error msg = send (Printf.sprintf "(error \"%s\")" (escape_sx_string msg))

(** Length-prefixed binary send — handles any content without escaping.
    Sends: (ok-len N)\n followed by exactly N bytes of raw data, then \n.
    Python reads the length line, then reads exactly N bytes. *)
let send_ok_blob s =
  let n = String.length s in
  Printf.printf "(ok-len %d)\n" n;
  print_string s;
  print_char '\n';
  flush stdout

(** Send a string value — use blob for anything that might contain
    newlines, quotes, or be large. *)
let send_ok_string s = send_ok_blob s

(** Send raw SX wire format — may contain newlines in string literals. *)
let send_ok_raw s = send_ok_blob s


(* ====================================================================== *)
(* IO bridge — primitives that yield to Python                             *)
(* ====================================================================== *)

(** Read a line from stdin (blocking). *)
let read_line_blocking () =
  try Some (input_line stdin)
  with End_of_file -> None

(** Read exactly N bytes from stdin (blocking). *)
let read_exact_bytes n =
  let buf = Bytes.create n in
  really_input stdin buf 0 n;
  Bytes.to_string buf

(** Read a length-prefixed blob from stdin.
    Expects the next line to be "(blob N)" where N is byte count,
    followed by exactly N bytes of raw data, then a newline. *)
let read_blob () =
  match read_line_blocking () with
  | None -> raise (Eval_error "read_blob: stdin closed")
  | Some line ->
    let line = String.trim line in
    match Sx_parser.parse_all line with
    | [List [Symbol "blob"; Number n]] ->
      let len = int_of_float n in
      let data = read_exact_bytes len in
      (* consume trailing newline *)
      (try ignore (input_line stdin) with End_of_file -> ());
      data
    | _ -> raise (Eval_error ("read_blob: expected (blob N), got: " ^ line))

(** Batch IO mode — collect requests during aser-slot, resolve after. *)
let io_batch_mode = ref false
let io_queue : (int * string * value list) list ref = ref []
let io_counter = ref 0

(** Module-level scope stacks — shared between make_server_env (aser
    scope-push!/pop!) and step-sf-context (via get-primitive "scope-peek"). *)
let _scope_stacks : (string, value list) Hashtbl.t = Hashtbl.create 8

let () = Sx_primitives.register "scope-push!" (fun args ->
  match args with
  | [String name; value] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    Hashtbl.replace _scope_stacks name (value :: stack); Nil
  | _ -> Nil)

let () = Sx_primitives.register "scope-pop!" (fun args ->
  match args with
  | [String name] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with _ :: rest -> Hashtbl.replace _scope_stacks name rest | [] -> ()); Nil
  | _ -> Nil)

let () = Sx_primitives.register "scope-peek" (fun args ->
  match args with
  | [String name] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with v :: _ -> v | [] -> Nil)
  | _ -> Nil)

let () = Sx_primitives.register "context" (fun args ->
  match args with
  | [String name] | [String name; _] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack, args with
     | v :: _, _ -> v
     | [], [_; default_val] -> default_val
     | [], _ -> Nil)
  | _ -> Nil)

(** collect! — lazy scope accumulator. Creates root scope if missing,
    emits value (deduplicates).  Used by cssx and spread components. *)
let () = Sx_primitives.register "collect!" (fun args ->
  match args with
  | [String name; value] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with
     | List items :: rest ->
       if not (List.mem value items) then
         Hashtbl.replace _scope_stacks name (List (items @ [value]) :: rest)
     | [] ->
       (* Lazy root scope — create with the value *)
       Hashtbl.replace _scope_stacks name [List [value]]
     | _ :: _ -> ());
    Nil
  | _ -> Nil)

let () = Sx_primitives.register "collected" (fun args ->
  match args with
  | [String name] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with List items :: _ -> List items | _ -> List [])
  | _ -> List [])

let () = Sx_primitives.register "clear-collected!" (fun args ->
  match args with
  | [String name] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with
     | _ :: rest -> Hashtbl.replace _scope_stacks name (List [] :: rest)
     | [] -> Hashtbl.replace _scope_stacks name [List []]);
    Nil
  | _ -> Nil)

let () = Sx_primitives.register "scope-emit!" (fun args ->
  match args with
  | [String name; value] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with
     | List items :: rest ->
       Hashtbl.replace _scope_stacks name (List (items @ [value]) :: rest)
     | Nil :: rest ->
       Hashtbl.replace _scope_stacks name (List [value] :: rest)
     | [] ->
       (* Lazy root scope *)
       Hashtbl.replace _scope_stacks name [List [value]]
     | _ :: _ -> ());
    Nil
  | _ -> Nil)

let () = Sx_primitives.register "clear-collected!" (fun args ->
  match args with
  | [String name] ->
    let stack = try Hashtbl.find _scope_stacks name with Not_found -> [] in
    (match stack with
     | _ :: rest -> Hashtbl.replace _scope_stacks name (List [] :: rest)
     | [] -> ());
    Nil
  | _ -> Nil)

(** Helpers safe to defer — pure functions whose results are only used
    as rendering output (inlined into SX wire format), not in control flow. *)
let batchable_helpers = [
  "highlight"; "component-source"
]

let is_batchable name args =
  name = "helper" &&
  match args with
  | String h :: _ -> List.mem h batchable_helpers
  | _ -> false

(** Send an io-request — batch mode returns placeholder, else blocks. *)
let io_request name args =
  if !io_batch_mode && is_batchable name args then begin
    incr io_counter;
    let id = !io_counter in
    io_queue := (id, name, args) :: !io_queue;
    (* Return SxExpr so serialize/inspect passes it through unquoted *)
    SxExpr (Printf.sprintf "(\xc2\xabIO:%d\xc2\xbb)" id)
  end else begin
    let args_str = String.concat " " (List.map serialize_value args) in
    send (Printf.sprintf "(io-request \"%s\" %s)" name args_str);
    (* Block on stdin for io-response *)
    match read_line_blocking () with
    | None -> raise (Eval_error "IO bridge: stdin closed while waiting for io-response")
    | Some line ->
      let exprs = Sx_parser.parse_all line in
      match exprs with
      | [List [Symbol "io-response"; value]] -> value
      | [List (Symbol "io-response" :: values)] ->
        (match values with
         | [v] -> v
         | _ -> List values)
      | _ -> raise (Eval_error ("IO bridge: unexpected response: " ^ line))
  end

(** Flush batched IO: send all requests, read all responses, replace placeholders. *)
let flush_batched_io result_str =
  let queue = List.rev !io_queue in
  io_queue := [];
  io_counter := 0;
  if queue = [] then result_str
  else begin
    (* Send all batched requests with IDs *)
    List.iter (fun (id, name, args) ->
      let args_str = String.concat " " (List.map serialize_value args) in
      send (Printf.sprintf "(io-request %d \"%s\" %s)" id name args_str)
    ) queue;
    send (Printf.sprintf "(io-done %d)" (List.length queue));
    (* Read all responses and replace placeholders *)
    let final = ref result_str in
    List.iter (fun (id, _, _) ->
      match read_line_blocking () with
      | Some line ->
        let exprs = Sx_parser.parse_all line in
        let value_str = match exprs with
          | [List [Symbol "io-response"; String s]]
          | [List [Symbol "io-response"; SxExpr s]] -> s
          | [List [Symbol "io-response"; v]] -> serialize_value v
          | _ -> "nil"
        in
        let placeholder = Printf.sprintf "(\xc2\xabIO:%d\xc2\xbb)" id in
        (* Replace all occurrences of this placeholder *)
        let plen = String.length placeholder in
        let buf = Buffer.create (String.length !final) in
        let pos = ref 0 in
        let s = !final in
        let slen = String.length s in
        while !pos <= slen - plen do
          if String.sub s !pos plen = placeholder then begin
            Buffer.add_string buf value_str;
            pos := !pos + plen
          end else begin
            Buffer.add_char buf s.[!pos];
            incr pos
          end
        done;
        if !pos < slen then
          Buffer.add_substring buf s !pos (slen - !pos);
        final := Buffer.contents buf
      | None -> raise (Eval_error "IO batch: stdin closed")
    ) queue;
    !final
  end

(** Bind IO primitives into the environment. *)
let setup_io_env env =
  let bind name fn =
    ignore (env_bind env name (NativeFn (name, fn)))
  in

  bind "query" (fun args ->
    match args with
    | service :: query_name :: rest ->
      io_request "query" (service :: query_name :: rest)
    | _ -> raise (Eval_error "query: expected (query service name ...)"));

  bind "action" (fun args ->
    match args with
    | service :: action_name :: rest ->
      io_request "action" (service :: action_name :: rest)
    | _ -> raise (Eval_error "action: expected (action service name ...)"));

  bind "request-arg" (fun args ->
    match args with
    | [name] -> io_request "request-arg" [name]
    | [name; default] ->
      let result = io_request "request-arg" [name] in
      if result = Nil then default else result
    | _ -> raise (Eval_error "request-arg: expected 1-2 args"));

  bind "request-method" (fun _args ->
    io_request "request-method" []);

  bind "ctx" (fun args ->
    match args with
    | [key] -> io_request "ctx" [key]
    | _ -> raise (Eval_error "ctx: expected 1 arg"));

  bind "call-lambda" (fun args ->
    match args with
    | [fn_val; List call_args; Env e] ->
      Sx_ref.eval_expr (List (fn_val :: call_args)) (Env e)
    | [fn_val; List call_args] ->
      Sx_ref.eval_expr (List (fn_val :: call_args)) (Env env)
    | _ -> raise (Eval_error "call-lambda: expected (fn args env?)"));

  (* Register HO forms as callable NativeFn — the CEK machine handles them
     as special forms, but the VM needs them as callable values in globals. *)
  let ho_via_cek name =
    bind name (fun args ->
      Sx_ref.eval_expr (List (Symbol name :: args)) (Env env))
  in
  List.iter ho_via_cek [
    "map"; "map-indexed"; "filter"; "reduce"; "some"; "every?"; "for-each";
  ];

  (* Generic helper call — dispatches to Python page helpers *)
  bind "helper" (fun args ->
    io_request "helper" args)


(* ====================================================================== *)
(* Environment setup                                                       *)
(* ====================================================================== *)

let make_server_env () =
  let env = make_env () in

  (* Evaluator bindings — same as run_tests.ml's make_test_env,
     but only the ones needed for rendering (not test helpers). *)
  let bind name fn =
    ignore (env_bind env name (NativeFn (name, fn)))
  in

  bind "assert" (fun args ->
    match args with
    | [cond] ->
      if not (sx_truthy cond) then raise (Eval_error "Assertion failed");
      Bool true
    | [cond; String msg] ->
      if not (sx_truthy cond) then raise (Eval_error ("Assertion error: " ^ msg));
      Bool true
    | [cond; msg] ->
      if not (sx_truthy cond) then
        raise (Eval_error ("Assertion error: " ^ value_to_string msg));
      Bool true
    | _ -> raise (Eval_error "assert: expected 1-2 args"));

  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"));

  (* HTML renderer — OCaml render module provides the shell renderer;
     adapter-html.sx provides the SX-level render-to-html *)
  Sx_render.setup_render_env env;

  (* Render-mode flags *)
  bind "set-render-active!" (fun _args -> Nil);
  bind "render-active?" (fun _args -> Bool true);

  (* Browser APIs — no-op stubs for SSR *)
  bind "local-storage-get" (fun _args -> Nil);
  bind "local-storage-set" (fun _args -> Nil);
  bind "dom-listen" (fun _args -> NativeFn ("noop", fun _ -> Nil));
  bind "dom-dispatch" (fun _args -> Nil);
  bind "dom-set-data" (fun _args -> Nil);
  bind "dom-get-data" (fun _args -> Nil);
  bind "event-detail" (fun _args -> Nil);
  bind "promise-then" (fun _args -> Nil);
  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 "schedule-idle" (fun _args -> Nil);
  bind "dom-query" (fun _args -> Nil);
  bind "dom-query-all" (fun _args -> List []);
  bind "dom-set-prop" (fun _args -> Nil);
  bind "dom-get-attr" (fun _args -> Nil);
  bind "dom-set-attr" (fun _args -> Nil);
  bind "dom-text-content" (fun _args -> String "");
  bind "dom-set-text-content" (fun _args -> Nil);
  bind "dom-body" (fun _args -> Nil);

  (* Raw HTML — platform primitives for adapter-html.sx *)
  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 "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 env))
      ) items; Nil
    | _ -> Nil);

  (* Scope stack — platform primitives for render-time dynamic scope.
     Used by aser for spread/provide/emit patterns.
     Module-level so step-sf-context can check it via get-primitive. *)
  let scope_stacks = _scope_stacks in
  bind "scope-push!" (fun args ->
    match args with
    | [String name; value] ->
      let stack = try Hashtbl.find scope_stacks name with Not_found -> [] in
      Hashtbl.replace scope_stacks name (value :: stack); Nil
    | _ -> Nil);
  bind "scope-pop!" (fun args ->
    match args with
    | [String name] ->
      let stack = try Hashtbl.find scope_stacks name with Not_found -> [] in
      (match stack with
       | _ :: rest -> Hashtbl.replace scope_stacks name rest
       | [] -> ()); Nil
    | _ -> Nil);
  bind "scope-peek" (fun args ->
    match args with
    | [String name] ->
      let stack = try Hashtbl.find scope_stacks name with Not_found -> [] in
      (match stack with v :: _ -> v | [] -> Nil)
    | _ -> Nil);
  (* scope-emit! / scope-peek — Hashtbl-based scope primitives for aser.
     Different names from emit!/emitted to avoid CEK special form conflict. *)
  bind "scope-emit!" (fun args ->
    match args with
    | [String name; value] ->
      let stack = try Hashtbl.find scope_stacks name with Not_found -> [] in
      (match stack with
       | List items :: rest ->
         Hashtbl.replace scope_stacks name (List (items @ [value]) :: rest)
       | Nil :: rest ->
         Hashtbl.replace scope_stacks name (List [value] :: rest)
       | _ :: _ -> ()
       | [] -> ()); Nil
    | _ -> Nil);

  (* context — scope lookup.  The CEK handles this as a special form
     by walking continuation frames, but compiled VM code needs it as
     a function that reads from the scope_stacks hashtable. *)
  let context_impl = NativeFn ("context", fun args ->
    match args with
    | [String name] | [String name; _] ->
      let stack = try Hashtbl.find scope_stacks name with Not_found -> [] in
      (match stack, args with
       | v :: _, _ -> v
       | [], [_; default_val] -> default_val
       | [], _ -> Nil)
    | _ -> Nil) in
  ignore (env_bind env "sx-context" context_impl);
  ignore (env_bind env "context" context_impl);

  (* qq-expand-runtime — quasiquote expansion at runtime.
     The bytecode compiler emits CALL_PRIM "qq-expand-runtime" for
     backtick expressions. Expands the template with the calling env. *)
  bind "qq-expand-runtime" (fun args ->
    match args with
    | [template] ->
      (* qq_expand needs an env — use the kernel env *)
      Sx_ref.qq_expand template (Env env)
    | [template; Env e] -> Sx_ref.qq_expand template (Env e)
    | _ -> Nil);

  (* Evaluator bridge — aser calls these spec functions.
     Route to the OCaml CEK machine. *)
  bind "eval-expr" (fun args ->
    match args with
    | [expr; Env e] -> Sx_ref.eval_expr expr (Env e)
    | [expr] -> Sx_ref.eval_expr expr (Env env)
    | _ -> raise (Eval_error "eval-expr: expected (expr env?)"));
  bind "trampoline" (fun args ->
    match args with
    | [v] ->
      (* sf-letrec returns thunks — resolve them *)
      let rec resolve v = match v with
        | Thunk (expr, env) -> resolve (Sx_ref.eval_expr expr (Env env))
        | _ -> v
      in resolve v
    | _ -> raise (Eval_error "trampoline: expected 1 arg"));
  bind "call-lambda" (fun args ->
    match args with
    | [fn_val; List call_args; Env e] ->
      Sx_ref.eval_expr (List (fn_val :: call_args)) (Env e)
    | [fn_val; List call_args] ->
      Sx_ref.eval_expr (List (fn_val :: call_args)) (Env env)
    | _ -> raise (Eval_error "call-lambda: expected (fn args env?)"));
  bind "cek-call" (fun args ->
    match args with
    | [fn_val; List call_args] -> Sx_ref.cek_call fn_val (List call_args)
    | [fn_val; Nil] -> Sx_ref.cek_call fn_val (List [])
    | [fn_val] -> Sx_ref.cek_call fn_val (List [])
    | _ -> Nil);
  bind "expand-macro" (fun args ->
    match args with
    | [Macro m; List macro_args; Env e] ->
      let body_env = { bindings = Hashtbl.create 16; parent = Some e } in
      List.iteri (fun i p ->
        let v = if i < List.length macro_args then List.nth macro_args i else Nil in
        Hashtbl.replace body_env.bindings p v
      ) m.m_params;
      Sx_ref.eval_expr m.m_body (Env body_env)
    | _ -> raise (Eval_error "expand-macro: expected (macro args env)"));

  (* Register <> as a special form — evaluates all children, returns list *)
  ignore (Sx_ref.register_special_form (String "<>") (NativeFn ("<>", fun args ->
    List (List.map (fun a -> Sx_ref.eval_expr a (Env env)) args))));


  (* Missing primitives that may be referenced *)
  bind "upcase" (fun args ->
    match args with
    | [String s] -> String (String.uppercase_ascii s)
    | _ -> raise (Eval_error "upcase: expected string"));

  bind "downcase" (fun args ->
    match args with
    | [String s] -> String (String.lowercase_ascii s)
    | _ -> raise (Eval_error "downcase: expected string"));

  bind "make-keyword" (fun args ->
    match args with
    | [String s] -> Keyword s
    | _ -> raise (Eval_error "make-keyword: expected string"));

  (* Type predicates and accessors — platform interface for aser *)
  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 "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));
  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 "spread?" (fun args ->
    match args with [Spread _] -> Bool true | _ -> Bool false);
  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] ->
      let pairs = Hashtbl.fold (fun k v acc -> (k, v) :: acc) d [] in
      Spread pairs
    | _ -> Nil);
  bind "is-html-tag?" (fun args ->
    match args with
    | [String s] -> Bool (Sx_render.is_html_tag s)
    | _ -> Bool false);

  (* Spec evaluator helpers needed by render.sx when loaded at runtime *)
  bind "random-int" (fun args ->
    match args with
    | [Number lo; Number hi] ->
      let lo = int_of_float lo and hi = int_of_float hi in
      Number (float_of_int (lo + Random.int (max 1 (hi - lo + 1))))
    | _ -> raise (Eval_error "random-int: expected (low high)"));

  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 "json-encode" (fun args -> io_request "helper" (String "json-encode" :: args));
  bind "into" (fun args -> io_request "helper" (String "into" :: args));

  bind "sleep" (fun args -> io_request "sleep" args);
  bind "set-response-status" (fun args -> io_request "set-response-status" args);
  bind "set-response-header" (fun args -> io_request "set-response-header" args);

  (* Application constructs — no-ops in the kernel, but needed so
     handler/page files load successfully (their define forms get evaluated) *)
  ignore (Sx_ref.register_special_form (String "defhandler") (NativeFn ("defhandler", fun _args -> Nil)));
  ignore (Sx_ref.register_special_form (String "defpage") (NativeFn ("defpage", fun _args -> Nil)));

  bind "cond-scheme?" (fun args ->
    match args with
    | [clauses] -> Sx_ref.cond_scheme_p clauses
    | _ -> Bool false);
  bind "is-else-clause?" (fun args ->
    match args with
    | [test] -> Sx_ref.is_else_clause test
    | _ -> Bool false);
  bind "primitive?" (fun args ->
    match args with
    | [String name] ->
      (* Check if name is bound in the env as a NativeFn *)
      (try match env_get env name with NativeFn _ -> Bool true | _ -> Bool false
       with _ -> Bool false)
    | _ -> Bool false);
  bind "get-primitive" (fun args ->
    match args with
    | [String name] ->
      (try env_get env name with _ -> Nil)
    | _ -> Nil);

  (* Character classification — platform primitives for spec/parser.sx *)
  bind "ident-start?" (fun args ->
    match args with
    | [String s] when String.length s = 1 ->
      let c = s.[0] in
      Bool (c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c = '_' || c = '~'
            || c = '!' || c = '?' || c = '+' || c = '-' || c = '*' || c = '/'
            || c = '=' || c = '<' || c = '>' || c = '&' || c = '|' || c = '%'
            || c = '^' || c = '$' || c = '#')
    | _ -> Bool false);
  bind "ident-char?" (fun args ->
    match args with
    | [String s] when String.length s = 1 ->
      let c = s.[0] in
      Bool (c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c = '_' || c = '~'
            || c = '!' || c = '?' || c = '+' || c = '-' || c = '*' || c = '/'
            || c = '=' || c = '<' || c = '>' || c = '&' || c = '|' || c = '%'
            || c = '^' || c = '$' || c = '#'
            || c >= '0' && c <= '9' || c = '.' || c = ':')
    | _ -> Bool false);
  bind "char-numeric?" (fun args ->
    match args with
    | [String s] when String.length s = 1 ->
      Bool (s.[0] >= '0' && s.[0] <= '9')
    | _ -> Bool false);
  bind "parse-number" (fun args ->
    match args with
    | [String s] ->
      (try Number (float_of_string s)
       with _ -> Nil)
    | _ -> Nil);

  bind "escape-string" (fun args ->
    match args with
    | [String s] ->
      let buf = Buffer.create (String.length s) in
      String.iter (fun c -> match c with
        | '"'  -> Buffer.add_string buf "\\\""
        | '\\' -> Buffer.add_string buf "\\\\"
        | '\n' -> Buffer.add_string buf "\\n"
        | '\r' -> Buffer.add_string buf "\\r"
        | '\t' -> Buffer.add_string buf "\\t"
        | c    -> Buffer.add_char buf c) s;
      String (Buffer.contents buf)
    | _ -> raise (Eval_error "escape-string: expected string"));

  bind "string-length" (fun args ->
    match args with
    | [String s] -> Number (float_of_int (String.length s))
    | _ -> raise (Eval_error "string-length: expected string"));

  bind "dict-get" (fun args ->
    match args with
    | [Dict d; String k] -> dict_get d k
    | [Dict d; Keyword k] -> dict_get d k
    | _ -> raise (Eval_error "dict-get: expected dict and key"));

  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: expected function and args"));

  bind "equal?" (fun args ->
    match args with
    | [a; b] -> Bool (a = b)
    | _ -> raise (Eval_error "equal?: expected 2 args"));

  bind "identical?" (fun args ->
    match args with
    | [a; b] -> Bool (a == b)
    | _ -> raise (Eval_error "identical?: expected 2 args"));

  bind "make-sx-expr" (fun args ->
    match args with
    | [String s] -> SxExpr s
    | _ -> raise (Eval_error "make-sx-expr: expected string"));

  bind "sx-expr-source" (fun args ->
    match args with
    | [SxExpr s] -> String s
    | [String s] -> String s
    | _ -> raise (Eval_error "sx-expr-source: expected sx-expr or string"));

  bind "make-continuation" (fun args ->
    match args with
    | [f] ->
      let k v = Sx_runtime.sx_call f [v] in
      Continuation (k, None)
    | _ -> raise (Eval_error "make-continuation: expected 1 arg"));

  bind "continuation?" (fun args ->
    match args with
    | [Continuation _] -> Bool true
    | [_] -> Bool false
    | _ -> raise (Eval_error "continuation?: expected 1 arg"));

  bind "make-symbol" (fun args ->
    match args with
    | [String s] -> Symbol s
    | [v] -> Symbol (value_to_string v)
    | _ -> raise (Eval_error "make-symbol: expected 1 arg"));

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

  (* Env operations *)
  bind "env-get" (fun args ->
    match args with
    | [Env e; String k] -> env_get e k
    | [Env e; Keyword k] -> env_get e k
    | _ -> raise (Eval_error "env-get: expected env and string"));

  bind "env-has?" (fun args ->
    match args with
    | [Env e; String k] -> Bool (env_has e k)
    | [Env e; Keyword k] -> Bool (env_has e k)
    | _ -> raise (Eval_error "env-has?: expected env and string"));

  bind "env-bind!" (fun args ->
    match args with
    | [Env e; String k; v] -> env_bind e k v
    | [Env e; Keyword k; v] -> env_bind e k v
    | _ -> raise (Eval_error "env-bind!: expected env, key, value"));

  bind "env-set!" (fun args ->
    match args with
    | [Env e; String k; v] -> env_set e k v
    | [Env e; Keyword k; v] -> env_set e k v
    | _ -> raise (Eval_error "env-set!: expected env, key, value"));

  bind "env-extend" (fun args ->
    match args with
    | [Env e] -> Env (env_extend e)
    | _ -> raise (Eval_error "env-extend: expected env"));

  bind "env-merge" (fun args ->
    match args with
    | [Env a; Env b] -> Env (env_merge a b)
    | _ -> raise (Eval_error "env-merge: expected 2 envs"));

  (* Strict mode state *)
  ignore (env_bind env "*strict*" (Bool false));
  ignore (env_bind env "*prim-param-types*" Nil);

  bind "set-strict!" (fun args ->
    match args with
    | [v] ->
      Sx_ref._strict_ref := v;
      ignore (env_set env "*strict*" v); Nil
    | _ -> raise (Eval_error "set-strict!: expected 1 arg"));

  bind "set-prim-param-types!" (fun args ->
    match args with
    | [v] ->
      Sx_ref._prim_param_types_ref := v;
      ignore (env_set env "*prim-param-types*" v); Nil
    | _ -> raise (Eval_error "set-prim-param-types!: expected 1 arg"));

  bind "component-param-types" (fun _args -> Nil);
  bind "component-set-param-types!" (fun _args -> Nil);

  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 env "component-has-children" has_children_impl);
  ignore (env_bind 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 "keyword-name" (fun args ->
    match args with
    | [Keyword k] -> String k
    | _ -> raise (Eval_error "keyword-name: expected keyword"));

  bind "symbol-name" (fun args ->
    match args with
    | [Symbol s] -> String s
    | _ -> raise (Eval_error "symbol-name: expected symbol"));

  (* IO primitives *)
  setup_io_env env;

  (* Initialize trampoline ref so HO primitives (map, filter, etc.)
     can call SX lambdas. Must be done here (not sx_runtime.ml)
     because Sx_ref is only available at the binary level. *)
  Sx_primitives._sx_trampoline_fn := (fun v ->
    match v with
    | Thunk (body, closure_env) -> Sx_ref.eval_expr body (Env closure_env)
    | other -> other);

  env


(* ====================================================================== *)
(* JIT hook registration                                                    *)
(* ====================================================================== *)

(** Register the JIT call hook.  Called once after the compiler is loaded
    into the kernel env.  The hook handles both cached execution (bytecode
    already compiled) and first-call compilation (invoke compiler.sx via
    CEK, cache result).  cek_call checks this before CEK dispatch. *)
let _jit_compiling = ref false  (* re-entrancy guard *)


let register_jit_hook env =
  Sx_ref.jit_call_hook := Some (fun f args ->
    match f with
    | Lambda l ->
      (match l.l_compiled with
       | Some cl when not (Sx_vm.is_jit_failed cl) ->
         (* Cached bytecode — run on VM, fall back to CEK on runtime error *)
         (try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
          with _ -> None)
       | Some _ -> None  (* compile failed — CEK handles *)
       | None ->
         if !_jit_compiling then None
         else begin
           let fn_name = match l.l_name with Some n -> n | None -> "?" in
           _jit_compiling := true;
           let t0 = Unix.gettimeofday () in
           let compiled = Sx_vm.jit_compile_lambda l env.bindings in
           let dt = Unix.gettimeofday () -. t0 in
           _jit_compiling := false;
           Printf.eprintf "[jit-hook] %s compile %s in %.3fs\n%!"
             fn_name (match compiled with Some _ -> "OK" | None -> "FAIL") dt;
           match compiled with
           | Some cl ->
             l.l_compiled <- Some cl;
             (* Run on VM, fall back to CEK on runtime error *)
             (try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
              with _ -> None)
           | None -> None
         end)
    | _ -> None)


(* ====================================================================== *)
(* Command dispatch                                                        *)
(* ====================================================================== *)

let rec dispatch env cmd =
  match cmd with
  | List [Symbol "ping"] ->
    send_ok_string "ocaml-cek"

  | List [Symbol "load"; String path] ->
    (try
      let exprs = Sx_parser.parse_file path in
      let count = ref 0 in
      List.iter (fun expr ->
        ignore (Sx_ref.eval_expr expr (Env env));
        incr count
      ) exprs;
      send_ok_value (Number (float_of_int !count))
    with
    | Eval_error msg -> send_error msg
    | Sys_error msg -> send_error ("File error: " ^ msg)
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "load-source"; String src] ->
    (try
      let exprs = Sx_parser.parse_all src in
      let count = ref 0 in
      List.iter (fun expr ->
        ignore (Sx_ref.eval_expr expr (Env env));
        incr count
      ) exprs;
      send_ok_value (Number (float_of_int !count))
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "eval-blob"] ->
    let src = read_blob () in
    dispatch env (List [Symbol "eval"; String src])

  | List [Symbol "eval"; String src] ->
    (try
      let exprs = Sx_parser.parse_all src in
      let result = List.fold_left (fun _acc expr ->
        Sx_ref.eval_expr expr (Env env)
      ) Nil exprs in
      (* Use ok-raw with natural list serialization — no (list ...) wrapping.
         This preserves the SX structure for Python to parse back. *)
      let rec raw_serialize = function
        | Nil -> "nil"
        | Bool true -> "true"
        | Bool false -> "false"
        | Number n ->
          if Float.is_integer n then string_of_int (int_of_float n)
          else Printf.sprintf "%g" n
        | String s -> "\"" ^ escape_sx_string s ^ "\""
        | Symbol s -> s
        | Keyword k -> ":" ^ k
        | List items | ListRef { contents = items } ->
          "(" ^ String.concat " " (List.map raw_serialize items) ^ ")"
        | Dict d ->
          let pairs = Hashtbl.fold (fun k v acc ->
            (Printf.sprintf ":%s %s" k (raw_serialize v)) :: acc) d [] in
          "{" ^ String.concat " " pairs ^ "}"
        | Component c -> "~" ^ c.c_name
        | Island i -> "~" ^ i.i_name
        | SxExpr s -> s
        | RawHTML s -> "\"" ^ escape_sx_string s ^ "\""
        | _ -> "nil"
      in
      send_ok_raw (raw_serialize result)
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "vm-reset-fn"; String name] ->
    (* Reset a function's JIT-compiled bytecode, forcing CEK interpretation.
       Used to work around JIT compilation bugs in specific functions. *)
    (match Hashtbl.find_opt env.bindings name with
     | Some (Lambda l) ->
       l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
       Printf.eprintf "[jit] reset %s (forced CEK)\n%!" name;
       send_ok ()
     | _ ->
       Printf.eprintf "[jit] reset %s: not found or not lambda\n%!" name;
       send_ok ())

  | List [Symbol "aser-blob"] ->
    (* Like aser but reads source as a binary blob. *)
    let src = read_blob () in
    dispatch env (List [Symbol "aser"; String src])

  | List [Symbol "aser-slot-blob"] ->
    (* Like aser-slot but reads source as a binary blob. *)
    let src = read_blob () in
    dispatch env (List [Symbol "aser-slot"; String src])

  | List [Symbol "aser"; String src] ->
    (* Evaluate and serialize as SX wire format.
       Calls the SX-defined aser function from adapter-sx.sx.
       aser is loaded into the kernel env via _ensure_components. *)
    (try
      let exprs = Sx_parser.parse_all src in
      let expr = match exprs with
        | [e] -> e
        | [] -> Nil
        | _ -> List (Symbol "<>" :: exprs)
      in
      (* Call (aser <quoted-expr> <env>) *)
      let call = List [Symbol "aser";
                        List [Symbol "quote"; expr];
                        Env env] in
      let result = Sx_ref.eval_expr call (Env env) in
      (* Send raw SX wire format without re-escaping.
         Use (ok-raw ...) so Python knows not to unescape. *)
      (match result with
       | String s | SxExpr s -> send_ok_raw s
       | _ -> send_ok_value result)
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "vm-compile-adapter"] ->
    (* Register JIT hook and pre-compile the compiler itself.
       The compile function running on the VM makes all subsequent
       JIT compilations near-instant (VM speed vs CEK speed). *)
    register_jit_hook env;
    Printf.eprintf "[jit] JIT hook registered (lazy compilation active)\n%!";
    (* Pre-compile the compiler: compile → make-emitter → make-pool
       and other compiler internals so they run on VM from the start *)
    let compiler_fns = ["compile"; "compile-module"; "compile-expr";
      "compile-symbol"; "compile-dict"; "compile-list"; "compile-if";
      "compile-when"; "compile-and"; "compile-or"; "compile-begin";
      "compile-let"; "compile-lambda"; "compile-define"; "compile-set";
      "compile-quote"; "compile-cond"; "compile-case"; "compile-case-clauses";
      "compile-thread"; "compile-thread-step"; "compile-defcomp";
      "compile-defmacro"; "compile-quasiquote"; "compile-call";
      "make-emitter"; "make-pool"; "emit-byte"; "emit-u16"; "emit-i16";
      "pool-add"; "resolve-scope"; "scope-resolve"] in
    let t0 = Unix.gettimeofday () in
    let count = ref 0 in
    List.iter (fun name ->
      match Hashtbl.find_opt env.bindings name with
      | Some (Lambda l) when l.l_compiled = None ->
        l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
        (match Sx_vm.jit_compile_lambda l env.bindings with
         | Some cl ->
           l.l_compiled <- Some cl;
           incr count
         | None -> ())
      | _ -> ()
    ) compiler_fns;
    let dt = Unix.gettimeofday () -. t0 in
    Printf.eprintf "[jit] Pre-compiled %d compiler functions in %.3fs\n%!" !count dt;
    send_ok ()

  | List [Symbol "aser-slot"; String src] ->
    (* Expand ALL components server-side. Uses batch IO mode.
       Calls aser via CEK — the JIT hook compiles it on first call. *)
    (try
      let exprs = Sx_parser.parse_all src in
      let expr = match exprs with
        | [e] -> e
        | [] -> Nil
        | _ -> List (Symbol "<>" :: exprs)
      in
      io_batch_mode := true;
      io_queue := [];
      io_counter := 0;
      let t0 = Unix.gettimeofday () in
      let expand_fn = NativeFn ("expand-components?", fun _args -> Bool true) in
      ignore (env_bind env "expand-components?" expand_fn);
      Printf.eprintf "[aser-slot] starting aser eval...\n%!";
      let result =
        let call = List [Symbol "aser";
                          List [Symbol "quote"; expr];
                          Env env] in
        let r = Sx_ref.eval_expr call (Env env) in
        Printf.eprintf "[aser-slot] aser eval returned\n%!";
        r
      in
      let t1 = Unix.gettimeofday () in
      io_batch_mode := false;
      Hashtbl.remove env.bindings "expand-components?";
      let result_str = match result with
        | String s | SxExpr s -> s
        | _ -> serialize_value result
      in
      let n_batched = List.length !io_queue in
      (* Flush batched IO: send requests, receive responses, replace placeholders *)
      let final = flush_batched_io result_str in
      let t2 = Unix.gettimeofday () in
      Printf.eprintf "[aser-slot] eval=%.1fs io_flush=%.1fs batched=%d result=%d chars\n%!"
        (t1 -. t0) (t2 -. t1) n_batched (String.length final);
      send_ok_raw final
    with
    | Eval_error msg ->
      io_batch_mode := false;
      io_queue := [];
      Hashtbl.remove env.bindings "expand-components?";
      send_error msg
    | exn ->
      io_batch_mode := false;
      io_queue := [];
      Hashtbl.remove env.bindings "expand-components?";
      send_error (Printexc.to_string exn))

  | List (Symbol "sx-page-full-blob" :: shell_kwargs) ->
    (* Like sx-page-full but reads page source as a length-prefixed blob
       from the next line(s), avoiding string-escape round-trip issues. *)
    let page_src = read_blob () in
    dispatch env (List (Symbol "sx-page-full" :: String page_src :: shell_kwargs))

  | List (Symbol "sx-page-full" :: String page_src :: shell_kwargs) ->
    (* Full page render: aser-slot body + render-to-html shell in ONE call.
       shell_kwargs are keyword pairs: :title "..." :csrf "..." etc.
       These are passed directly to ~shared:shell/sx-page-shell. *)
    (try
      (* Phase 1: aser-slot the page body *)
      let exprs = Sx_parser.parse_all page_src in
      let expr = match exprs with
        | [e] -> e | [] -> Nil | _ -> List (Symbol "<>" :: exprs)
      in
      io_batch_mode := true;
      io_queue := [];
      io_counter := 0;
      let t0 = Unix.gettimeofday () in
      let expand_fn = NativeFn ("expand-components?", fun _args -> Bool true) in
      ignore (env_bind env "expand-components?" expand_fn);
      let body_result =
          let call = List [Symbol "aser";
                            List [Symbol "quote"; expr];
                            Env env] in
          Sx_ref.eval_expr call (Env env)
      in
      let t1 = Unix.gettimeofday () in
      io_batch_mode := false;
      Hashtbl.remove env.bindings "expand-components?";
      let body_str = match body_result with
        | String s | SxExpr s -> s
        | _ -> serialize_value body_result
      in
      let body_final = flush_batched_io body_str in
      let t2 = Unix.gettimeofday () in
      (* Phase 1b: render the aser'd SX to HTML for isomorphic SSR.
         The aser output is flat (all components expanded, just HTML tags),
         so render-to-html is cheap — no component lookups needed. *)
      let body_html =
        try
          let body_exprs = Sx_parser.parse_all body_final in
          let body_expr = match body_exprs with
            | [e] -> e | [] -> Nil | _ -> List (Symbol "<>" :: body_exprs)
          in
          Sx_render.render_to_html body_expr env
        with e ->
          Printf.eprintf "[ssr] render-to-html failed: %s\n%!" (Printexc.to_string e);
          ""  (* fallback: client renders from SX source. Islands with
                 reactive state may fail SSR — client hydrates them. *)
      in
      let t2b = Unix.gettimeofday () in
      (* Phase 2: render shell with body + all kwargs.
         Resolve symbol references (e.g. __shell-component-defs) to their
         values from the env — these were pre-injected by the bridge. *)
      let resolved_kwargs = List.map (fun v ->
        match v with
        | Symbol s ->
          (try env_get env s
           with _ -> try Sx_primitives.get_primitive s with _ -> v)
        | _ -> v
      ) shell_kwargs in
      let shell_args = Keyword "page-sx" :: String body_final
                     :: Keyword "body-html" :: String body_html
                     :: resolved_kwargs in
      let shell_call = List (Symbol "~shared:shell/sx-page-shell" :: shell_args) in
      let html = Sx_render.render_to_html shell_call env in
      let t3 = Unix.gettimeofday () in
      Printf.eprintf "[sx-page-full] aser=%.3fs io=%.3fs ssr=%.3fs shell=%.3fs total=%.3fs body=%d ssr=%d html=%d\n%!"
        (t1 -. t0) (t2 -. t1) (t2b -. t2) (t3 -. t2b) (t3 -. t0)
        (String.length body_final) (String.length body_html) (String.length html);
      send_ok_string html
    with
    | Eval_error msg ->
      io_batch_mode := false;
      io_queue := [];
      Hashtbl.remove env.bindings "expand-components?";
      send_error msg
    | exn ->
      io_batch_mode := false;
      io_queue := [];
      Hashtbl.remove env.bindings "expand-components?";
      send_error (Printexc.to_string exn))

  | List [Symbol "render"; String src] ->
    (try
      let exprs = Sx_parser.parse_all src in
      let expr = match exprs with
        | [e] -> e
        | [] -> Nil
        | _ -> List (Symbol "do" :: exprs)
      in
      let html = Sx_render.render_to_html expr env in
      send_ok_string html
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "vm-exec"; code_val] ->
    (* Execute a bytecode module on the VM.
       code_val is a dict with {bytecode, pool} from compiler.sx *)
    (try
      let code = Sx_vm.code_from_value code_val in
      let globals = Hashtbl.create 256 in
      Hashtbl.iter (fun k v -> Hashtbl.replace globals k v) env.bindings;
      let result = Sx_vm.execute_module code globals in
      send_ok_value result
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "vm-load-module"; code_val] ->
    (* Execute a compiled module on the VM. The module's defines
       are stored in the kernel env, replacing Lambda values with
       NativeFn VM closures. This is how compiled code gets wired
       into the CEK dispatch — the CEK calls NativeFn directly. *)
    (try
      let code = Sx_vm.code_from_value code_val in
      (* VM uses the LIVE kernel env — defines go directly into it *)
      let _result = Sx_vm.execute_module code env.bindings in
      (* Count how many defines the module added *)
      send_ok ()
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "vm-compile"] ->
    (* Compile all named lambdas in env to bytecode.
       Called after all .sx files are loaded. *)
    (try
      if not (Hashtbl.mem env.bindings "compile") then
        send_error "compiler not loaded"
      else begin
        let compile_fn = Hashtbl.find env.bindings "compile" in
        let count = ref 0 in
        let failed = ref 0 in
        let names = Hashtbl.fold (fun k _ acc -> k :: acc) env.bindings [] in
        List.iter (fun name ->
          match Hashtbl.find_opt env.bindings name with
          | Some (Lambda lam) when lam.l_name <> None
              && lam.l_closure.parent = None ->
            (try
              let quoted = List [Symbol "quote"; lam.l_body] in
              let result = Sx_ref.eval_expr (List [compile_fn; quoted]) (Env env) in
              match result with
              | Dict d when Hashtbl.mem d "bytecode" ->
                let code = Sx_vm.code_from_value result in
                (* Live env reference — NOT a snapshot. Functions see
                   current bindings, including later-defined functions. *)
                let live_env = env.bindings in
                (* Original lambda for CEK fallback *)
                let orig_lambda = Lambda lam in
                let fn = NativeFn ("vm:" ^ name, fun args ->
                  try
                    Sx_vm.call_closure
                      { vm_code = code; vm_upvalues = [||]; vm_name = lam.l_name;
                        vm_env_ref = live_env }
                      args live_env
                  with
                  | _ ->
                    (* Any VM error — fall back to CEK *)
                    Sx_ref.eval_expr (List (orig_lambda :: args)) (Env env)) in
                Hashtbl.replace env.bindings name fn;
                incr count
              | _ -> incr failed
            with e ->
              if !failed < 3 then
                Printf.eprintf "[vm] FAIL %s: %s\n  body: %s\n%!"
                  name (Printexc.to_string e)
                  (String.sub (inspect lam.l_body) 0
                    (min 200 (String.length (inspect lam.l_body))));
              incr failed)
          | _ -> ()
        ) names;
        Printf.eprintf "[vm] Compiled %d functions (%d failed)\n%!" !count !failed;
        send_ok_value (Number (float_of_int !count))
      end
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "reset"] ->
    (* Clear all bindings and rebuild env.
       We can't reassign env, so clear and re-populate. *)
    Hashtbl.clear env.bindings;
    let fresh = make_server_env () in
    Hashtbl.iter (fun k v -> Hashtbl.replace env.bindings k v) fresh.bindings;
    send_ok ()

  | _ ->
    send_error ("Unknown command: " ^ inspect cmd)


(* ====================================================================== *)
(* Main loop                                                               *)
(* ====================================================================== *)

(* ====================================================================== *)
(* CLI mode — one-shot render/aser from stdin                              *)
(* ====================================================================== *)

let cli_load_files env files =
  List.iter (fun path ->
    if Sys.file_exists path then begin
      let exprs = Sx_parser.parse_file path in
      List.iter (fun expr ->
        ignore (Sx_ref.eval_expr expr (Env env))
      ) exprs
    end
  ) files

let cli_mode mode =
  let env = make_server_env () in
  (* Load spec + adapter files for aser modes *)
  let base = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
  let web_base = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in
  let spec_files = [
    Filename.concat base "parser.sx";
    Filename.concat base "render.sx";
    Filename.concat web_base "adapter-sx.sx";
  ] in
  (if mode = "aser" || mode = "aser-slot" then
    cli_load_files env spec_files);
  (* Load any files passed via --load *)
  let load_files = ref [] in
  let args = Array.to_list Sys.argv in
  let rec scan = function
    | "--load" :: path :: rest -> load_files := path :: !load_files; scan rest
    | _ :: rest -> scan rest
    | [] -> ()
  in scan args;
  cli_load_files env (List.rev !load_files);
  (* Read SX from stdin *)
  let buf = Buffer.create 4096 in
  (try while true do
    let line = input_line stdin in
    Buffer.add_string buf line;
    Buffer.add_char buf '\n'
  done with End_of_file -> ());
  let src = String.trim (Buffer.contents buf) in
  if src = "" then exit 0;
  (try
    match mode with
    | "render" ->
      let exprs = Sx_parser.parse_all src in
      let expr = match exprs with
        | [e] -> e | [] -> Nil | _ -> List (Symbol "do" :: exprs) in
      let html = Sx_render.render_to_html expr env in
      print_string html; flush stdout
    | "aser" ->
      let exprs = Sx_parser.parse_all src in
      let expr = match exprs with
        | [e] -> e | [] -> Nil | _ -> List (Symbol "<>" :: exprs) in
      let call = List [Symbol "aser";
                        List [Symbol "quote"; expr];
                        Env env] in
      let result = Sx_ref.eval_expr call (Env env) in
      (match result with
       | String s | SxExpr s -> print_string s
       | Dict d when Hashtbl.mem d "__aser_sx" ->
         (match Hashtbl.find d "__aser_sx" with
          | String s | SxExpr s -> print_string s
          | v -> print_string (serialize_value v))
       | _ -> print_string (serialize_value result));
      flush stdout
    | "aser-slot" ->
      ignore (env_bind env "expand-components?" (NativeFn ("expand-components?", fun _args -> Bool true)));
      let exprs = Sx_parser.parse_all src in
      let expr = match exprs with
        | [e] -> e | [] -> Nil | _ -> List (Symbol "<>" :: exprs) in
      let call = List [Symbol "aser";
                        List [Symbol "quote"; expr];
                        Env env] in
      let result = Sx_ref.eval_expr call (Env env) in
      (match result with
       | String s | SxExpr s -> print_string s
       | Dict d when Hashtbl.mem d "__aser_sx" ->
         (match Hashtbl.find d "__aser_sx" with
          | String s | SxExpr s -> print_string s
          | v -> print_string (serialize_value v))
       | _ -> print_string (serialize_value result));
      flush stdout
    | _ ->
      Printf.eprintf "Unknown CLI mode: %s\n" mode; exit 1
  with
  | Eval_error msg ->
    Printf.eprintf "Error: %s\n" msg; exit 1
  | exn ->
    Printf.eprintf "Error: %s\n" (Printexc.to_string exn); exit 1)


let () =
  (* Check for CLI mode flags *)
  let args = Array.to_list Sys.argv in
  if List.mem "--render" args then cli_mode "render"
  else if List.mem "--aser-slot" args then cli_mode "aser-slot"
  else if List.mem "--aser" args then cli_mode "aser"
  else begin
    (* Normal persistent server mode *)
    let env = make_server_env () in
    send "(ready)";
    (* Main command loop *)
    try
      while true do
        match read_line_blocking () with
        | None -> exit 0  (* stdin closed *)
        | Some line ->
          let line = String.trim line in
          if line = "" then ()  (* skip blank lines *)
          else begin
            let exprs = Sx_parser.parse_all line in
            match exprs with
            | [cmd] -> dispatch env cmd
            | _ -> send_error ("Expected single command, got " ^ string_of_int (List.length exprs))
          end
      done
    with
    | End_of_file -> ()
  end