(** SX coroutine subprocess server.

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

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

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

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

open Sx_types


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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

  (* HTML renderer *)
  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-sx-expr" (fun args ->
    match args with
    | [String s] -> SxExpr s
    | _ -> raise (Eval_error "make-sx-expr: expected string"));

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  bind "component-params" (fun args ->
    match args with
    | [Component c] -> List (List.map (fun s -> String s) c.c_params)
    | _ -> 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;

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

  env


(* ====================================================================== *)
(* VM adapter — compiled aser functions in isolated globals                 *)
(* ====================================================================== *)

(** Compiled adapter globals — separate from kernel env.
    Contains compiled aser functions + reads from kernel env for
    components, helpers, and other runtime bindings. *)
let vm_adapter_globals : (string, value) Hashtbl.t option ref = ref None

(** Compile adapter-sx.sx and store in vm_adapter_globals.
    Called from vm-compile-adapter command. *)
let compile_adapter env =
  if not (Hashtbl.mem env.bindings "compile") then
    raise (Eval_error "compiler not loaded")
  else begin
    let compile_fn = Hashtbl.find env.bindings "compile" in
    (* Find and parse adapter-sx.sx *)
    let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found ->
      try Filename.concat (Sys.getenv "SX_SPEC_DIR") "../web"
      with Not_found -> "web" in
    let adapter_path = Filename.concat web_dir "adapter-sx.sx" in
    if not (Sys.file_exists adapter_path) then
      raise (Eval_error ("adapter-sx.sx not found: " ^ adapter_path));
    let exprs = Sx_parser.parse_file adapter_path in
    (* Compile each define's body *)
    let globals = Hashtbl.create 64 in
    (* Seed with kernel env for component/helper lookups *)
    Hashtbl.iter (fun k v -> Hashtbl.replace globals k v) env.bindings;
    let compiled = ref 0 in
    List.iter (fun expr ->
      match expr with
      | List (Symbol "define" :: Symbol name :: rest) ->
        (* Find the body — skip :effects annotations *)
        let rec find_body = function
          | Keyword _ :: _ :: rest -> find_body rest
          | body :: _ -> body
          | [] -> Nil
        in
        let body = find_body rest in
        (try
          let quoted = List [Symbol "quote"; 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 outer_code = Sx_vm.code_from_value result in
            Printf.eprintf "[vm] %s: outer bc=%d consts=%d inner_type=%s\n%!"
              name (Array.length outer_code.Sx_vm.bytecode)
              (Array.length outer_code.Sx_vm.constants)
              (if Array.length outer_code.Sx_vm.constants > 0 then
                type_of outer_code.Sx_vm.constants.(0) else "empty");
            let bc = outer_code.Sx_vm.bytecode in
            if Array.length bc >= 4 && bc.(0) = 51 then begin
              (* The compiled define body is (fn ...) which compiles to
                 OP_CLOSURE + [upvalue descriptors] + OP_RETURN.
                 Extract the inner code object from constants[idx]. *)
              let idx = bc.(1) lor (bc.(2) lsl 8) in
              let code =
                if idx < Array.length outer_code.Sx_vm.constants then begin
                  let inner_val = outer_code.Sx_vm.constants.(idx) in
                  try Sx_vm.code_from_value inner_val
                  with e ->
                    Printf.eprintf "[vm] inner code_from_value failed for %s: %s\n%!"
                      name (Printexc.to_string e);
                    raise e
                end else outer_code
              in
              let cl = { Sx_vm.code; upvalues = [||]; name = Some name;
                          env_ref = globals } in
              Hashtbl.replace globals name
                (NativeFn ("vm:" ^ name, fun args ->
                  Sx_vm.call_closure cl args globals));
              incr compiled
            end else begin
              (* Not a lambda — constant expression (e.g. (list ...)).
                 Execute once and store the resulting value directly. *)
              let value = Sx_vm.execute_module outer_code globals in
              Hashtbl.replace globals name value;
              Printf.eprintf "[vm] %s: constant (type=%s)\n%!" name (type_of value);
              incr compiled
            end
          | _ -> ()  (* non-dict result — skip *)
        with e ->
          Printf.eprintf "[vm] FAIL adapter %s: %s\n%!" name (Printexc.to_string e))

      | _ ->
        (* Non-define expression — evaluate on CEK to set up constants *)
        (try ignore (Sx_ref.eval_expr expr (Env env)) with _ -> ())
    ) exprs;
    vm_adapter_globals := Some globals;
    Printf.eprintf "[vm] Compiled adapter: %d functions\n%!" !compiled
  end


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

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

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

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

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

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

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

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

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

  | List [Symbol "vm-compile-adapter"] ->
    (* Compile adapter-sx.sx to VM bytecode with isolated globals *)
    (try
      compile_adapter env;
      send_ok ()
    with
    | Eval_error msg -> send_error msg
    | exn -> send_error (Printexc.to_string exn))

  | List [Symbol "aser-slot"; String src] ->
    (* Expand ALL components server-side. Uses batch IO mode.
       Routes through VM if adapter is compiled, else CEK. *)
    (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);
      let result = match !vm_adapter_globals with
        | Some globals ->
          Hashtbl.replace globals "expand-components?" expand_fn;
          let aser_fn = try Hashtbl.find globals "aser"
            with Not_found -> raise (Eval_error "VM: aser not compiled") in
          let r = match aser_fn with
            | NativeFn (_, fn) -> fn [expr; Env env]
            | _ -> raise (Eval_error "VM: aser not a function")
          in
          Hashtbl.remove globals "expand-components?";
          r
        | None ->
          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 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 = match !vm_adapter_globals with
        | Some globals ->
          Hashtbl.replace globals "expand-components?" expand_fn;
          let aser_fn = try Hashtbl.find globals "aser"
            with Not_found -> raise (Eval_error "VM: aser not compiled") in
          let r = match aser_fn with
            | NativeFn (_, fn) -> fn [expr; Env env]
            | _ -> raise (Eval_error "VM: aser not a function")
          in
          Hashtbl.remove globals "expand-components?";
          r
        | None ->
          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 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 :: resolved_kwargs in
      let shell_call = List (Symbol "~shared:shell/sx-page-shell" :: shell_args) in
      let html = Sx_render.render_to_html shell_call env in
      let t3 = Unix.gettimeofday () in
      Printf.eprintf "[sx-page-full] aser=%.3fs io=%.3fs shell=%.3fs total=%.3fs body=%d html=%d\n%!"
        (t1 -. t0) (t2 -. t1) (t3 -. t2) (t3 -. t0)
        (String.length body_final) (String.length html);
      send_ok_string html
    with
    | Eval_error msg ->
      io_batch_mode := false;
      io_queue := [];
      Hashtbl.remove env.bindings "expand-components?";
      send_error msg
    | exn ->
      io_batch_mode := false;
      io_queue := [];
      Hashtbl.remove env.bindings "expand-components?";
      send_error (Printexc.to_string exn))

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

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

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

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

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

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


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

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

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

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


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