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

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

let rec handle_tool name args =
  let open Yojson.Safe.Util in
  match name with
  | "sx_read_tree" ->
    let tree = parse_file (args |> member "file" |> to_string) in
    text_result (value_to_string (call_sx "annotate-tree" [tree]))

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

  | _ -> 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 parts = match new_tree with
         | List items | ListRef { contents = items } ->
           List.map (fun expr -> Sx_runtime.value_to_str expr) items
         | _ -> [Sx_runtime.value_to_str new_tree]
       in
       let source = String.concat "\n\n" parts ^ "\n" 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 tool_definitions = `List [
  tool "sx_read_tree" "Read an .sx file as an annotated tree with path labels on every node. Use this to understand structure before editing."
    [file_prop] ["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"];
]

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