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

(** SX coroutine subprocess server.

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

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

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

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

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

(** Send an io-request and block until io-response arrives. *)
let io_request name args =
  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))

(** 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]
    | _ -> raise (Eval_error "request-arg: expected 1 arg"));

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


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

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

  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
      send_ok_value result
    with
    | Eval_error msg -> send_error msg
    | exn -> 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 "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                                                               *)
(* ====================================================================== *)

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