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 Capitalizer" 
  ~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 "CapitalizeServer is starting up!\n";
  flush stderr;
  Log.info "CapitalizeServer is starting up!"

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

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

(* Define and register a capitalize tool *)
let _ = add_tool server
  ~name:"capitalize"
  ~description:"Capitalizes the provided text"
  ~schema_properties:[
    ("text", "string", "The text to capitalize")
  ]
  ~schema_required:["text"]
  (fun args ->
    try
      let text = get_string_param args "text" in
      let capitalized_text = String.uppercase_ascii text in
      TextContent.yojson_of_t TextContent.{
        text = capitalized_text;
        annotations = None
      }
    with
    | Failure msg ->
      Log.error (Printf.sprintf "Error in capitalize tool: %s" msg);
      TextContent.yojson_of_t TextContent.{
        text = Printf.sprintf "Error: %s" msg;
        annotations = None
      }
  )

(* Define and register a resource example *)
let _ = add_resource server
  ~uri_template:"greeting://{name}"
  ~description:"Get a greeting for a name"
  ~mime_type:"text/plain"
  (fun params ->
    match params with
    | [name] -> Printf.sprintf "Hello, %s! Welcome to the OCaml MCP server." name
    | _ -> "Hello, world! Welcome to the OCaml MCP server."
  )

(* Define and register a prompt example *)
let _ = add_prompt server
  ~name:"capitalize-prompt"
  ~description:"A prompt to help with text capitalization"
  ~arguments:[
    ("text", Some "The text to be capitalized", true)
  ]
  (fun args ->
    let text = 
      try
        List.assoc "text" args
      with
      | Not_found -> "No text provided"
    in
    [
      Prompt.{ 
        role = `User; 
        content = make_text_content "Please help me capitalize the following text:"
      };
      Prompt.{ 
        role = `User; 
        content = make_text_content text
      };
      Prompt.{ 
        role = `Assistant; 
        content = make_text_content "Here's the capitalized version:"
      };
      Prompt.{ 
        role = `Assistant; 
        content = make_text_content (String.uppercase_ascii text)
      }
    ]
  )

(* 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: capitalize_sdk [--http]" in
  Arg.parse args (fun _ -> ()) usage_msg;

  (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
  Printf.fprintf stderr "Starting CapitalizeServer...\n";
  flush stderr;
  Log.info "Starting CapitalizeServer...";
  
  (* Configure the server with appropriate capabilities *)
  ignore (configure_server server ());
  
  (* Create and start MCP server with the selected transport *)
  let mcp_server = create ~server ~transport:!transport_type () in
  start mcp_server