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

(** MCP server for SX tree tools — structural reading and editing of .sx files.

    Stdio JSON-RPC transport following the MCP specification.
    Loads tree-tools.sx into the SX evaluator and exposes comprehension
    and editing functions as MCP tools. *)

open Sx_types

(* ------------------------------------------------------------------ *)
(* SX evaluator setup — minimal env for parser + tree-tools            *)
(* ------------------------------------------------------------------ *)

let env = ref (make_env ())

let load_sx_file e path =
  let src = In_channel.with_open_text path In_channel.input_all in
  let exprs = Sx_parser.parse_all src in
  List.iter (fun expr ->
    try ignore (Sx_ref.cek_call
      (NativeFn ("eval", fun args ->
        match args with
        | [ex] -> Sx_ref.eval_expr ex (Env e)
        | _ -> Nil))
      (List [expr]))
    with _ ->
      (* Fallback: direct eval *)
      ignore (Sx_ref.eval_expr expr (Env e))
  ) exprs

let setup_env () =
  let e = make_env () in
  (* Primitives are auto-registered at module init *)
  (* Trampoline ref for HO primitives *)
  Sx_primitives._sx_trampoline_fn := (fun v ->
    match v with
    | Thunk (body, closure_env) -> Sx_ref.eval_expr body (Env closure_env)
    | other -> other);
  (* Character classification for parser *)
  let bind name fn = ignore (env_bind e name (NativeFn (name, fn))) in
  bind "is-whitespace?" (fun args -> match args with
    | [String s] when String.length s = 1 ->
      let c = s.[0] in Bool (c = ' ' || c = '\t' || c = '\n' || c = '\r')
    | _ -> Bool false);
  bind "is-digit?" (fun args -> match args with
    | [String s] when String.length s = 1 ->
      Bool (s.[0] >= '0' && s.[0] <= '9')
    | _ -> Bool false);
  bind "is-alpha?" (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'))
    | _ -> Bool false);
  bind "char-code" (fun args -> match args with
    | [String s] when String.length s > 0 -> Number (float_of_int (Char.code s.[0]))
    | _ -> Number 0.0);
  bind "code-char" (fun args -> match args with
    | [Number n] -> String (String.make 1 (Char.chr (int_of_float n)))
    | _ -> String "");
  bind "parse-number" (fun args -> match args with
    | [String s] -> (try Number (float_of_string s) with _ -> Nil)
    | _ -> Nil);
  bind "identical?" (fun args -> match args with
    | [a; b] -> Bool (a == b)
    | _ -> Bool false);
  (* Character classification for SX 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 = '%' ||
            Char.code c > 127)
    | _ -> 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 >= '0' && c <= '9') ||
            c = '_' || c = '~' || c = '*' || c = '+' || c = '-' ||
            c = '>' || c = '<' || c = '=' || c = '/' || c = '!' ||
            c = '?' || c = '&' || c = '.' || c = ':' || c = '#' ||
            c = ',' || c = '@' || c = '^' || c = '%' ||
            Char.code c > 127)
    | _ -> Bool false);
  bind "make-keyword" (fun args -> match args with
    | [String s] -> Keyword s | _ -> 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"
        | '\t' -> Buffer.add_string buf "\\t"
        | '\r' -> Buffer.add_string buf "\\r"
        | c -> Buffer.add_char buf c) s;
      String (Buffer.contents buf)
    | _ -> String "");
  bind "sx-expr-source" (fun args -> match args with
    | [SxExpr s] -> String s | _ -> String "");
  (* Runtime functions needed by tree-tools *)
  bind "symbol-name" (fun args -> match args with
    | [Symbol s] -> String s | _ -> String "");
  bind "keyword-name" (fun args -> match args with
    | [Keyword k] -> String k | _ -> String "");
  bind "make-symbol" (fun args -> match args with
    | [String s] -> Symbol s | _ -> Nil);
  bind "type-of" (fun args -> match args with
    | [v] -> String (type_of v) | _ -> String "nil");
  bind "list?" (fun args -> match args with
    | [List _ | ListRef _] -> Bool true | _ -> Bool false);
  bind "nil?" (fun args -> match args with
    | [v] -> Bool (is_nil v) | _ -> Bool true);
  bind "string?" (fun args -> match args with
    | [String _] -> Bool true | _ -> Bool false);
  bind "number?" (fun args -> match args with
    | [Number _] -> Bool true | _ -> Bool false);
  bind "callable?" (fun args -> match args with
    | [NativeFn _ | Lambda _ | Component _ | Island _] -> Bool true | _ -> Bool false);
  bind "empty?" (fun args -> match args with
    | [List []] | [ListRef { contents = [] }] -> Bool true
    | [Nil] -> Bool true | _ -> Bool false);
  bind "contains?" (fun args -> match args with
    | [String s; String sub] ->
      let rec find i =
        if i > String.length s - String.length sub then false
        else if String.sub s i (String.length sub) = sub then true
        else find (i + 1)
      in Bool (String.length sub = 0 || find 0)
    | [List l; v] | [ListRef { contents = l }; v] ->
      Bool (List.exists (fun x -> x = v) l)
    | _ -> Bool false);
  bind "starts-with?" (fun args -> match args with
    | [String s; String prefix] ->
      Bool (String.length s >= String.length prefix &&
            String.sub s 0 (String.length prefix) = prefix)
    | _ -> Bool false);
  bind "append!" (fun args -> match args with
    | [ListRef r; v] -> r := !r @ [v]; v
    | _ -> Nil);
  bind "map-indexed" (fun args -> match args with
    | [f; List l] | [f; ListRef { contents = l }] ->
      List (List.mapi (fun i x -> Sx_ref.cek_call f (List [Number (float_of_int i); x])) l)
    | _ -> List []);
  bind "trim" (fun args -> match args with
    | [String s] -> String (String.trim s) | _ -> String "");
  bind "split" (fun args -> match args with
    | [String s; String d] ->
      List (List.map (fun p -> String p) (String.split_on_char d.[0] s))
    | _ -> List []);
  (* sx-parse — use the native OCaml parser for bootstrapping *)
  bind "sx-parse" (fun args -> match args with
    | [String s] -> List (Sx_parser.parse_all s)
    | _ -> List []);
  bind "sx-serialize" (fun args -> match args with
    | [v] -> String (Sx_runtime.value_to_str v)
    | _ -> String "");
  (* Load parser.sx for the SX-level sx-parse/sx-serialize *)
  let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
  let lib_dir = try Sys.getenv "SX_LIB_DIR" with Not_found -> "lib" in
  (try load_sx_file e (Filename.concat spec_dir "parser.sx")
   with exn -> Printf.eprintf "[mcp] Warning: parser.sx load failed: %s\n%!" (Printexc.to_string exn));
  (* Load tree-tools *)
  (try load_sx_file e (Filename.concat lib_dir "tree-tools.sx")
   with exn -> Printf.eprintf "[mcp] Error: tree-tools.sx load failed: %s\n%!" (Printexc.to_string exn); exit 1);
  Printf.eprintf "[mcp] SX tree-tools loaded\n%!";
  env := e

(* ------------------------------------------------------------------ *)
(* Call SX tree-tools functions                                        *)
(* ------------------------------------------------------------------ *)

let call_sx fn_name args =
  let e = !env in
  let fn = env_get e fn_name in
  Sx_ref.cek_call fn (List args)

let parse_file path =
  let src = In_channel.with_open_text path In_channel.input_all in
  let exprs = Sx_parser.parse_all src in
  List exprs

let parse_path_str s =
  (* Parse SX path string: "(0 3 2)" or "0 3 2" → SX list of numbers *)
  let exprs = Sx_parser.parse_all s in
  match exprs with
  | [List items] ->
    (* (0 3 2) → list of numbers *)
    List (List.map (fun x -> match x with Number _ -> x | _ -> Number 0.0) items)
  | _ ->
    (* Bare numbers: "0 3 2" → parsed as separate exprs *)
    List (List.map (fun x -> match x with Number _ -> x | _ -> Number 0.0) exprs)

let _json_to_path j =
  let open Yojson.Safe.Util in
  parse_path_str (to_string j)

(* Resolve path: if it contains ">", use resolve-named-path; else parse as index path *)
let resolve_path tree path_str =
  if String.contains path_str '>' then
    call_sx "resolve-named-path" [tree; String path_str]
  else
    parse_path_str path_str

let value_to_string v =
  match v with
  | String s -> s
  | _ -> Sx_runtime.value_to_str v

let text_result s =
  `Assoc [("content", `List [`Assoc [
    ("type", `String "text");
    ("text", `String s)
  ]])]

let error_result s =
  `Assoc [("content", `List [`Assoc [
    ("type", `String "text");
    ("text", `String s)
  ]]);
  ("isError", `Bool true)]

(* ------------------------------------------------------------------ *)
(* Recursive .sx file discovery                                        *)
(* ------------------------------------------------------------------ *)

let glob_sx_files dir =
  let results = ref [] in
  let rec walk path =
    if Sys.is_directory path then
      let entries = Sys.readdir path in
      Array.iter (fun e -> walk (Filename.concat path e)) entries
    else if Filename.check_suffix path ".sx" then
      results := path :: !results
  in
  (try walk dir with Sys_error _ -> ());
  List.sort String.compare !results

let relative_path ~base path =
  let blen = String.length base in
  if String.length path > blen && String.sub path 0 blen = base then
    let rest = String.sub path (blen + 1) (String.length path - blen - 1) in
    rest
  else path

(* ------------------------------------------------------------------ *)
(* Pretty printer                                                      *)
(* ------------------------------------------------------------------ *)

let pp_atom = Sx_types.inspect

(* Estimate single-line width of a value *)
let rec est_width = function
  | Nil -> 3 | Bool true -> 4 | Bool false -> 5
  | Number n -> String.length (if Float.is_integer n then string_of_int (int_of_float n) else Printf.sprintf "%g" n)
  | String s -> String.length s + 2
  | Symbol s -> String.length s
  | Keyword k -> String.length k + 1
  | SxExpr s -> String.length s + 2
  | List items | ListRef { contents = items } ->
    2 + List.fold_left (fun acc x -> acc + est_width x + 1) 0 items
  | _ -> 10

let pretty_print_value ?(max_width=80) v =
  let buf = Buffer.create 4096 in
  let rec pp indent v =
    match v with
    | List items | ListRef { contents = items } when items <> [] ->
      if est_width v <= max_width - indent then
        (* Fits on one line *)
        Buffer.add_string buf (pp_atom v)
      else begin
        (* Multi-line *)
        Buffer.add_char buf '(';
        let head = List.hd items in
        Buffer.add_string buf (pp_atom head);
        let child_indent = indent + 2 in
        let rest = List.tl items in
        (* Special case: keyword args stay on same line as their value *)
        let rec emit = function
          | [] -> ()
          | Keyword k :: v :: rest ->
            Buffer.add_char buf '\n';
            Buffer.add_string buf (String.make child_indent ' ');
            Buffer.add_char buf ':';
            Buffer.add_string buf k;
            Buffer.add_char buf ' ';
            pp child_indent v;
            emit rest
          | item :: rest ->
            Buffer.add_char buf '\n';
            Buffer.add_string buf (String.make child_indent ' ');
            pp child_indent item;
            emit rest
        in
        emit rest;
        Buffer.add_char buf ')'
      end
    | _ -> Buffer.add_string buf (pp_atom v)
  in
  pp 0 v;
  Buffer.contents buf

let pretty_print_file exprs =
  String.concat "\n\n" (List.map pretty_print_value exprs) ^ "\n"

(* ------------------------------------------------------------------ *)
(* Tool handlers                                                       *)
(* ------------------------------------------------------------------ *)

let rec handle_tool name args =
  let open Yojson.Safe.Util in
  match name with
  | "sx_read_tree" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let focus = args |> member "focus" |> to_string_option in
    let max_depth = args |> member "max_depth" |> to_int_option in
    let max_lines = args |> member "max_lines" |> to_int_option in
    let offset = args |> member "offset" |> to_int_option |> Option.value ~default:0 in
    (match focus with
     | Some pattern ->
       (* Focus mode: expand matching subtrees, collapse rest *)
       text_result (value_to_string (call_sx "annotate-focused" [tree; String pattern]))
     | None ->
       match max_lines with
       | Some limit ->
         (* Paginated mode *)
         text_result (value_to_string (call_sx "annotate-paginated"
           [tree; Number (float_of_int offset); Number (float_of_int limit)]))
       | None ->
         match max_depth with
         | Some depth ->
           (* Depth-limited mode *)
           text_result (value_to_string (call_sx "summarise" [tree; Number (float_of_int depth)]))
         | None ->
           (* Auto mode: full tree if small, summarise if large *)
           let full = value_to_string (call_sx "annotate-tree" [tree]) in
           let line_count = 1 + String.fold_left (fun n c -> if c = '\n' then n + 1 else n) 0 full in
           if line_count <= 200 then text_result full
           else
             let summary = value_to_string (call_sx "summarise" [tree; Number 2.0]) in
             text_result (Printf.sprintf ";; File has %d lines — showing depth-2 summary. Use max_depth, max_lines, or focus to control output.\n%s" line_count summary))

  | "sx_summarise" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    let depth = args |> member "depth" |> to_int in
    text_result (value_to_string (call_sx "summarise" [tree; Number (float_of_int depth)]))

  | "sx_read_subtree" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    text_result (value_to_string (call_sx "read-subtree" [tree; path]))

  | "sx_get_context" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    text_result (value_to_string (call_sx "get-context" [tree; path]))

  | "sx_find_all" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    let pattern = args |> member "pattern" |> to_string in
    let results = call_sx "find-all" [tree; String pattern] in
    let lines = match results with
      | List items | ListRef { contents = items } ->
        List.map (fun item ->
          match item with
          | List [p; s] | ListRef { contents = [p; s] } ->
            value_to_string (call_sx "path-str" [p]) ^ " " ^ value_to_string s
          | _ -> value_to_string item
        ) items
      | _ -> [value_to_string results]
    in
    text_result (String.concat "\n" lines)

  | "sx_get_siblings" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    text_result (value_to_string (call_sx "get-siblings" [tree; path]))

  | "sx_validate" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    text_result (value_to_string (call_sx "validate" [tree]))

  | "sx_replace_node" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    let src = args |> member "new_source" |> to_string in
    write_edit file (call_sx "replace-node" [tree; path; String src])

  | "sx_insert_child" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    let index = args |> member "index" |> to_int in
    let src = args |> member "new_source" |> to_string in
    write_edit file (call_sx "insert-child" [tree; path; Number (float_of_int index); String src])

  | "sx_delete_node" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    write_edit file (call_sx "delete-node" [tree; path])

  | "sx_wrap_node" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let path = resolve_path tree (args |> member "path" |> to_string) in
    let wrapper = args |> member "wrapper" |> to_string in
    write_edit file (call_sx "wrap-node" [tree; path; String wrapper])

  | "sx_format_check" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let warnings = call_sx "lint-file" [tree] in
    (match warnings with
     | List [] | ListRef { contents = [] } -> text_result "OK — no issues found"
     | List items | ListRef { contents = items } ->
       text_result (String.concat "\n" (List.map value_to_string items))
     | _ -> text_result (value_to_string warnings))

  | "sx_macroexpand" ->
    let file = args |> member "file" |> to_string_option in
    let expr_str = args |> member "expr" |> to_string in
    (* Create a fresh env with file definitions loaded *)
    let e = !env in
    (* Optionally load a file's definitions to get its macros *)
    (match file with
     | Some f ->
       (try load_sx_file e f
        with exn -> Printf.eprintf "[mcp] Warning: failed to load %s: %s\n%!" f (Printexc.to_string exn))
     | None -> ());
    let exprs = Sx_parser.parse_all expr_str in
    let result = List.fold_left (fun _acc expr ->
      Sx_ref.eval_expr expr (Env e)
    ) Nil exprs in
    text_result (Sx_types.inspect result)

  | "sx_build" ->
    let target = args |> member "target" |> to_string_option |> Option.value ~default:"js" in
    let full = args |> member "full" |> to_bool_option |> Option.value ~default:false in
    let project_dir = try Sys.getenv "SX_PROJECT_DIR" with Not_found ->
      let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
      Filename.dirname spec_dir
    in
    let cmd = match target with
      | "ocaml" ->
        Printf.sprintf "cd %s/hosts/ocaml && eval $(opam env 2>/dev/null) && dune build 2>&1" project_dir
      | "js" | _ ->
        let extra = if full then " --extensions continuations --spec-modules types" else "" in
        Printf.sprintf "cd %s && python3 hosts/javascript/cli.py%s --output shared/static/scripts/sx-browser.js 2>&1" project_dir extra
    in
    let ic = Unix.open_process_in cmd in
    let lines = ref [] in
    (try while true do lines := input_line ic :: !lines done with End_of_file -> ());
    let status = Unix.close_process_in ic in
    let output = String.concat "\n" (List.rev !lines) in
    (match status with
     | Unix.WEXITED 0 -> text_result (Printf.sprintf "OK — %s build succeeded\n%s" target (String.trim output))
     | _ -> error_result (Printf.sprintf "%s build failed:\n%s" target output))

  | "sx_test" ->
    let host = args |> member "host" |> to_string_option |> Option.value ~default:"js" in
    let full = args |> member "full" |> to_bool_option |> Option.value ~default:false in
    let project_dir = try Sys.getenv "SX_PROJECT_DIR" with Not_found ->
      (* Walk up from spec dir to find project root *)
      let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
      Filename.dirname spec_dir
    in
    let cmd = match host with
      | "ocaml" ->
        Printf.sprintf "cd %s/hosts/ocaml && eval $(opam env 2>/dev/null) && dune exec bin/run_tests.exe%s 2>&1"
          project_dir (if full then " -- --full" else "")
      | "js" | _ ->
        Printf.sprintf "cd %s && node hosts/javascript/run_tests.js%s 2>&1"
          project_dir (if full then " --full" else "")
    in
    let ic = Unix.open_process_in cmd in
    let lines = ref [] in
    (try while true do lines := input_line ic :: !lines done with End_of_file -> ());
    ignore (Unix.close_process_in ic);
    let all_lines = List.rev !lines in
    (* Extract summary and failures *)
    let fails = List.filter (fun l -> let t = String.trim l in
      String.length t > 5 && String.sub t 0 4 = "FAIL") all_lines in
    let summary = List.find_opt (fun l -> try let _ = Str.search_forward (Str.regexp "Results:") l 0 in true with Not_found -> false) all_lines in
    let result = match summary with
      | Some s ->
        if fails = [] then s
        else s ^ "\n\nFailures:\n" ^ String.concat "\n" fails
      | None ->
        let last_n = List.filteri (fun i _ -> i >= List.length all_lines - 5) all_lines in
        String.concat "\n" last_n
    in
    text_result result

  | "sx_pretty_print" ->
    let file = args |> member "file" |> to_string in
    let exprs = Sx_parser.parse_all (In_channel.with_open_text file In_channel.input_all) in
    let source = pretty_print_file exprs in
    Out_channel.with_open_text file (fun oc -> output_string oc source);
    text_result (Printf.sprintf "OK — reformatted %s (%d bytes, %d forms)" file (String.length source) (List.length exprs))

  | "sx_write_file" ->
    let file = args |> member "file" |> to_string in
    let source = args |> member "source" |> to_string in
    (* Validate by parsing first *)
    (try
      let exprs = Sx_parser.parse_all source in
      if exprs = [] then error_result "Source parsed to empty — nothing to write"
      else begin
        let output = pretty_print_file exprs in
        Out_channel.with_open_text file (fun oc -> output_string oc output);
        text_result (Printf.sprintf "OK — wrote %d bytes (%d top-level forms) to %s" (String.length output) (List.length exprs) file)
      end
    with e -> error_result (Printf.sprintf "Parse error — file not written: %s" (Printexc.to_string e)))

  | "sx_rename_symbol" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let old_name = args |> member "old_name" |> to_string in
    let new_name = args |> member "new_name" |> to_string in
    let new_tree = call_sx "rename-symbol" [tree; String old_name; String new_name] in
    let count = call_sx "count-renames" [tree; String old_name] in
    let count_str = value_to_string count in
    write_edit file (Dict (let d = Hashtbl.create 2 in Hashtbl.replace d "ok" new_tree; d))
    |> (fun result ->
      match result with
      | `Assoc [("content", `List [`Assoc [("type", _); ("text", `String s)]])] when not (String.starts_with ~prefix:"Error" s) ->
        text_result (Printf.sprintf "Renamed %s occurrences of '%s' → '%s' in %s" count_str old_name new_name file)
      | other -> other)

  | "sx_replace_by_pattern" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let pattern = args |> member "pattern" |> to_string in
    let src = args |> member "new_source" |> to_string in
    let all = args |> member "all" |> to_bool_option |> Option.value ~default:false in
    if all then
      write_edit file (call_sx "replace-all-by-pattern" [tree; String pattern; String src])
    else
      write_edit file (call_sx "replace-by-pattern" [tree; String pattern; String src])

  | "sx_insert_near" ->
    let file = args |> member "file" |> to_string in
    let tree = parse_file file in
    let pattern = args |> member "pattern" |> to_string in
    let position = args |> member "position" |> to_string_option |> Option.value ~default:"after" in
    let src = args |> member "new_source" |> to_string in
    write_edit file (call_sx "insert-near-pattern" [tree; String pattern; String position; String src])

  | "sx_rename_across" ->
    let dir = args |> member "dir" |> to_string in
    let old_name = args |> member "old_name" |> to_string in
    let new_name = args |> member "new_name" |> to_string in
    let dry_run = args |> member "dry_run" |> to_bool_option |> Option.value ~default:false in
    let files = glob_sx_files dir in
    let results = List.filter_map (fun path ->
      let rel = relative_path ~base:dir path in
      try
        let tree = parse_file path in
        let count = call_sx "count-renames" [tree; String old_name] in
        match count with
        | Number n when n > 0.0 ->
          if dry_run then
            Some (Printf.sprintf "%s: %d occurrences (dry run)" rel (int_of_float n))
          else begin
            let new_tree = call_sx "rename-symbol" [tree; String old_name; String new_name] in
            let items = match new_tree with
              | List items | ListRef { contents = items } -> items
              | _ -> [new_tree]
            in
            let source = pretty_print_file items in
            Out_channel.with_open_text path (fun oc -> output_string oc source);
            Some (Printf.sprintf "%s: %d occurrences renamed" rel (int_of_float n))
          end
        | _ -> None
      with _ -> None
    ) files in
    if results = [] then text_result (Printf.sprintf "No occurrences of '%s' found" old_name)
    else text_result (String.concat "\n" results)

  | "sx_comp_list" ->
    let dir = args |> member "dir" |> to_string in
    let files = glob_sx_files dir in
    let all_lines = List.concat_map (fun path ->
      let rel = relative_path ~base:dir path in
      try
        let exprs = Sx_parser.parse_all (In_channel.with_open_text path In_channel.input_all) in
        List.filter_map (fun expr ->
          match expr with
          | List (Symbol head :: Symbol name :: rest) | ListRef { contents = Symbol head :: Symbol name :: rest } ->
            (match head with
             | "defcomp" | "defisland" | "defmacro" | "defpage" | "define" ->
               let params = match rest with
                 | List ps :: _ | ListRef { contents = ps } :: _ ->
                   String.concat " " (List.map Sx_runtime.value_to_str ps)
                 | _ -> ""
               in
               Some (Printf.sprintf "%-10s %-40s %-50s %s" head name rel params)
             | _ -> None)
          | _ -> None
        ) exprs
      with _ -> []
    ) files in
    if all_lines = [] then text_result "(no definitions found)"
    else text_result (Printf.sprintf "%-10s %-40s %-50s %s\n%s" "TYPE" "NAME" "FILE" "PARAMS" (String.concat "\n" all_lines))

  | "sx_find_across" ->
    let dir = args |> member "dir" |> to_string in
    let pattern = args |> member "pattern" |> to_string in
    let files = glob_sx_files dir in
    let all_lines = List.concat_map (fun path ->
      let rel = relative_path ~base:dir path in
      try
        let tree = parse_file path in
        let results = call_sx "find-all" [tree; String pattern] in
        (match results with
         | List items | ListRef { contents = items } ->
           List.map (fun item ->
             match item with
             | List [p; s] | ListRef { contents = [p; s] } ->
               rel ^ " " ^ value_to_string (call_sx "path-str" [p]) ^ " " ^ value_to_string s
             | _ -> rel ^ " " ^ value_to_string item
           ) items
         | _ -> [])
      with _ -> []
    ) files in
    if all_lines = [] then text_result "(no matches)"
    else text_result (String.concat "\n" all_lines)

  | "sx_diff" ->
    let file_a = args |> member "file_a" |> to_string in
    let file_b = args |> member "file_b" |> to_string in
    let tree_a = parse_file file_a in
    let tree_b = parse_file file_b in
    text_result (value_to_string (call_sx "tree-diff" [tree_a; tree_b]))

  | "sx_comp_usage" ->
    let dir = args |> member "dir" |> to_string in
    let name = args |> member "name" |> to_string in
    let files = glob_sx_files dir in
    let all_lines = List.concat_map (fun path ->
      let rel = relative_path ~base:dir path in
      try
        let tree = parse_file path in
        let results = call_sx "find-all" [tree; String name] in
        (match results with
         | List items | ListRef { contents = items } ->
           List.map (fun item ->
             match item with
             | List [p; s] | ListRef { contents = [p; s] } ->
               rel ^ " " ^ value_to_string (call_sx "path-str" [p]) ^ " " ^ value_to_string s
             | _ -> rel ^ " " ^ value_to_string item
           ) items
         | _ -> [])
      with _ -> []
    ) files in
    if all_lines = [] then text_result "(no usages found)"
    else text_result (String.concat "\n" all_lines)

  | "sx_eval" ->
    let expr_str = args |> member "expr" |> to_string in
    let exprs = Sx_parser.parse_all expr_str in
    let e = !env in
    let result = List.fold_left (fun _acc expr ->
      Sx_ref.eval_expr expr (Env e)
    ) Nil exprs in
    text_result (Sx_runtime.value_to_str result)

  | _ -> error_result ("Unknown tool: " ^ name)

and write_edit file result =
  match result with
  | Dict d ->
    (match Hashtbl.find_opt d "ok" with
     | Some new_tree ->
       let items = match new_tree with
         | List items | ListRef { contents = items } -> items
         | _ -> [new_tree]
       in
       let source = pretty_print_file items in
       Out_channel.with_open_text file (fun oc -> output_string oc source);
       text_result (Printf.sprintf "OK — wrote %d bytes to %s" (String.length source) file)
     | None ->
       let err = match Hashtbl.find_opt d "error" with
         | Some (String s) -> s | Some v -> value_to_string v | None -> "Unknown error"
       in
       error_result ("Error: " ^ err))
  | _ -> error_result "Unexpected result type"

(* ------------------------------------------------------------------ *)
(* MCP tool definitions                                                *)
(* ------------------------------------------------------------------ *)

let tool name desc props required =
  `Assoc [
    ("name", `String name);
    ("description", `String desc);
    ("inputSchema", `Assoc [
      ("type", `String "object");
      ("required", `List (List.map (fun r -> `String r) required));
      ("properties", `Assoc props)])]

let file_prop = ("file", `Assoc [("type", `String "string"); ("description", `String "Path to .sx file")])
let path_prop = ("path", `Assoc [("type", `String "string"); ("description", `String "SX path, e.g. \"(0 2 1)\"")])
let dir_prop = ("dir", `Assoc [("type", `String "string"); ("description", `String "Directory to scan recursively")])

let tool_definitions = `List [
  tool "sx_read_tree" "Read an .sx file as an annotated tree with path labels. Auto-summarises large files (>200 lines). Use focus to expand only matching subtrees, max_depth for depth limit, or max_lines+offset for pagination."
    [file_prop;
     ("focus", `Assoc [("type", `String "string"); ("description", `String "Pattern — expand matching subtrees, collapse rest")]);
     ("max_depth", `Assoc [("type", `String "integer"); ("description", `String "Depth limit (like summarise)")]);
     ("max_lines", `Assoc [("type", `String "integer"); ("description", `String "Max lines to return (pagination)")]);
     ("offset", `Assoc [("type", `String "integer"); ("description", `String "Line offset for pagination (default 0)")])]
    ["file"];
  tool "sx_summarise" "Folded structural overview of an .sx file. Use to orient before drilling into a region."
    [file_prop; ("depth", `Assoc [("type", `String "integer"); ("description", `String "Max depth (0=heads only)")])] ["file"; "depth"];
  tool "sx_read_subtree" "Expand a specific subtree by path. Use after summarise to drill in."
    [file_prop; path_prop] ["file"; "path"];
  tool "sx_get_context" "Show enclosing chain from root to a target node."
    [file_prop; path_prop] ["file"; "path"];
  tool "sx_find_all" "Search for nodes matching a pattern. Returns paths and summaries."
    [file_prop; ("pattern", `Assoc [("type", `String "string"); ("description", `String "Search pattern")])] ["file"; "pattern"];
  tool "sx_get_siblings" "Show siblings of a node with target marked."
    [file_prop; path_prop] ["file"; "path"];
  tool "sx_validate" "Check structural integrity of an .sx file."
    [file_prop] ["file"];
  tool "sx_replace_node" "Replace node at path with new SX source. Fragment is parsed before file is touched."
    [file_prop; path_prop; ("new_source", `Assoc [("type", `String "string"); ("description", `String "New SX source")])] ["file"; "path"; "new_source"];
  tool "sx_insert_child" "Insert new child at index within a list node."
    [file_prop; path_prop; ("index", `Assoc [("type", `String "integer"); ("description", `String "Insert position")]); ("new_source", `Assoc [("type", `String "string"); ("description", `String "New SX source")])] ["file"; "path"; "index"; "new_source"];
  tool "sx_delete_node" "Remove node at path. Siblings shift to fill gap."
    [file_prop; path_prop] ["file"; "path"];
  tool "sx_wrap_node" "Wrap node in a new form. Use _ as placeholder, e.g. \"(when cond _)\"."
    [file_prop; path_prop; ("wrapper", `Assoc [("type", `String "string"); ("description", `String "Wrapper with _ placeholder")])] ["file"; "path"; "wrapper"];
  tool "sx_eval" "Evaluate an SX expression. Environment has parser + tree-tools + primitives."
    [("expr", `Assoc [("type", `String "string"); ("description", `String "SX expression to evaluate")])] ["expr"];
  tool "sx_find_across" "Search for a pattern across all .sx files under a directory. Returns file paths, tree paths, and summaries."
    [dir_prop; ("pattern", `Assoc [("type", `String "string"); ("description", `String "Search pattern")])] ["dir"; "pattern"];
  tool "sx_comp_list" "List all definitions (defcomp, defisland, defmacro, defpage, define) across .sx files in a directory."
    [dir_prop] ["dir"];
  tool "sx_comp_usage" "Find all uses of a component or symbol name across .sx files in a directory."
    [dir_prop; ("name", `Assoc [("type", `String "string"); ("description", `String "Component or symbol name to search for")])] ["dir"; "name"];
  tool "sx_diff" "Structural diff between two .sx files. Reports ADDED, REMOVED, CHANGED nodes with paths."
    [("file_a", `Assoc [("type", `String "string"); ("description", `String "Path to first .sx file")]);
     ("file_b", `Assoc [("type", `String "string"); ("description", `String "Path to second .sx file")])] ["file_a"; "file_b"];
  tool "sx_format_check" "Lint an .sx file for common issues: empty let bindings, missing bodies, duplicate params, structural problems."
    [file_prop] ["file"];
  tool "sx_macroexpand" "Evaluate an SX expression with a file's definitions loaded. Use to test macros — the file's defmacro forms are available."
    [("file", `Assoc [("type", `String "string"); ("description", `String "Optional .sx file to load for macro/component definitions")]);
     ("expr", `Assoc [("type", `String "string"); ("description", `String "Expression to expand/evaluate")])]
    ["expr"];
  tool "sx_build" "Build the SX runtime. Target \"js\" (default) builds sx-browser.js, \"ocaml\" runs dune build. Set full=true for extensions+types."
    [("target", `Assoc [("type", `String "string"); ("description", `String "Build target: \"js\" (default) or \"ocaml\"")]);
     ("full", `Assoc [("type", `String "boolean"); ("description", `String "Include extensions and type system (default: false)")])]
    [];
  tool "sx_test" "Run SX test suite. Returns pass/fail summary and any failures."
    [("host", `Assoc [("type", `String "string"); ("description", `String "Test host: \"js\" (default) or \"ocaml\"")]);
     ("full", `Assoc [("type", `String "boolean"); ("description", `String "Run full test suite including extensions (default: false)")])]
    [];
  tool "sx_pretty_print" "Reformat an .sx file with indentation. Short forms stay on one line, longer forms break across lines."
    [file_prop] ["file"];
  tool "sx_write_file" "Create or overwrite an .sx file. Source is parsed first — malformed SX is rejected and the file is not touched."
    [file_prop;
     ("source", `Assoc [("type", `String "string"); ("description", `String "SX source to write")])]
    ["file"; "source"];
  tool "sx_rename_symbol" "Rename all occurrences of a symbol in an .sx file. Structural — only renames symbols, not strings."
    [file_prop;
     ("old_name", `Assoc [("type", `String "string"); ("description", `String "Current symbol name")]);
     ("new_name", `Assoc [("type", `String "string"); ("description", `String "New symbol name")])]
    ["file"; "old_name"; "new_name"];
  tool "sx_replace_by_pattern" "Find nodes matching a pattern and replace with new source. Set all=true to replace all matches (default: first only)."
    [file_prop;
     ("pattern", `Assoc [("type", `String "string"); ("description", `String "Search pattern to match")]);
     ("new_source", `Assoc [("type", `String "string"); ("description", `String "Replacement SX source")]);
     ("all", `Assoc [("type", `String "boolean"); ("description", `String "Replace all matches (default: first only)")])]
    ["file"; "pattern"; "new_source"];
  tool "sx_insert_near" "Insert new source before or after the first node matching a pattern. No path needed."
    [file_prop;
     ("pattern", `Assoc [("type", `String "string"); ("description", `String "Pattern to find insertion point")]);
     ("new_source", `Assoc [("type", `String "string"); ("description", `String "SX source to insert")]);
     ("position", `Assoc [("type", `String "string"); ("description", `String "\"before\" or \"after\" (default: after)")])]
    ["file"; "pattern"; "new_source"];
  tool "sx_rename_across" "Rename a symbol across all .sx files in a directory. Use dry_run=true to preview without writing."
    [dir_prop;
     ("old_name", `Assoc [("type", `String "string"); ("description", `String "Current symbol name")]);
     ("new_name", `Assoc [("type", `String "string"); ("description", `String "New symbol name")]);
     ("dry_run", `Assoc [("type", `String "boolean"); ("description", `String "Preview changes without writing (default: false)")])]
    ["dir"; "old_name"; "new_name"];
]

(* ------------------------------------------------------------------ *)
(* JSON-RPC dispatch                                                   *)
(* ------------------------------------------------------------------ *)

let dispatch method_name params =
  match method_name with
  | "initialize" ->
    `Assoc [
      ("protocolVersion", `String "2024-11-05");
      ("capabilities", `Assoc [("tools", `Assoc [])]);
      ("serverInfo", `Assoc [
        ("name", `String "sx-tree-tools");
        ("version", `String "0.1.0")])]
  | "notifications/initialized" -> `Null
  | "tools/list" -> `Assoc [("tools", tool_definitions)]
  | "tools/call" ->
    let open Yojson.Safe.Util in
    let name = params |> member "name" |> to_string in
    let args = params |> member "arguments" in
    (try handle_tool name args
     with e -> error_result ("Error: " ^ Printexc.to_string e))
  | _ -> `Null

(* ------------------------------------------------------------------ *)
(* Stdio JSON-RPC main loop                                            *)
(* ------------------------------------------------------------------ *)

let () =
  setup_env ();
  try while true do
    let line = input_line stdin in
    if String.length line > 0 then begin
      let json = Yojson.Safe.from_string line in
      let open Yojson.Safe.Util in
      let meth = json |> member "method" |> to_string_option |> Option.value ~default:"" in
      let params = json |> member "params" in
      let id = json |> member "id" in
      let result = dispatch meth params in
      if id <> `Null then begin
        let resp = `Assoc [
          ("jsonrpc", `String "2.0");
          ("id", id);
          ("result", result)] in
        print_string (Yojson.Safe.to_string resp);
        print_char '\n';
        flush stdout
      end
    end
  done
  with End_of_file -> ()