(** 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_ok_string s = send (Printf.sprintf "(ok \"%s\")" (escape_sx_string s))
let send_error msg = send (Printf.sprintf "(error \"%s\")" (escape_sx_string msg))


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

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

(** 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;
    (* Placeholder starts with ( so aser inlines it as pre-serialized SX *)
    String (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?)"));

  (* 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 *)
  Sx_render.setup_render_env env;

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

  (* Scope stack — platform primitives for render-time dynamic scope.
     Used by aser for spread/provide/emit patterns. *)
  let scope_stacks : (string, value list) Hashtbl.t = Hashtbl.create 8 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);

  (* 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] -> v  (* CEK never produces thunks *)
    | _ -> 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 "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 -> Bool false);
  bind "spread-attrs" (fun _args -> Dict (Hashtbl.create 0));
  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)
    | [Number n] -> Number (Float.round n)
    | _ -> 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-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)
    | _ -> Nil);

  bind "component-body" (fun args ->
    match args with
    | [Component c] -> c.c_body
    | _ -> Nil);

  bind "component-has-children" (fun args ->
    match args with
    | [Component c] -> Bool c.c_has_children
    | _ -> Bool false);

  bind "component-affinity" (fun args ->
    match args with
    | [Component c] -> String c.c_affinity
    | _ -> 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;

  env


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

let 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"; 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 (Printf.sprintf "(ok-raw %s)" (raw_serialize result))
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | 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 (Printf.sprintf "(ok-raw %s)" s)
       | _ -> send_ok_value result)
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "aser-slot"; String src] ->
    (* Like aser but expands ALL components server-side, not just
       server-affinity ones.  Uses batch IO mode: batchable helper
       calls (highlight etc.) return placeholders during evaluation,
       then all IO is flushed concurrently after the aser completes. *)
    (try
      ignore (env_bind env "expand-components?" (NativeFn ("expand-components?", fun _args -> Bool true)));
      (* Enable batch IO mode *)
      io_batch_mode := true; Sx_ref._cek_steps := 0;
      io_queue := [];
      io_counter := 0;
      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 t0 = Unix.gettimeofday () in
      let result = Sx_ref.eval_expr call (Env env) 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 cek_steps=%d\n%!"
        (t1 -. t0) (t2 -. t1) n_batched (String.length final) !Sx_ref._cek_steps;
      send (Printf.sprintf "(ok-raw %s)" 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 "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-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 ->
            (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
                let globals_snapshot = Hashtbl.copy env.bindings in
                Hashtbl.iter (fun k v ->
                  Hashtbl.replace globals_snapshot k v) lam.l_closure.bindings;
                (* VM closure with CEK fallback on error *)
                let orig_lambda = Lambda lam in
                let fn = NativeFn ("vm:" ^ name, fun args ->
                  try
                    Sx_vm.execute_closure
                      { Sx_vm.code; name = lam.l_name } args globals_snapshot
                  with _ ->
                    (* Fall back to CEK machine *)
                    Sx_ref.cek_call orig_lambda (List args)) in
                Hashtbl.replace env.bindings name fn;
                incr count
              | _ -> incr failed
            with e ->
              if !failed < 5 then
                Printf.eprintf "[vm] FAIL %s: %s\n%!" name (Printexc.to_string e);
              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
       | _ -> 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
       | _ -> 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