open Mcp
open Jsonrpc

(* Logging utilities *)
let log_debug msg =
  Printf.eprintf "[DEBUG] %s\n" msg;
  flush stderr

let log_error msg =
  Printf.eprintf "[ERROR] %s\n" msg;
  flush stderr

(* Server state *)
let protocol_version = "2024-11-05"
let server_info = Implementation.{ name = "ocaml-mcp-capitalizer"; version = "0.1.0" }
let server_capabilities = `Assoc [
  (* We support tools *)
  ("tools", `Assoc [
    ("listChanged", `Bool true)
  ]);
  (* We don't support resources - make this explicit *)
  ("resources", `Assoc [
    ("listChanged", `Bool false);
    ("subscribe", `Bool false)
  ]);
  (* We don't support prompts - make this explicit *)
  ("prompts", `Assoc [
    ("listChanged", `Bool false)
  ])
]

(* Tool implementation *)
module CapitalizeTool = struct
  let name = "capitalize"
  let description = "Capitalizes the provided text"
  let input_schema = `Assoc [
    ("type", `String "object");
    ("properties", `Assoc [
      ("text", `Assoc [
        ("type", `String "string");
        ("description", `String "The text to capitalize")
      ])
    ]);
    ("required", `List [`String "text"])
  ]

  let call json =
    match json with
    | `Assoc fields ->
      (match List.assoc_opt "text" fields with
       | Some (`String text) ->
         let capitalized_text = String.uppercase_ascii text in
         let content = TextContent.{ 
           text = capitalized_text;
           annotations = None 
         } in
         `Assoc [
           ("content", `List [TextContent.yojson_of_t content]);
           ("isError", `Bool false)
         ]
       | _ -> 
         let error_content = TextContent.{
           text = "Missing or invalid 'text' parameter";
           annotations = None
         } in
         `Assoc [
           ("content", `List [TextContent.yojson_of_t error_content]);
           ("isError", `Bool true)
         ])
    | _ ->
      let error_content = TextContent.{
        text = "Invalid arguments format";
        annotations = None
      } in
      `Assoc [
        ("content", `List [TextContent.yojson_of_t error_content]);
        ("isError", `Bool true)
      ]
end

(* Handle tool listing *)
let list_tools () =
  let tool = `Assoc [
    ("name", `String CapitalizeTool.name);
    ("description", `String CapitalizeTool.description);
    ("inputSchema", CapitalizeTool.input_schema)
  ] in
  `Assoc [
    ("tools", `List [tool])
  ]

(* Handle tool calls *)
let call_tool name args =
  if name = CapitalizeTool.name then
    CapitalizeTool.call args
  else
    let error_content = TextContent.{
      text = Printf.sprintf "Unknown tool: %s" name;
      annotations = None
    } in
    `Assoc [
      ("content", `List [TextContent.yojson_of_t error_content]);
      ("isError", `Bool true)
    ]

(* Handle initialization *)
let handle_initialize id json =
  try
    log_debug (Printf.sprintf "Processing initialize request with id: %s" 
      (match id with 
       | `Int i -> string_of_int i 
       | `String s -> s));
    
    log_debug (Printf.sprintf "Initialize params: %s" 
      (match json with 
       | Some j -> Yojson.Safe.to_string j 
       | None -> "null"));
    
    let _ = match json with
      | Some params -> 
          log_debug "Parsing initialize request params...";
          let req = Initialize.Request.t_of_yojson params in
          log_debug (Printf.sprintf "Client info: %s v%s" req.client_info.name req.client_info.version);
          log_debug (Printf.sprintf "Client protocol version: %s" req.protocol_version);
          
          (* Check protocol version compatibility *)
          if req.protocol_version <> protocol_version then
            log_debug (Printf.sprintf "Protocol version mismatch: client=%s server=%s - will use server version" 
              req.protocol_version protocol_version);
          
          req
      | None -> 
          log_error "Missing params for initialize request";
          raise (Json.Of_json ("Missing params for initialize request", `Null))
    in
    
    log_debug "Creating initialize response...";
    let result = Initialize.Result.create
      ~capabilities:server_capabilities
      ~server_info
      ~protocol_version
      ~instructions:"This server provides a tool to capitalize text."
      ()
    in
    
    log_debug "Serializing initialize response...";
    let response = create_response ~id ~result:(Initialize.Result.yojson_of_t result) in
    log_debug "Initialize response created successfully";
    response
  with
  | Json.Of_json (msg, _) ->
    log_error (Printf.sprintf "JSON error in initialize: %s" msg);
    create_error ~id ~code:(-32602) ~message:("Invalid params: " ^ msg) ()
  | exc ->
    log_error (Printf.sprintf "Exception in initialize: %s" (Printexc.to_string exc));
    log_error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()));
    create_error ~id ~code:(-32603) ~message:("Internal error: " ^ (Printexc.to_string exc)) ()

(* Handle tools/list *)
let handle_list_tools id =
  log_debug "Processing tools/list request";
  let result = list_tools () in
  log_debug (Printf.sprintf "Tools list result: %s" (Yojson.Safe.to_string result));
  create_response ~id ~result

(* Handle tools/call *)
let handle_call_tool id json =
  try
    log_debug (Printf.sprintf "Processing tool call request with id: %s" 
      (match id with 
       | `Int i -> string_of_int i 
       | `String s -> s));
    
    log_debug (Printf.sprintf "Tool call params: %s" 
      (match json with 
       | Some j -> Yojson.Safe.to_string j 
       | None -> "null"));
    
    match json with
    | Some (`Assoc params) ->
      let name = match List.assoc_opt "name" params with
        | Some (`String name) -> 
            log_debug (Printf.sprintf "Tool name: %s" name);
            name
        | _ -> 
            log_error "Missing or invalid 'name' parameter in tool call";
            raise (Json.Of_json ("Missing or invalid 'name' parameter", `Assoc params))
      in
      let args = match List.assoc_opt "arguments" params with
        | Some (args) -> 
            log_debug (Printf.sprintf "Tool arguments: %s" (Yojson.Safe.to_string args));
            args
        | _ -> 
            log_debug "No arguments provided for tool call, using empty object";
            `Assoc [] (* Empty arguments is valid *)
      in
      log_debug (Printf.sprintf "Calling tool: %s" name);
      let result = call_tool name args in
      log_debug (Printf.sprintf "Tool call result: %s" (Yojson.Safe.to_string result));
      create_response ~id ~result
    | _ ->
      log_error "Invalid params format for tools/call";
      create_error ~id ~code:(-32602) ~message:"Invalid params for tools/call" ()
  with
  | Json.Of_json (msg, _) ->
    log_error (Printf.sprintf "JSON error in tool call: %s" msg);
    create_error ~id ~code:(-32602) ~message:("Invalid params: " ^ msg) ()
  | exc ->
    log_error (Printf.sprintf "Exception in tool call: %s" (Printexc.to_string exc));
    log_error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()));
    create_error ~id ~code:(-32603) ~message:("Internal error: " ^ (Printexc.to_string exc)) ()

(* Handle ping *)
let handle_ping id =
  create_response ~id ~result:(`Assoc [])

(* Process a single message *)
let process_message message =
  try
    log_debug "Parsing message as JSONRPC message...";
    match JSONRPCMessage.t_of_yojson message with
    | JSONRPCMessage.Request req ->
      log_debug (Printf.sprintf "Received request with method: %s" req.method_);
      (match req.method_ with
       | "initialize" -> 
         log_debug "Processing initialize request";
         Some (handle_initialize req.id req.params)
       | "tools/list" -> 
         log_debug "Processing tools/list request";
         Some (handle_list_tools req.id)
       | "tools/call" -> 
         log_debug "Processing tools/call request";
         Some (handle_call_tool req.id req.params)
       | "ping" -> 
         log_debug "Processing ping request";
         Some (handle_ping req.id)
       | _ -> 
         log_error (Printf.sprintf "Unknown method received: %s" req.method_);
         Some (create_error ~id:req.id ~code:(-32601) ~message:("Method not found: " ^ req.method_) ()))
    | JSONRPCMessage.Notification notif ->
      log_debug (Printf.sprintf "Received notification with method: %s" notif.method_);
      (match notif.method_ with
       | "notifications/initialized" ->
         log_debug "Client initialization complete - Server is now ready to receive requests";
         log_debug (Printf.sprintf "Notification params: %s" 
           (match notif.params with
            | Some p -> Yojson.Safe.to_string p
            | None -> "null"));
         None
       | _ ->
         log_debug (Printf.sprintf "Ignoring notification: %s" notif.method_);
         None)
    | JSONRPCMessage.Response _ ->
      log_error "Unexpected response message received";
      None
    | JSONRPCMessage.Error _ ->
      log_error "Unexpected error message received";
      None
  with
  | exc ->
    log_error (Printf.sprintf "Exception during message processing: %s" (Printexc.to_string exc));
    log_error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()));
    log_error (Printf.sprintf "Message was: %s" (Yojson.Safe.to_string message));
    None

(* Main loop *)
let rec read_message () =
  try
    log_debug "Attempting to read line from stdin...";
    let line = read_line () in
    if line = "" then (
      log_debug "Empty line received, ignoring";
      None
    ) else (
      log_debug (Printf.sprintf "Raw input: %s" line);
      try
        let json = Yojson.Safe.from_string line in
        log_debug "Successfully parsed JSON";
        Some json
      with
      | Yojson.Json_error msg -> 
        log_error (Printf.sprintf "Error parsing JSON: %s" msg);
        log_error (Printf.sprintf "Input was: %s" line);
        read_message ()
    )
  with
  | End_of_file -> 
    log_debug "End of file received on stdin";
    None
  | Sys_error msg -> 
    log_error (Printf.sprintf "System error while reading: %s" msg);
    None
  | exc ->
    log_error (Printf.sprintf "Exception while reading: %s" (Printexc.to_string exc));
    None

let () =
  try
    (* Enable exception backtraces *)
    Printexc.record_backtrace true;
    
    (* Enable line buffering for stdout *)
    set_binary_mode_out stdout false;
    
    log_debug "MCP Capitalizer server started";
    log_debug (Printf.sprintf "Protocol version: %s" protocol_version);
    log_debug (Printf.sprintf "Server info: %s v%s" server_info.name server_info.version);
    
    (* Print environment info for debugging *)
    log_debug "Environment variables:";
    Unix.environment()
    |> Array.iter (fun s -> 
      try
        let i = String.index s '=' in
        let name = String.sub s 0 i in
        if String.length name > 0 then
          log_debug (Printf.sprintf "  %s" s)
      with Not_found -> ()
    );
    
    let rec server_loop count =
      log_debug (Printf.sprintf "Waiting for message #%d..." count);
      match read_message () with
      | Some json ->
        log_debug (Printf.sprintf "Received message: %s" (Yojson.Safe.to_string json));
        (match process_message json with
         | Some response ->
           let response_json = JSONRPCMessage.yojson_of_t response in
           let response_str = Yojson.Safe.to_string response_json in
           log_debug (Printf.sprintf "Sending response: %s" response_str);
           (* Make sure we emit properly formatted JSON on a single line with a newline at the end *)
           Printf.printf "%s\n" response_str;
           flush stdout;
           (* Give the client a moment to process the response *)
           Unix.sleepf 0.01;
           server_loop (count + 1)
         | None -> 
           log_debug "No response needed for this message";
           server_loop (count + 1))
      | None -> 
        log_debug "End of input stream, terminating server";
        ()
    in
    
    log_debug "Starting server loop...";
    log_debug "Waiting for the initialize request...";
    
    (* Set up signal handler to gracefully exit *)
    Sys.(set_signal sigint (Signal_handle (fun _ -> 
      log_debug "Received interrupt signal, exiting...";
      exit 0
    )));
    
    server_loop 1;
    log_debug "Server terminated normally";
  with
  | End_of_file ->
    log_error "Unexpected end of file";
  | Sys_error msg ->
    log_error (Printf.sprintf "System error: %s" msg);
  | Unix.Unix_error(err, func, arg) ->
    log_error (Printf.sprintf "Unix error in %s(%s): %s" func arg (Unix.error_message err));
  | exc ->
    log_error (Printf.sprintf "Unhandled exception: %s" (Printexc.to_string exc));
    log_error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()))