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.debugf "Client info: %s v%s" 
          req_data.client_info.name req_data.client_info.version;
        Log.debugf "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.debugf "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.debugf "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 = ErrorCode.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) ()


(* 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.debugf "Found tool: %s" name;
    
    (* Call the tool handler *)
    match tool.handler ctx args with
    | Ok result -> 
        Log.debug "Tool execution succeeded";
        result
    | Error err -> Tool.handle_execution_error err
  with
  | Not_found -> Tool.handle_unknown_tool_error name
  | exn -> Tool.handle_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)
          ?progress_token:req.progress_token
          ~lifespan_context:[("tools/call", `Assoc params)] 
          () 
        in
                
        (* Execute the tool *)
        let result = execute_tool server ctx name args in
        
        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.debugf "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.debugf "Processing message: %s" (Yojson.Safe.to_string message);
    match JSONRPCMessage.t_of_yojson message with
    | JSONRPCMessage.Request req ->
      Log.debugf "Received request with method: %s" (Method.to_string req.meth);
      (match req.meth with
       | Method.Initialize -> handle_initialize server req
       | Method.ToolsList -> handle_tools_list server req
       | Method.ToolsCall -> handle_tools_call server req
       | Method.PromptsList -> handle_prompts_list server req
       | Method.ResourcesList -> handle_resources_list server req
       | _ -> 
          Log.errorf "Unknown method received: %s" (Method.to_string req.meth);
          Some (create_jsonrpc_error req.id MethodNotFound ("Method not found: " ^ (Method.to_string req.meth)) ()))
    | JSONRPCMessage.Notification notif ->
      Log.debugf "Received notification with method: %s" (Method.to_string notif.meth);
      (match notif.meth with
       | Method.Initialized -> handle_initialized notif
       | _ ->
          Log.debugf "Ignoring notification: %s" (Method.to_string notif.meth);
          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.errorf "JSON error: %s" msg;
    (* Can't respond with error because we don't have a request ID *)
    None
  | Yojson.Json_error msg -> 
    Log.errorf "JSON parse error: %s" msg;
    (* Can't respond with error because we don't have a request ID *)
    None
  | exc ->
    Log.errorf "Exception during message processing: %s" (Printexc.to_string exc);
    Log.errorf "Backtrace: %s" (Printexc.get_backtrace());
    Log.errorf "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.debugf "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.errorf "Error parsing JSON: %s" msg;
        Log.errorf "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.debugf "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.debugf "Starting MCP server: %s v%s" (name server) (version server);
  Log.debugf "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.errorf "I/O error while reading: %s" (Printexc.to_string exn);
      ()
  | exn ->
      Log.errorf "Exception while reading: %s" (Printexc.to_string exn);
      ()