open Mcp
open Mcp_sdk
open Mcp_server

(* Helper for extracting string value from JSON *)
let get_string_param json name =
  match json with
  | `Assoc fields ->
    (match List.assoc_opt name fields with
    | Some (`String value) -> value
    | _ -> raise (Failure (Printf.sprintf "Missing or invalid parameter: %s" name)))
  | _ -> raise (Failure "Expected JSON object")

(* Create a server *)
let server = create_server 
  ~name:"OCaml MCP Completion Example" 
  ~version:"0.1.0" 
  ~protocol_version:"2024-11-05" 
  ()

(* Define startup and shutdown hooks *)
let startup () =
  (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
  Printf.fprintf stderr "CompletionServer is starting up!\n";
  flush stderr;
  Log.info "CompletionServer is starting up!"

let shutdown () =
  Printf.fprintf stderr "CompletionServer is shutting down. Goodbye!\n";
  flush stderr;
  Log.info "CompletionServer is shutting down. Goodbye!"

(* Register the hooks *)
let () =
  set_startup_hook server startup;
  set_shutdown_hook server shutdown

(* Database of programming languages and their features *)
let languages = [
  ("ocaml", ["functional"; "static typing"; "pattern matching"; "modules"; "type inference"]);
  ("python", ["dynamic typing"; "interpreted"; "object-oriented"; "high-level"; "scripting"]);
  ("rust", ["memory safety"; "performance"; "static typing"; "ownership"; "zero-cost abstractions"]);
  ("javascript", ["dynamic typing"; "interpreted"; "prototypes"; "single-threaded"; "event-driven"]);
  ("go", ["concurrency"; "garbage collection"; "simplicity"; "static typing"; "compiled"]);
]

(* Helper function to create a completion response *)
let create_completion values ?(has_more=false) ?(total=None) () =
  Completion.Result.{
    completion = {
      values;
      has_more = Some has_more;
      total;
    };
    meta = None;
  }

(* Define and register a tool that handles completions *)
let _ = add_tool server
  ~name:"complete"
  ~description:"Handles completion requests for programming languages and features"
  ~schema_properties:[
    ("argument_name", "string", "The name of the argument to complete");
    ("argument_value", "string", "The partial value to complete");
  ]
  ~schema_required:["argument_name"; "argument_value"]
  (fun args ->
    try
      let argument_name = get_string_param args "argument_name" in
      let argument_value = get_string_param args "argument_value" in
      
      Log.info (Printf.sprintf "Completion request for %s = %s" argument_name argument_value);
      
      (* Handle different completion requests *)
      let result = 
        match argument_name with
        | "language" ->
            (* Complete programming language names *)
            let matches = 
              List.filter 
                (fun (lang, _) -> 
                  let lang_lower = String.lowercase_ascii lang in
                  let arg_lower = String.lowercase_ascii argument_value in
                  String.length lang_lower >= String.length arg_lower &&
                  String.sub lang_lower 0 (String.length arg_lower) = arg_lower)
                languages
            in
            let values = List.map fst matches in
            create_completion values ~has_more:false ~total:(Some (List.length values)) ()
            
        | "feature" ->
            (* Complete programming language features *)
            let all_features = 
              List.flatten (List.map snd languages) |>
              List.sort_uniq String.compare
            in
            let matches =
              List.filter
                (fun feature -> 
                  let feature_lower = String.lowercase_ascii feature in
                  let arg_lower = String.lowercase_ascii argument_value in
                  String.length feature_lower >= String.length arg_lower &&
                  String.sub feature_lower 0 (String.length arg_lower) = arg_lower)
                all_features
            in
            create_completion matches ~has_more:false ~total:(Some (List.length matches)) ()
            
        | _ ->
            (* Default completions for unknown arguments *)
            create_completion ["unknown argument"] ()
      in
      
      (* Convert to JSON and return *)
      TextContent.yojson_of_t TextContent.{
        text = Yojson.Safe.to_string (Completion.Result.to_result result);
        annotations = None
      }
    with
    | Failure msg ->
        Log.error (Printf.sprintf "Error handling completion request: %s" msg);
        TextContent.yojson_of_t TextContent.{
          text = Printf.sprintf "Error: %s" msg;
          annotations = None
        }
  )

(* Define and register a prompt that provides programming language info *)
let _ = add_prompt server
  ~name:"language-info-prompt"
  ~description:"A prompt that provides information about programming languages"
  ~arguments:[
    ("language", Some "Name of the programming language", true);
  ]
  (fun args ->
    let language = 
      try List.assoc "language" args
      with Not_found -> "ocaml" (* Default to OCaml *)
    in
    
    let features =
      try 
        let features = List.assoc (String.lowercase_ascii language) languages in
        String.concat ", " features
      with Not_found -> "unknown language"
    in
    
    [
      Prompt.{ 
        role = `User; 
        content = make_text_content (Printf.sprintf "Tell me about the %s programming language" language)
      };
      Prompt.{ 
        role = `Assistant; 
        content = make_text_content (Printf.sprintf "%s is a programming language with the following features: %s" language features)
      }
    ]
  )

(* Main function *)
let () =
  (* Parse command line arguments *)
  let transport_type = ref Stdio in
  let args = [
    ("--http", Arg.Unit (fun () -> transport_type := Http), 
     "Start server with HTTP transport (default is stdio)");
  ] in
  let usage_msg = "Usage: completion_example [--http]" in
  Arg.parse args (fun _ -> ()) usage_msg;

  (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
  Printf.fprintf stderr "Starting CompletionServer...\n";
  flush stderr;
  Log.info "Starting CompletionServer...";
  
  (* Set custom capabilities to indicate support for completions *)
  let capabilities = `Assoc [
    ("completions", `Assoc []);  (* Indicate support for completions *)
    ("prompts", `Assoc [
      ("listChanged", `Bool true)
    ]);
    ("resources", `Assoc [
      ("listChanged", `Bool true);
      ("subscribe", `Bool true)
    ]);
    ("tools", `Assoc [
      ("listChanged", `Bool true)
    ])
  ] in
  set_capabilities server capabilities;
  
  (* Create and start MCP server with the selected transport *)
  let mcp_server = create ~server ~transport:!transport_type () in
  start mcp_server