(** 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 ^ "\""
  | SxExpr s -> s
  | Spread pairs ->
    let items = List.map (fun (k, v) ->
      Printf.sprintf ":%s %s" k (serialize_value v)) pairs in
    "(make-spread {" ^ String.concat " " items ^ "})"
  | _ -> "nil"

(** Request epoch — monotonically increasing, set by (epoch N) from Python.
    All responses are tagged with the current epoch so Python can discard
    stale messages from previous requests. Makes pipe desync impossible. *)
let current_epoch = ref 0

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

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

(** Length-prefixed binary send — handles any content without escaping.
    Sends: (ok-len EPOCH 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 %d)\n" !current_epoch 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

(* Scope stacks and cookies — all primitives registered in sx_scope.ml.
   We just reference the shared state for the IO bridge. *)
module Sx_scope = Sx.Sx_scope
let _request_cookies = Sx_scope.request_cookies
let _scope_stacks = Sx_scope.scope_stacks

(** 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

(** Read an io-response from stdin, discarding stale messages from old epochs. *)
let rec read_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
    (* Epoch-tagged: (io-response EPOCH value) *)
    | [List [Symbol "io-response"; Number n; value]]
      when int_of_float n = !current_epoch -> value
    | [List (Symbol "io-response" :: Number n :: values)]
      when int_of_float n = !current_epoch ->
      (match values with [v] -> v | _ -> List values)
    (* Legacy untagged: (io-response value) — accept for backwards compat *)
    | [List [Symbol "io-response"; value]] -> value
    | [List (Symbol "io-response" :: values)] ->
      (match values with [v] -> v | _ -> List values)
    (* Stale epoch or unexpected — discard and retry *)
    | _ ->
      Printf.eprintf "[io] discarding stale message (%d chars, epoch=%d)\n%!"
        (String.length line) !current_epoch;
      read_io_response ()

(** 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 %d \"%s\" %s)" !current_epoch name args_str);
    read_io_response ()
  end

(** Read a batched io-response, discarding stale epoch messages. *)
let read_batched_io_response () =
  let rec loop () =
    match read_line_blocking () with
    | None -> raise (Eval_error "IO batch: stdin closed")
    | Some line ->
      let exprs = Sx_parser.parse_all line in
      match exprs with
      (* Epoch-tagged: (io-response EPOCH value) *)
      | [List [Symbol "io-response"; Number n; String s]]
        when int_of_float n = !current_epoch -> s
      | [List [Symbol "io-response"; Number n; SxExpr s]]
        when int_of_float n = !current_epoch -> s
      | [List [Symbol "io-response"; Number n; v]]
        when int_of_float n = !current_epoch -> serialize_value v
      (* Legacy untagged *)
      | [List [Symbol "io-response"; String s]]
      | [List [Symbol "io-response"; SxExpr s]] -> s
      | [List [Symbol "io-response"; v]] -> serialize_value v
      (* Stale — discard and retry *)
      | _ ->
        Printf.eprintf "[io-batch] discarding stale message (%d chars)\n%!"
          (String.length line);
        loop ()
  in
  loop ()

(** 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, tagged with epoch *)
    List.iter (fun (id, name, args) ->
      let args_str = String.concat " " (List.map serialize_value args) in
      send (Printf.sprintf "(io-request %d %d \"%s\" %s)" !current_epoch id name args_str)
    ) queue;
    send (Printf.sprintf "(io-done %d %d)" !current_epoch (List.length queue));
    (* Read all responses and replace placeholders *)
    let final = ref result_str in
    List.iter (fun (id, _, _) ->
      let value_str = read_batched_io_response () 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
    ) 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 ->
    (* Use cek_call instead of eval_expr to avoid re-evaluating
       already-evaluated args.  eval_expr copies Dict values (signals)
       during evaluation, so mutations in the lambda body would affect
       the copy, not the original. *)
    match args with
    | [fn_val; List call_args; Env _e] ->
      Sx_ref.cek_call fn_val (List call_args)
    | [fn_val; List call_args] ->
      Sx_ref.cek_call fn_val (List call_args)
    | _ -> 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                                                       *)
(* ====================================================================== *)

(* ---- Browser API stubs (no-op for SSR) ---- *)
let setup_browser_stubs env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  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 "promise-delayed" (fun args ->
    match args with _ :: [v] -> 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);
  bind "dom-create-element" (fun _args -> Nil);
  bind "dom-append" (fun _args -> Nil);
  bind "create-text-node" (fun _args -> Nil);
  bind "render-to-dom" (fun _args -> Nil);
  bind "set-render-active!" (fun _args -> Nil);
  bind "render-active?" (fun _args -> Bool true)

(* ---- Scope primitives: bind into env for VM visibility ---- *)
let setup_scope_env env =
  List.iter (fun name ->
    try ignore (env_bind env name (Sx_primitives.get_primitive name))
    with _ -> ()
  ) ["scope-push!"; "scope-pop!"; "scope-peek"; "context";
     "collect!"; "collected"; "clear-collected!";
     "scope-emit!"; "emit!"; "emitted"; "scope-emitted";
     "scope-collected"; "scope-clear-collected!";
     "provide-push!"; "provide-pop!";
     "get-cookie"; "set-cookie"];
  ignore (env_bind env "sx-context" (Sx_primitives.get_primitive "context"))

(* ---- CEK evaluator bridge ---- *)
let setup_evaluator_bridge env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  bind "eval-expr" (fun args ->
    match args with
    | [expr; e] -> Sx_ref.eval_expr expr (Env (Sx_runtime.unwrap_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] ->
      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 ->
    (* Use cek_call instead of eval_expr to avoid re-evaluating
       already-evaluated args.  eval_expr copies Dict values (signals)
       during evaluation, so mutations in the lambda body would affect
       the copy, not the original. *)
    match args with
    | [fn_val; List call_args; Env _e] ->
      Sx_ref.cek_call fn_val (List call_args)
    | [fn_val; List call_args] ->
      Sx_ref.cek_call fn_val (List call_args)
    | _ -> 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)"));
  bind "qq-expand-runtime" (fun args ->
    match args with
    | [template] -> Sx_ref.qq_expand template (Env env)
    | [template; Env e] -> Sx_ref.qq_expand template (Env e)
    | _ -> Nil);
  bind "register-special-form!" (fun args ->
    match args with
    | [String name; handler] ->
      ignore (Sx_ref.register_special_form (String name) handler); Nil
    | _ -> raise (Eval_error "register-special-form!: expected (name handler)"));
  ignore (env_bind env "*custom-special-forms*" Sx_ref.custom_special_forms);
  ignore (Sx_ref.register_special_form (String "<>") (NativeFn ("<>", fun args ->
    List (List.map (fun a -> Sx_ref.eval_expr a (Env env)) args))))

(* ---- Type predicates and introspection ---- *)
let setup_introspection env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  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 | _ -> 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 "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 "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)

(* ---- Type operations, string/number/env helpers ---- *)
(* ---- Core runtime operations (assert, append!, apply, etc.) ---- *)
let setup_core_operations env =
  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"));
  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);
  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 "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 "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] -> Bool (Sx_primitives.is_primitive name ||
        (try (match env_get env name with NativeFn _ -> true | _ -> false) with _ -> false))
    | _ -> Bool false);
  bind "get-primitive" (fun args ->
    match args with
    | [String name] -> (try Sx_primitives.get_primitive name with _ -> try env_get env name with _ -> Nil)
    | _ -> Nil);
  bind "make-continuation" (fun args ->
    match args with [f] -> Continuation ((fun v -> Sx_runtime.sx_call f [v]), None) | _ -> raise (Eval_error "make-continuation: expected 1 arg"))

(* ---- Type constructors and symbol operations ---- *)
let setup_type_constructors env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  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"));
  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"));
  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 "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 "sx-serialize" (fun args -> match args with [v] -> String (inspect v) | _ -> raise (Eval_error "sx-serialize: expected 1 arg"));
  bind "is-html-tag?" (fun args -> match args with [String s] -> Bool (Sx_render.is_html_tag s) | _ -> Bool false);
  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 "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 "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 "parse-number" (fun args -> match args with [String s] -> (try Number (float_of_string s) with _ -> Nil) | _ -> Nil)

(* ---- Character classification (platform primitives for spec/parser.sx) ---- *)
let setup_character_classification env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  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)

(* ---- Env operations (env-get, env-has?, env-bind!, etc.) ---- *)
let setup_env_operations env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  let uw = Sx_runtime.unwrap_env in
  bind "env-get" (fun args -> match args with [e; String k] -> Sx_types.env_get (uw e) k | [e; Keyword k] -> Sx_types.env_get (uw e) k | _ -> raise (Eval_error "env-get: expected env and string"));
  bind "env-has?" (fun args -> match args with [e; String k] -> Bool (Sx_types.env_has (uw e) k) | [e; Keyword k] -> Bool (Sx_types.env_has (uw e) k) | _ -> raise (Eval_error "env-has?: expected env and string"));
  bind "env-bind!" (fun args -> match args with [e; String k; v] -> Sx_types.env_bind (uw e) k v | [e; Keyword k; v] -> Sx_types.env_bind (uw e) k v | _ -> raise (Eval_error "env-bind!: expected env, key, value"));
  bind "env-set!" (fun args -> match args with [e; String k; v] -> Sx_types.env_set (uw e) k v | [e; Keyword k; v] -> Sx_types.env_set (uw e) k v | _ -> raise (Eval_error "env-set!: expected env, key, value"));
  bind "env-extend" (fun args -> match args with [e] -> Env (Sx_types.env_extend (uw e)) | _ -> raise (Eval_error "env-extend: expected env"));
  bind "env-merge" (fun args -> match args with [a; b] -> Sx_runtime.env_merge a b | _ -> raise (Eval_error "env-merge: expected 2 envs"))

(* ---- Strict mode (gradual type system support) ---- *)
let setup_strict_mode env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  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)

(* ---- IO helpers routed to Python bridge ---- *)
let setup_io_bridges env =
  let bind name fn = ignore (env_bind env name (NativeFn (name, fn))) in
  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)

(* ---- HTML tag functions (div, span, h1, ...) ---- *)
let setup_html_tags env =
  List.iter (fun tag ->
    ignore (env_bind env tag
      (NativeFn ("html:" ^ tag, fun args -> List (Symbol tag :: args))))
  ) Sx_render.html_tags


(* ====================================================================== *)
(* Compose environment                                                      *)
(* ====================================================================== *)

let make_server_env () =
  let env = make_env () in
  Sx_render.setup_render_env env;
  setup_browser_stubs env;
  setup_scope_env env;
  setup_evaluator_bridge env;
  setup_introspection env;
  setup_core_operations env;
  setup_type_constructors env;
  setup_character_classification env;
  setup_env_operations env;
  setup_strict_mode env;
  setup_io_bridges env;
  setup_html_tags env;
  setup_io_env env;
  (* Initialize trampoline ref so HO primitives (map, filter, etc.)
     can call SX lambdas. Must be done here because Sx_ref is only
     available at the binary level, not in the library. *)
  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


(* ====================================================================== *)
(* SX render-to-html — calls adapter-html.sx via CEK                       *)
(* ====================================================================== *)

(** Render an SX expression to HTML using the SX adapter (adapter-html.sx).
    Falls back to Sx_render.render_to_html if the SX adapter isn't loaded. *)
let sx_render_to_html expr env =
  if env_has env "render-to-html" then
    let fn = env_get env "render-to-html" in
    let result = Sx_ref.cek_call fn (List [expr; Env env]) in
    match result with String s -> s | _ -> Sx_runtime.value_to_str result
  else
    Sx_render.render_to_html expr 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 *)

(* JIT compilation is lazy-only: every named lambda gets one compile
   attempt on first call. Failures are sentineled (never retried). *)

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.
            Mark as failed so we don't retry on every call. *)
         (try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
          with e ->
            let fn_name = match l.l_name with Some n -> n | None -> "?" in
            Printf.eprintf "[jit] DISABLED %s — %s\n%!" fn_name (Printexc.to_string e);
            l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
            None)
       | Some _ -> None  (* compile failed or disabled — CEK handles *)
       | None ->
         let fn_name = match l.l_name with Some n -> n | None -> "?" in
         if !_jit_compiling then None
         else begin
           _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] %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;
             (try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
              with e ->
                Printf.eprintf "[jit] DISABLED %s — %s\n%!" fn_name (Printexc.to_string e);
                l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
                None)
           | None -> None
         end)
    | _ -> None)


(* ====================================================================== *)
(* Re-assert host-provided extension points after loading .sx files.
   evaluator.sx defines *custom-special-forms* and register-special-form!
   which shadow the native bindings from setup_evaluator_bridge. *)
let rebind_host_extensions env =
  Hashtbl.replace env.bindings "register-special-form!"
    (NativeFn ("register-special-form!", fun args ->
      match args with
      | [String name; handler] ->
        ignore (Sx_ref.register_special_form (String name) handler); Nil
      | _ -> raise (Eval_error "register-special-form!: expected (name handler)")));
  ignore (env_bind env "*custom-special-forms*" Sx_ref.custom_special_forms)

(* 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;
      (* Rebind host extension points after .sx load — evaluator.sx
         defines *custom-special-forms* which shadows the native dict *)
      rebind_host_extensions env;
      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
       | List items | ListRef { contents = items } ->
         (* List of SxExprs from map/filter — join them as a fragment *)
         let parts = List.filter_map (fun v -> match v with
           | SxExpr s -> Some s
           | String s -> Some ("\"" ^ escape_sx_string s ^ "\"")
           | Nil -> None
           | v -> Some (serialize_value v)) items in
         if parts = [] then send_ok_raw ""
         else send_ok_raw (String.concat " " parts)
       | _ -> send_ok_value result)
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "vm-compile-adapter"] ->
    (* Register lazy JIT hook — all named lambdas compile on first call.
       Pre-compile compiler internals so subsequent JIT compilations
       run at VM speed, not CEK speed. *)
    register_jit_hook env;
    let t0 = Unix.gettimeofday () in
    let count = ref 0 in
    let compiler_names = [
      "compile"; "compile-module"; "compile-expr"; "compile-symbol";
      "compile-dict"; "compile-list"; "compile-if"; "compile-when";
      "compile-and"; "compile-or"; "compile-begin"; "compile-let";
      "compile-letrec"; "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-qq-expr";
      "compile-qq-list"; "compile-call";
      "make-emitter"; "make-pool"; "make-scope"; "pool-add";
      "scope-define-local"; "scope-resolve";
      "emit-byte"; "emit-u16"; "emit-i16"; "emit-op"; "emit-const";
      "current-offset"; "patch-i16";
    ] 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_names;
    let dt = Unix.gettimeofday () -. t0 in
    Printf.eprintf "[jit] Pre-compiled %d compiler functions in %.3fs (lazy JIT active for all)\n%!" !count dt;
    send_ok ()

  | List [Symbol "set-request-cookies"; Dict cookies] ->
    (* Set request cookies for get-cookie primitive.
       Called by Python bridge before each page render. *)
    Hashtbl.clear _request_cookies;
    Hashtbl.iter (fun k v ->
      match v with String s -> Hashtbl.replace _request_cookies k s | _ -> ()
    ) cookies;
    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_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_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_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; vm_closure_env = None }
                      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;
  (* Rebind after load in case .sx files shadowed host extension points *)
  rebind_host_extensions env

let cli_mode mode =
  let env = make_server_env () in
  (* Load spec + adapter files for aser modes *)
  let spec_base = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
  let lib_base = try Sys.getenv "SX_LIB_DIR" with Not_found -> "lib" in
  let web_base = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in
  let render_files = [
    Filename.concat spec_base "parser.sx";
    Filename.concat spec_base "render.sx";
    Filename.concat web_base "adapter-html.sx";
    Filename.concat web_base "adapter-sx.sx";
    Filename.concat web_base "web-forms.sx";
  ] in
  (* Load spec + adapter files for rendering CLI modes *)
  (if mode = "aser" || mode = "aser-slot" || mode = "render" then
    cli_load_files env render_files);
  ignore lib_base; (* available for --load paths *)
  (* 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_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 test_mode () =
  let env = make_server_env () in
  (* Load spec + lib + adapter stack *)
  let spec_base = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
  let lib_base = try Sys.getenv "SX_LIB_DIR" with Not_found -> "lib" in
  let web_base = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in
  let files = [
    Filename.concat spec_base "parser.sx";
    Filename.concat spec_base "render.sx";
    Filename.concat lib_base "compiler.sx";
    Filename.concat web_base "signals.sx";
    Filename.concat web_base "adapter-html.sx";
    Filename.concat web_base "adapter-sx.sx";
    Filename.concat web_base "web-forms.sx";
  ] in
  cli_load_files env files;
  (* Register JIT *)
  register_jit_hook env;
  (* Load any --load files *)
  let load_files = ref [] in
  let eval_exprs = ref [] in
  let args = Array.to_list Sys.argv in
  let rec scan = function
    | "--load" :: path :: rest -> load_files := path :: !load_files; scan rest
    | "--eval" :: expr :: rest -> eval_exprs := expr :: !eval_exprs; scan rest
    | _ :: rest -> scan rest
    | [] -> ()
  in scan args;
  cli_load_files env (List.rev !load_files);
  if !eval_exprs <> [] then
    List.iter (fun src ->
      try
        let exprs = Sx_parser.parse_all src in
        let result = List.fold_left (fun _ e ->
          Sx_ref.eval_expr e (Env env)) Nil exprs in
        Printf.printf "%s\n%!" (serialize_value result)
      with
      | Eval_error msg -> Printf.eprintf "Error: %s\n%!" msg; exit 1
      | exn -> Printf.eprintf "Error: %s\n%!" (Printexc.to_string exn); exit 1
    ) (List.rev !eval_exprs)
  else begin
    (* Read 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 begin
      try
        let exprs = Sx_parser.parse_all src in
        let result = List.fold_left (fun _ e ->
          Sx_ref.eval_expr e (Env env)) Nil exprs in
        Printf.printf "%s\n%!" (serialize_value result)
      with
      | Eval_error msg -> Printf.eprintf "Error: %s\n%!" msg; exit 1
      | exn -> Printf.eprintf "Error: %s\n%!" (Printexc.to_string exn); exit 1
    end
  end

let () =
  (* Check for CLI mode flags *)
  let args = Array.to_list Sys.argv in
  if List.mem "--test" args then test_mode ()
  else 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 *)
          (* Discard stale io-responses from previous requests. *)
          else if String.length line > 14
               && String.sub line 0 14 = "(io-response " then
            Printf.eprintf "[sx-server] discarding stale io-response (%d chars)\n%!"
              (String.length line)
          else begin
            let exprs = Sx_parser.parse_all line in
            match exprs with
            (* Epoch marker: (epoch N) — set current epoch, read next command *)
            | [List [Symbol "epoch"; Number n]] ->
              current_epoch := int_of_float n
            | [cmd] -> dispatch env cmd
            | _ -> send_error ("Expected single command, got " ^ string_of_int (List.length exprs))
          end
      done
    with
    | End_of_file -> ()
  end