open Mcp
open Jsonrpc
open Mcp_sdk

(* Process initialize request *)
let handle_initialize server req =
  Log.debug "Processing initialize request";
  let result = match req.JSONRPCMessage.params with
    | Some params -> 
        let req_data = Initialize.Request.t_of_yojson params in
        Log.debug (Printf.sprintf "Client info: %s v%s" 
          req_data.client_info.name req_data.client_info.version);
        Log.debug (Printf.sprintf "Client protocol version: %s" req_data.protocol_version);
        
        (* Create initialize response *)
        let result = Initialize.Result.create
          ~capabilities:(capabilities server)
          ~server_info:Implementation.{ 
            name = name server; 
            version = version server 
          }
          ~protocol_version:(protocol_version server)
          ~instructions:(Printf.sprintf "This server provides tools for %s." (name server))
          ()
        in
        Initialize.Result.yojson_of_t result
    | None -> 
        Log.error "Missing params for initialize request";
        `Assoc [("error", `String "Missing params for initialize request")]
  in
  Some (create_response ~id:req.id ~result)

(* Process tools/list request *)
let handle_tools_list server (req:JSONRPCMessage.request) =
  Log.debug "Processing tools/list request";
  let tool_list = List.map Tool.to_json (tools server) in
  let result = `Assoc [("tools", `List tool_list)] in
  Some (create_response ~id:req.id ~result)

(* Process prompts/list request *)
let handle_prompts_list server (req:JSONRPCMessage.request) =
  Log.debug "Processing prompts/list request";
  let prompt_list = List.map Prompt.to_json (prompts server) in
  let result = `Assoc [("prompts", `List prompt_list)] in
  Some (create_response ~id:req.id ~result)

(* Process resources/list request *)
let handle_resources_list server (req:JSONRPCMessage.request) =
  Log.debug "Processing resources/list request";
  let resource_list = List.map Resource.to_json (resources server) in
  let result = `Assoc [("resources", `List resource_list)] in
  Some (create_response ~id:req.id ~result)

(* Extract the tool name from params *)
let extract_tool_name params =
  match List.assoc_opt "name" params with
    | Some (`String name) -> 
        Log.debug (Printf.sprintf "Tool name: %s" name);
        Some name
    | _ -> 
        Log.error "Missing or invalid 'name' parameter in tool call";
        None

(* Extract the tool arguments from params *)
let extract_tool_arguments params =
  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 *)

(* Create a proper JSONRPC error with code and data *)
let create_jsonrpc_error id code message ?data () =
  let error_code = error_code_to_int code in
  let error_data = match data with
    | Some d -> d
    | None -> `Null
  in
  create_error ~id ~code:error_code ~message ~data:(Some error_data) ()

(* Create a tool error result with structured content *)
let create_tool_error_result error =
  create_tool_result [TextContent error] ~is_error:true

(* Handle tool execution errors *)
let handle_tool_execution_error err =
  Log.error (Printf.sprintf "Tool execution failed: %s" err);
  create_tool_error_result (Printf.sprintf "Error executing tool: %s" err)

(* Handle unknown tool error *)
let handle_unknown_tool_error name =
  Log.error (Printf.sprintf "Unknown tool: %s" name);
  create_tool_error_result (Printf.sprintf "Unknown tool: %s" name)

(* Handle general tool execution exception *)
let handle_tool_execution_exception exn =
  Log.error (Printf.sprintf "Exception executing tool: %s" (Printexc.to_string exn));
  create_tool_error_result (Printf.sprintf "Internal error: %s" (Printexc.to_string exn))

(* Execute a tool *)
let execute_tool server ctx name args =
  try
    let tool = List.find (fun t -> t.Tool.name = name) (tools server) in
    Log.debug (Printf.sprintf "Found tool: %s" name);
    
    (* Call the tool handler *)
    match tool.handler ctx args with
    | Ok result -> 
        Log.debug "Tool execution succeeded";
        result
    | Error err -> handle_tool_execution_error err
  with
  | Not_found -> handle_unknown_tool_error name
  | exn -> handle_tool_execution_exception exn

(* Process tools/call request *)
let handle_tools_call server req =
  Log.debug "Processing tools/call request";
  match req.JSONRPCMessage.params with
  | Some (`Assoc params) ->
    (match extract_tool_name params with
    | Some name ->
        let args = extract_tool_arguments params in
        
        (* Create context for this request *)
        let ctx = Context.create 
          ?request_id:(Some req.id) (* Store request ID for progress reporting *)
          ~lifespan_context:[("tools/call", `Assoc params)] (* Store params for reference *)
          () 
        in
        
        (* Set progress token if present *)
        ctx.progress_token <- req.progress_token;
        
        (* Execute the tool *)
        let result = execute_tool server ctx name args in
        
        (* Process progress messages if any *)
        let progress_msg = Context.report_progress ctx 1.0 1.0 in
        (match progress_msg with
         | Some msg -> Log.debug "Progress complete notification would be sent here";
         | None -> ());
        
        Some (create_response ~id:req.id ~result)
    | None ->
        Some (create_jsonrpc_error req.id InvalidParams "Missing tool name parameter" ()))
  | _ ->
    Log.error "Invalid params format for tools/call";
    Some (create_jsonrpc_error req.id InvalidParams "Invalid params format for tools/call" ())

(* Process ping request *)
let handle_ping (req:JSONRPCMessage.request) =
  Log.debug "Processing ping request";
  Some (create_response ~id:req.JSONRPCMessage.id ~result:(`Assoc []))

(* Handle notifications/initialized *)
let handle_initialized (notif:JSONRPCMessage.notification) =
  Log.debug "Client initialization complete - Server is now ready to receive requests";
  Log.debug (Printf.sprintf "Notification params: %s" 
    (match notif.JSONRPCMessage.params with
      | Some p -> Yojson.Safe.to_string p
      | None -> "null"));
  None

(* Process a single message using the MCP SDK *)
let process_message server message =
  try
    Log.debug (Printf.sprintf "Processing message: %s" (Yojson.Safe.to_string 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" -> handle_initialize server req
       | "tools/list" -> handle_tools_list server req
       | "tools/call" -> handle_tools_call server req
       | "prompts/list" -> handle_prompts_list server req
       | "resources/list" -> handle_resources_list server req
       | "ping" -> handle_ping req
       | _ -> 
          Log.error (Printf.sprintf "Unknown method received: %s" req.method_);
          Some (create_jsonrpc_error req.id MethodNotFound ("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" -> handle_initialized notif
       | _ ->
          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
  | Json.Of_json (msg, _) ->
    Log.error (Printf.sprintf "JSON error: %s" msg);
    (* Can't respond with error because we don't have a request ID *)
    None
  | Yojson.Json_error msg -> 
    Log.error (Printf.sprintf "JSON parse error: %s" msg);
    (* Can't respond with error because we don't have a request ID *)
    None
  | 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

(* Extract a request ID from a potentially malformed message *)
let extract_request_id json =
  try
    match json with
    | `Assoc fields ->
      (match List.assoc_opt "id" fields with
       | Some (`Int id) -> Some (`Int id)
       | Some (`String id) -> Some (`String id)
       | _ -> None)
    | _ -> None
  with _ -> None

(* Handle processing for an input line *)
let process_input_line server line =
  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";
      
      (* Process the message *)
      process_message server json
    with
    | Yojson.Json_error msg -> begin
        Log.error (Printf.sprintf "Error parsing JSON: %s" msg);
        Log.error (Printf.sprintf "Input was: %s" line);
        None
    end
  )

(* Send a response to the client *)
let send_response stdout 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);
  
  (* Write the response followed by a newline *)
  Eio.Flow.copy_string response_str stdout;
  Eio.Flow.copy_string "\n" stdout

(* Run the MCP server with the given server configuration *)
let run_server env server =
  let stdin = Eio.Stdenv.stdin env in
  let stdout = Eio.Stdenv.stdout env in

  Log.debug (Printf.sprintf "Starting MCP server: %s v%s" (name server) (version server));
  Log.debug (Printf.sprintf "Protocol version: %s" (protocol_version server));

  (* Enable exception backtraces *)
  Printexc.record_backtrace true;
  
  let buf = Eio.Buf_read.of_flow stdin ~initial_size:100 ~max_size:1_000_000 in
  
  (* Main processing loop *)
  try
    while true do
      Log.debug "Waiting for message...";
      let line = Eio.Buf_read.line buf in
      
      (* Process the input and send response if needed *)
      match process_input_line server line with
      | Some response -> send_response stdout response
      | None -> Log.debug "No response needed for this message"
    done
  with
  | End_of_file -> 
      Log.debug "End of file received on stdin";
      ()
  | Eio.Exn.Io _ as exn ->
      Log.error (Printf.sprintf "I/O error while reading: %s" (Printexc.to_string exn));
      ()
  | exn ->
      Log.error (Printf.sprintf "Exception while reading: %s" (Printexc.to_string exn));
      ()