(** Example MCP server using the Server_session API.

    This example demonstrates how to create a simple MCP server that provides:
    - A few example tools (add, echo)
    - A simple resource (example://greeting)
    - Basic ping support

    Note: This is a template/demonstration. A real MCP server would typically:
    1. Read JSON-RPC messages from stdin
    2. Write JSON-RPC responses to stdout
    3. Use proper error handling and logging

    To use this as a real server, you would need to:
    - Create a custom Transport that reads/writes to stdio
    - Handle process lifecycle properly
    - Add comprehensive error handling
*)

open Mcp

(* Helper to find field in Jsont.Object *)
let find_field name fields =
  List.find_map (fun ((n, _), value) ->
    if n = name then Some value else None
  ) fields

(* Example tool: Add two numbers *)
let handle_add_tool ~name:_ ~arguments =
  let open Messages.Tools in
  match arguments with
  | Some (Jsont.Object (fields, _)) ->
      let a = match find_field "a" fields with
        | Some (Jsont.Number (n, _)) -> n
        | _ -> 0.0
      in
      let b = match find_field "b" fields with
        | Some (Jsont.Number (n, _)) -> n
        | _ -> 0.0
      in
      let result = a +. b in
      let content = [
        Content.text (Printf.sprintf "Result: %.2f" result)
      ] in
      make_call_result ~content ()
  | _ ->
      make_call_result
        ~content:[Content.text "Invalid arguments"]
        ~is_error:true
        ()

(* Example tool: Echo a message *)
let handle_echo_tool ~name:_ ~arguments =
  let open Messages.Tools in
  match arguments with
  | Some (Jsont.Object (fields, _)) ->
      (match find_field "message" fields with
       | Some (Jsont.String (msg, _)) ->
           let content = [Content.text msg] in
           make_call_result ~content ()
       | _ ->
           make_call_result
             ~content:[Content.text "No message provided"]
             ~is_error:true
             ())
  | _ ->
      make_call_result
        ~content:[Content.text "Invalid arguments"]
        ~is_error:true
        ()

(* Tool dispatcher *)
let call_tool ~name ~arguments =
  match name with
  | "add" -> handle_add_tool ~name ~arguments
  | "echo" -> handle_echo_tool ~name ~arguments
  | _ ->
      Messages.Tools.make_call_result
        ~content:[Content.text (Printf.sprintf "Unknown tool: %s" name)]
        ~is_error:true
        ()

(* List available tools *)
let list_tools ~cursor:_ =
  let open Messages.Tools in
  let meta = Jsont.Meta.none in
  let tools = [
    make_tool
      ~name:"add"
      ~description:"Add two numbers together"
      ~input_schema:(Jsont.Object ([
        (("type", meta), Jsont.String ("object", meta));
        (("properties", meta), Jsont.Object ([
          (("a", meta), Jsont.Object ([(("type", meta), Jsont.String ("number", meta))], meta));
          (("b", meta), Jsont.Object ([(("type", meta), Jsont.String ("number", meta))], meta));
        ], meta));
        (("required", meta), Jsont.Array ([Jsont.String ("a", meta); Jsont.String ("b", meta)], meta));
      ], meta))
      ();
    make_tool
      ~name:"echo"
      ~description:"Echo a message back"
      ~input_schema:(Jsont.Object ([
        (("type", meta), Jsont.String ("object", meta));
        (("properties", meta), Jsont.Object ([
          (("message", meta), Jsont.Object ([(("type", meta), Jsont.String ("string", meta))], meta));
        ], meta));
        (("required", meta), Jsont.Array ([Jsont.String ("message", meta)], meta));
      ], meta))
      ();
  ] in
  make_list_result ~tools ()

(* Example resource *)
let read_resource ~uri =
  let open Messages.Resources in
  match uri with
  | "example://greeting" ->
      let contents = [
        make_text_contents
          ~uri
          ~text:"Hello from the MCP server!"
          ~mime_type:"text/plain"
          ()
      ] in
      make_read_result ~contents
  | _ ->
      failwith (Printf.sprintf "Unknown resource: %s" uri)

let list_resources ~cursor:_ =
  let open Messages.Resources in
  let resources = [
    make_resource
      ~uri:"example://greeting"
      ~name:"Greeting"
      ~description:"A simple greeting message"
      ~mime_type:"text/plain"
      ()
  ] in
  make_list_result ~resources ()

(* Ping handler *)
let ping () = ()

(* Main server function *)
let run_server env =
  let open Eio in

  Switch.run @@ fun sw ->
    (* Create stdio transport for the server *)
    (* Note: For a real server, you would typically use a transport that
       reads from stdin and writes to stdout. Here we create a simple
       in-process transport for demonstration. *)
    let params = {
      Transport_stdio.command = "cat";  (* Echo back for demo *)
      args = [];
      env = None;
      max_buffer_size = None;
    } in
    let transport = Transport_stdio.create
      ~sw
      ~process_mgr:(Eio.Stdenv.process_mgr env)
      params
    in

    (* Configure server *)
    let config = {
      Server_session.server_info = Capabilities.Implementation.make
        ~name:"example-mcp-server"
        ~version:"1.0.0";
      server_capabilities = Capabilities.Server.make
        ~tools:(Capabilities.Tools.make ())
        ~resources:(Capabilities.Resources.make ())
        ()
      ;
      instructions = Some "Example MCP server with basic tools and resources";
    } in

    (* Set up handlers *)
    let handlers = {
      Server_session.list_resources = Some list_resources;
      list_resource_templates = None;
      read_resource = Some read_resource;
      subscribe_resource = None;
      unsubscribe_resource = None;
      list_tools = Some list_tools;
      call_tool = Some call_tool;
      list_prompts = None;
      get_prompt = None;
      complete = None;
      ping = Some ping;
    } in

    (* Create and run server *)
    let server = Server_session.create
      ~sw
      ~transport
      config
      handlers
    in

    (* Log server startup *)
    let client_info = Server_session.client_info server in
    Eio.traceln "MCP server started, connected to client: %s v%s"
      client_info.Capabilities.Implementation.name
      client_info.Capabilities.Implementation.version;

    (* The server runs in the background via the Session's receive loop *)
    (* The switch will keep the server alive until it's explicitly closed or an error occurs *)
    (* For this example, we just let it run until the process is terminated *)
    ()
    (* Note: In a real server, you would wait for some termination signal or condition *)

let () =
  Eio_main.run @@ fun env ->
    try
      run_server env
    with
    | exn ->
        Printf.eprintf "Server error: %s\n%!" (Printexc.to_string exn);
        exit 1