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 Resource Template 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 "ResourceTemplateServer is starting up!\n";
  flush stderr;
  Log.info "ResourceTemplateServer is starting up!"

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

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

(* Database of "documents" *)
let documents = [
  ("doc1", "This is the first document content");
  ("doc2", "This document contains information about OCaml");
  ("doc3", "MCP protocol is a standard for LLM-based applications");
  ("doc4", "Resource templates allow for parameterized resources");
]

(* Define and register a resource template for documents *)
let _ = add_resource server
  ~uri_template:"document://{id}"
  ~description:"Get a document by ID"
  ~mime_type:"text/plain"
  (fun params ->
    match params with
    | [id] -> 
        begin
          try
            let content = List.assoc id documents in
            content
          with Not_found -> 
            Printf.sprintf "Error: Document '%s' not found" id
        end
    | _ -> "Error: Invalid document ID"
  )

(* Define and register a list documents resource *)
let _ = add_resource server
  ~uri_template:"documents://list"
  ~description:"List all available documents"
  ~mime_type:"text/plain"
  (fun _ ->
    let doc_list = 
      String.concat "\n" 
        (List.map (fun (id, _) -> Printf.sprintf "- %s" id) documents)
    in
    Printf.sprintf "Available Documents:\n%s" doc_list
  )

(* Define and register a tool that uses resource references *)
let _ = add_tool server
  ~name:"get_document"
  ~description:"Gets a document by ID using resource references"
  ~schema_properties:[
    ("document_id", "string", "The ID of the document to retrieve");
  ]
  ~schema_required:["document_id"]
  (fun args ->
    try
      let doc_id = get_string_param args "document_id" in
      
      (* Create a resource reference *)
      let ref = ResourceReference.{ uri = Printf.sprintf "document://%s" doc_id } in
      (* Convert to JSON for logging purposes *)
      let _ = ResourceReference.yojson_of_t ref in
      
      (* Return the reference *)
      CallToolResult.yojson_of_t CallToolResult.{
        content = [
          Text TextContent.{ text = Printf.sprintf "Resource reference for document %s:" doc_id; annotations = None };
          Resource EmbeddedResource.{ 
            resource = `Text TextResourceContents.{
              uri = Printf.sprintf "document://%s" doc_id;
              text = (try List.assoc doc_id documents with Not_found -> "Not found");
              mime_type = Some "text/plain"
            };
            annotations = None
          }
        ];
        is_error = false;
        meta = None
      }
    with
    | Failure msg ->
      Log.error (Printf.sprintf "Error in get_document tool: %s" msg);
      CallToolResult.yojson_of_t CallToolResult.{
        content = [
          Text TextContent.{ 
            text = Printf.sprintf "Error: %s" msg; 
            annotations = None 
          }
        ];
        is_error = true;
        meta = None
      }
  )

(* Define and register a prompt that uses resource templates *)
let _ = add_prompt server
  ~name:"document-prompt"
  ~description:"A prompt that references document resources"
  ~arguments:[
    ("document_id", Some "ID of the document to include in the prompt", true);
  ]
  (fun args ->
    let doc_id = 
      try List.assoc "document_id" args
      with Not_found -> "doc1" (* Default to doc1 *)
    in
    
    let doc_text = 
      try List.assoc doc_id documents
      with Not_found -> Printf.sprintf "Document '%s' not found" doc_id
    in
    
    [
      Prompt.{ 
        role = `User; 
        content = make_text_content (Printf.sprintf "Please summarize the following document (ID: %s):" doc_id)
      };
      Prompt.{ 
        role = `User; 
        content = Resource EmbeddedResource.{
          resource = `Text TextResourceContents.{
            uri = Printf.sprintf "document://%s" doc_id;
            text = doc_text;
            mime_type = Some "text/plain"
          };
          annotations = None
        }
      };
      Prompt.{ 
        role = `Assistant; 
        content = make_text_content "I'll help summarize this document for you."
      }
    ]
  )

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

  (* Instead of printing directly to stdout which messes up the JSON-RPC protocol,
     use the logging system which sends output to stderr *)
  Log.info "Starting ResourceTemplateServer...";
  
  (* 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