open Mcp
open Mcp_sdk

(* MCP Server module for handling JSON-RPC communication *)

(** Server types *)
type transport_type = 
  | Stdio  (* Read/write to stdin/stdout *)
  | Http   (* HTTP server - to be implemented *)

type t = {
  server: Mcp_sdk.server;
  transport: transport_type;
  mutable running: bool;
}

(** Process a single message *)
let process_message server message =
  try
    Log.debug "Parsing message as JSONRPC message...";
    match JSONRPCMessage.t_of_yojson message with
    | JSONRPCMessage.Request req -> begin
        Log.debug (Printf.sprintf "Received request with method: %s" req.method_);
        match req.method_ with
        | "initialize" -> begin
            Log.debug "Processing initialize request";
            let result = match req.params with
            | Some params -> begin
                Log.debug "Parsing initialize request params...";
                let req_params = Initialize.Request.t_of_yojson params in
                Log.debug (Printf.sprintf "Client info: %s v%s" 
                             req_params.client_info.name 
                             req_params.client_info.version);
                Log.debug (Printf.sprintf "Client protocol version: %s" req_params.protocol_version);
                
                (* Check protocol version compatibility *)
                if req_params.protocol_version <> server.protocol_version then begin
                  Log.debug (Printf.sprintf "Protocol version mismatch: client=%s server=%s" 
                              req_params.protocol_version server.protocol_version);
                end;
                
                Initialize.Result.create
                  ~capabilities:server.capabilities
                  ~server_info:Implementation.{ name = server.name; version = server.version }
                  ~protocol_version:server.protocol_version
                  ?instructions:(Some "MCP Server") (* TODO: Allow customization *)
                  ()
              end
            | None -> begin
                Log.error "Missing params for initialize request";
                Initialize.Result.create
                  ~capabilities:server.capabilities
                  ~server_info:Implementation.{ name = server.name; version = server.version }
                  ~protocol_version:server.protocol_version
                  ()
              end
            in
            Some (create_response ~id:req.id ~result:(Initialize.Result.yojson_of_t result))
          end
         
        | "tools/list" -> begin
            Log.debug "Processing tools/list request";
            let tools_json = List.map Mcp_sdk.Tool.to_json server.tools in
            let result = `Assoc [("tools", `List tools_json)] in
            Some (create_response ~id:req.id ~result)
          end
         
        | "tools/call" -> begin
            Log.debug "Processing tools/call request";
            match req.params with
            | Some (`Assoc params) -> begin
                let name = match List.assoc_opt "name" params with
                | Some (`String name) -> begin
                    Log.debug (Printf.sprintf "Tool name: %s" name);
                    name
                  end
                | _ -> begin
                    Log.error "Missing or invalid 'name' parameter in tool call";
                    failwith "Missing or invalid 'name' parameter"
                  end
                in
                let args = match List.assoc_opt "arguments" params with
                | Some args -> begin
                    Log.debug (Printf.sprintf "Tool arguments: %s" (Yojson.Safe.to_string args));
                    args
                  end
                | _ -> begin
                    Log.debug "No arguments provided for tool call, using empty object";
                    `Assoc [] (* Empty arguments is valid *)
                  end
                in
                let progress_token = req.progress_token in
                
                (* Find the tool *)
                let tool_opt = List.find_opt (fun t -> t.Mcp_sdk.Tool.name = name) server.tools in
                match tool_opt with
                | Some tool -> begin
                    Log.debug (Printf.sprintf "Found tool: %s" name);
                    let ctx = Mcp_sdk.Context.create 
                        ?request_id:(Some req.id) 
                        ~lifespan_context:server.lifespan_context 
                        () 
                    in
                    ctx.progress_token <- progress_token;
                    
                    (* Call the handler *)
                    let result = match tool.handler ctx args with
                    | Ok json -> begin
                        `Assoc [
                          ("content", `List [Mcp.yojson_of_content (Text (TextContent.{ 
                            text = Yojson.Safe.to_string json;
                            annotations = None 
                          }))]);
                          ("isError", `Bool false)
                        ]
                      end
                    | Error err -> begin
                        `Assoc [
                          ("content", `List [Mcp.yojson_of_content (Text (TextContent.{ 
                            text = err;
                            annotations = None 
                          }))]);
                          ("isError", `Bool true)
                        ]
                      end
                    in
                    Some (create_response ~id:req.id ~result)
                  end
                | None -> begin
                    Log.error (Printf.sprintf "Tool not found: %s" name);
                    let error_content = TextContent.{
                      text = Printf.sprintf "Unknown tool: %s" name;
                      annotations = None
                    } in
                    let result = `Assoc [
                      ("content", `List [Mcp.yojson_of_content (Text error_content)]);
                      ("isError", `Bool true)
                    ] in
                    Some (create_response ~id:req.id ~result)
                  end
              end
            | _ -> begin
                Log.error "Invalid params format for tools/call";
                Some (create_error ~id:req.id ~code:ErrorCode.invalid_params ~message:"Invalid params for tools/call" ())
              end
          end
            
        | "resources/list" -> begin
            Log.debug "Processing resources/list request";
            if server.resources <> [] then begin
              let resources_json = List.map Mcp_sdk.Resource.to_json server.resources in
              let result = `Assoc [("resources", `List resources_json)] in
              Some (create_response ~id:req.id ~result)
            end else begin
              Some (create_error ~id:req.id ~code:ErrorCode.method_not_found ~message:"Resources not supported" ())
            end
          end
           
        | "prompts/list" -> begin
            Log.debug "Processing prompts/list request";
            if server.prompts <> [] then begin
              let prompts_json = List.map Mcp_sdk.Prompt.to_json server.prompts in
              let result = `Assoc [("prompts", `List prompts_json)] in
              Some (create_response ~id:req.id ~result)
            end else begin
              Some (create_error ~id:req.id ~code:ErrorCode.method_not_found ~message:"Prompts not supported" ())
            end
          end
           
        | "prompts/get" -> begin
            Log.debug "Processing prompts/get request";
            if server.prompts <> [] then begin
              match req.params with
              | Some (`Assoc params) -> begin
                  (* Extract prompt name *)
                  let name = match List.assoc_opt "name" params with
                  | Some (`String name) -> begin 
                      Log.debug (Printf.sprintf "Prompt name: %s" name);
                      name
                    end
                  | _ -> begin
                      Log.error "Missing or invalid 'name' parameter in prompt request";
                      failwith "Missing or invalid 'name' parameter"
                    end
                  in
                  
                  (* Extract arguments if any *)
                  let arguments = match List.assoc_opt "arguments" params with
                  | Some (`Assoc args) -> begin
                      Log.debug (Printf.sprintf "Prompt arguments: %s" (Yojson.Safe.to_string (`Assoc args)));
                      List.map (fun (k, v) -> 
                        match v with 
                        | `String s -> begin (k, s) end
                        | _ -> begin (k, Yojson.Safe.to_string v) end
                      ) args
                    end
                  | _ -> begin
                      []
                    end
                  in
                  
                  (* Find the prompt *)
                  let prompt_opt = List.find_opt (fun p -> p.Mcp_sdk.Prompt.name = name) server.prompts in
                  match prompt_opt with
                  | Some prompt -> begin
                      Log.debug (Printf.sprintf "Found prompt: %s" name);
                      let ctx = Mcp_sdk.Context.create 
                          ?request_id:(Some req.id) 
                          ~lifespan_context:server.lifespan_context 
                          () 
                      in
                      
                      (* Call the prompt handler *)
                      match prompt.handler ctx arguments with
                      | Ok messages -> begin
                          Log.debug (Printf.sprintf "Prompt handler returned %d messages" (List.length messages));
                          
                          (* Important: We need to directly use yojson_of_message which preserves MIME types *)
                          let messages_json = List.map Prompt.yojson_of_message messages in
                          
                          (* Debug output *)
                          Log.debug (Printf.sprintf "Messages JSON: %s" (Yojson.Safe.to_string (`List messages_json)));
                          
                          (* Verify one message if available to check structure *)
                          if List.length messages > 0 then begin
                            let first_msg = List.hd messages in
                            let content_debug = match first_msg.content with
                            | Text t -> begin
                                Printf.sprintf "Text content: %s" t.text
                              end
                            | Image i -> begin
                                Printf.sprintf "Image content (mime: %s)" i.mime_type
                              end
                            | Audio a -> begin
                                Printf.sprintf "Audio content (mime: %s)" a.mime_type
                              end
                            | Resource r -> begin
                                "Resource content"
                              end
                            in
                            Log.debug (Printf.sprintf "First message content type: %s" content_debug);
                          end;
                          
                          let result = `Assoc [
                            ("messages", `List messages_json);
                            ("description", match prompt.description with 
                              | Some d -> begin `String d end 
                              | None -> begin `Null end)
                          ] in
                          Some (create_response ~id:req.id ~result)
                        end
                      | Error err -> begin
                          Log.error (Printf.sprintf "Error processing prompt: %s" err);
                          Some (create_error ~id:req.id ~code:ErrorCode.internal_error ~message:err ())
                        end
                    end
                  | None -> begin
                      Log.error (Printf.sprintf "Prompt not found: %s" name);
                      Some (create_error ~id:req.id ~code:ErrorCode.invalid_params ~message:(Printf.sprintf "Prompt not found: %s" name) ())
                    end
                end
              | _ -> begin
                  Log.error "Invalid params format for prompts/get";
                  Some (create_error ~id:req.id ~code:ErrorCode.invalid_params ~message:"Invalid params format" ())
                end
            end else begin
              Some (create_error ~id:req.id ~code:ErrorCode.method_not_found ~message:"Prompts not supported" ())
            end
          end
           
        | "ping" -> begin
            Log.debug "Processing ping request";
            Some (create_response ~id:req.id ~result:(`Assoc []))
          end
         
        | _ -> begin
            Log.error (Printf.sprintf "Unknown method received: %s" req.method_);
            Some (create_error ~id:req.id ~code:ErrorCode.method_not_found ~message:("Method not found: " ^ req.method_) ())
          end
      end
         
    | JSONRPCMessage.Notification notif -> begin
        Log.debug (Printf.sprintf "Received notification with method: %s" notif.method_);
        match notif.method_ with
        | "notifications/initialized" -> begin
            Log.debug "Client initialization complete - Server is now ready to receive requests";
            None
          end
        | _ -> begin
            Log.debug (Printf.sprintf "Ignoring notification: %s" notif.method_);
            None
          end
      end
         
    | JSONRPCMessage.Response _ -> begin
        Log.error "Unexpected response message received";
        None
      end
      
    | JSONRPCMessage.Error _ -> begin
        Log.error "Unexpected error message received";
        None
      end
  with
  | Failure msg -> begin
      Log.error (Printf.sprintf "JSON error in message processing: %s" msg);
      None
    end
  | exc -> begin
      Log.error (Printf.sprintf "Exception during message processing: %s" (Printexc.to_string exc));
      Log.error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()));
      None
    end

(** Read a single message from stdin *)
let read_stdio_message () =
  try
    Log.debug "Reading line from stdin...";
    let line = read_line () in
    if line = "" then begin
      Log.debug "Empty line received, ignoring";
      None
    end else begin
      Log.debug (Printf.sprintf "Raw input: %s" (String.sub line 0 (min 100 (String.length line))));
      try
        let json = Yojson.Safe.from_string line in
        Log.debug "Successfully parsed JSON";
        Some json
      with
      | Yojson.Json_error msg -> begin
          Log.error (Printf.sprintf "Error parsing JSON: %s" msg);
          Log.error (Printf.sprintf "Input was: %s" (String.sub line 0 (min 100 (String.length line))));
          None
        end
    end
  with
  | End_of_file -> begin
      Log.debug "End of file received on stdin";
      None
    end
  | Sys_error msg -> begin
      Log.error (Printf.sprintf "System error while reading: %s" msg);
      None
    end
  | exc -> begin
      Log.error (Printf.sprintf "Exception while reading: %s" (Printexc.to_string exc));
      None
    end

(** Run stdio server with enhanced error handling *)
let rec run_stdio_server mcp_server =
  try begin
    if not mcp_server.running then begin
      Log.debug "Server stopped";
      ()
    end else begin
      match read_stdio_message () with
      | Some json -> begin
          Log.debug "Processing message...";
          try begin
            match process_message mcp_server.server json with
            | Some response -> begin
                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" 
                            (String.sub response_str 0 (min 100 (String.length response_str))));
                Printf.printf "%s\n" response_str;
                flush stdout;
                (* Give client time to process *)
                Unix.sleepf 0.01;
              end
            | None -> begin
                Log.debug "No response needed"
              end
            end with
            | exc -> begin
                Log.error (Printf.sprintf "ERROR in message processing: %s" (Printexc.to_string exc));
                Log.error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()));
                (* Try to extract ID and send an error response *)
                try begin
                  let id_opt = match Yojson.Safe.Util.member "id" json with
                    | `Int i -> Some (`Int i)
                    | `String s -> Some (`String s)
                    | _ -> None
                  in
                  match id_opt with
                  | Some id -> begin
                      let error_resp = create_error ~id ~code:ErrorCode.internal_error ~message:(Printexc.to_string exc) () in
                      let error_json = JSONRPCMessage.yojson_of_t error_resp in
                      let error_str = Yojson.Safe.to_string error_json in
                      Printf.printf "%s\n" error_str;
                      flush stdout;
                    end
                  | None -> begin
                      Log.error "Could not extract request ID to send error response"
                    end
                  end with
                  | e -> begin
                      Log.error (Printf.sprintf "Failed to send error response: %s" (Printexc.to_string e))
                    end
              end;
          run_stdio_server mcp_server
        end
      | None -> begin
          if mcp_server.running then begin
            (* No message received, but server is still running *)
            Unix.sleepf 0.1;  (* Small sleep to prevent CPU spinning *)
            run_stdio_server mcp_server
          end else begin
            Log.debug "Server stopped during message processing"
          end
        end
    end
  end with
  | exc -> begin
      Log.error (Printf.sprintf "FATAL ERROR in server main loop: %s" (Printexc.to_string exc));
      Log.error (Printf.sprintf "Backtrace: %s" (Printexc.get_backtrace()));
      (* Try to continue anyway *)
      if mcp_server.running then begin
        Unix.sleepf 0.1;
        run_stdio_server mcp_server
      end
    end

(** Create an MCP server *)
let create ~server ~transport () =
  { server; transport; running = false }

(** HTTP server placeholder (to be fully implemented) *)
let run_http_server mcp_server port =
  Log.info (Printf.sprintf "%s HTTP server starting on port %d" mcp_server.server.name port);
  Log.info "HTTP transport is a placeholder and not fully implemented yet";
  
  (* This would be where we'd set up cohttp server *)
  (*
  let callback _conn req body =
    let uri = req |> Cohttp.Request.uri in
    let meth = req |> Cohttp.Request.meth |> Cohttp.Code.string_of_method in
    
    (* Handle only POST /jsonrpc endpoint *)
    match (meth, Uri.path uri) with
    | "POST", "/jsonrpc" ->
      (* Read the body *)
      Cohttp_lwt.Body.to_string body >>= fun body_str ->
      
      (* Parse JSON *)
      let json = try Some (Yojson.Safe.from_string body_str) with _ -> None in
      match json with
      | Some json_msg ->
        (* Process the message *)
        let response_opt = process_message mcp_server.server json_msg in
        (match response_opt with
        | Some response ->
          let response_json = JSONRPCMessage.yojson_of_t response in
          let response_str = Yojson.Safe.to_string response_json in
          Cohttp_lwt_unix.Server.respond_string
            ~status:`OK
            ~body:response_str
            ~headers:(Cohttp.Header.init_with "Content-Type" "application/json")
            ()
        | None ->
          Cohttp_lwt_unix.Server.respond_string
            ~status:`OK
            ~body:"{}"
            ~headers:(Cohttp.Header.init_with "Content-Type" "application/json")
            ())
      | None ->
        Cohttp_lwt_unix.Server.respond_string
          ~status:`Bad_request
          ~body:"{\"error\":\"Invalid JSON\"}"
          ~headers:(Cohttp.Header.init_with "Content-Type" "application/json")
          ()
    | _ ->
      (* Return 404 for any other routes *)
      Cohttp_lwt_unix.Server.respond_string
        ~status:`Not_found
        ~body:"Not found"
        ()
  in
  
  (* Create and start the server *)
  let server = Cohttp_lwt_unix.Server.create
    ~mode:(`TCP (`Port port))
    (Cohttp_lwt_unix.Server.make ~callback ())
  in
  
  (* Run the server *)
  Lwt_main.run server
  *)
  
  (* For now, just wait until the server is stopped *)
  while mcp_server.running do
    Unix.sleep 1
  done

(** Start the server based on transport type *)
let start server =
  server.running <- true;
  
  (* Run startup hook if provided *)
  (match server.server.startup_hook with
  | Some hook -> begin hook () end
  | None -> begin () end);
  
  (* Install signal handler *)
  Sys.(set_signal sigint (Signal_handle (fun _ -> 
    Log.debug "Received interrupt signal, stopping server...";
    server.running <- false
  )));
  
  match server.transport with
  | Stdio -> begin
      (* Setup stdout and stderr *)
      set_binary_mode_out stdout false;
      Log.info (Printf.sprintf "%s server started with stdio transport" server.server.name);
      
      (* Run the server loop *)
      run_stdio_server server
    end
  | Http -> begin
      (* HTTP server placeholder *)
      run_http_server server 8080
    end

(** Stop the server *)
let stop server =
  Log.info "Stopping server...";
  server.running <- false;
  
  (* Run shutdown hook if provided *)
  match server.server.shutdown_hook with
  | Some hook -> begin hook () end
  | None -> begin () end