open Mcp
open Jsonrpc

(* SDK version *)
let version = "0.1.0"

(* Logging utilities *)
module Log = struct
  type level = Debug | Info | Warning | Error

  let string_of_level = function
    | Debug -> "DEBUG"
    | Info -> "INFO"
    | Warning -> "WARNING"
    | Error -> "ERROR"

  let log level msg =
    Printf.eprintf "[%s] %s\n" (string_of_level level) msg;
    flush stderr

  let debug = log Debug
  let info = log Info
  let warning = log Warning
  let error = log Error
end

(* Context for tools and resources *)
module Context = struct
  type t = {
    request_id: RequestId.t option;
    lifespan_context: (string * Json.t) list;
    progress_token: ProgressToken.t option;
  }

  let create ?request_id ?progress_token ?(lifespan_context=[]) () =
    { request_id; lifespan_context; progress_token }

  let get_context_value ctx key =
    List.assoc_opt key ctx.lifespan_context
  
  let report_progress ctx value total =
    match ctx.progress_token, ctx.request_id with
    | Some token, Some _id ->
      let params = `Assoc [
        ("progress", `Float value);
        ("total", `Float total);
        ("progressToken", ProgressToken.yojson_of_t token)
      ] in
      Some (create_notification ~method_:"notifications/progress" ~params:(Some params) ())
    | _ -> None
end

(* Tools for the MCP server *)
module Tool = struct
  type handler = Context.t -> Json.t -> (Json.t, string) result

  type t = {
    name: string;
    description: string option;
    input_schema: Json.t;  (* JSON Schema *)
    handler: handler;
  }

  let create ~name ?description ~input_schema ~handler () = 
    { name; description; input_schema; handler }

  let to_json tool =
    let assoc = [
      ("name", `String tool.name);
      ("inputSchema", tool.input_schema);
    ] in
    let assoc = match tool.description with
      | Some desc -> ("description", `String desc) :: assoc
      | None -> assoc
    in
    `Assoc assoc
end

(* Resources for the MCP server *)
module Resource = struct
  type handler = Context.t -> string list -> (string, string) result

  type t = {
    uri_template: string;
    description: string option;
    mime_type: string option;
    handler: handler;
  }

  let create ~uri_template ?description ?mime_type ~handler () = 
    { uri_template; description; mime_type; handler }

  let to_json resource =
    let assoc = [
      ("uriTemplate", `String resource.uri_template);
    ] in
    let assoc = match resource.description with
      | Some desc -> ("description", `String desc) :: assoc
      | None -> assoc
    in
    let assoc = match resource.mime_type with
      | Some mime -> ("mimeType", `String mime) :: assoc
      | None -> assoc
    in
    `Assoc assoc
end

(* Prompts for the MCP server *)
module Prompt = struct
  type argument = {
    name: string;
    description: string option;
    required: bool;
  }

  type message = {
    role: Role.t;
    content: content;
  }

  type handler = Context.t -> (string * string) list -> (message list, string) result

  type t = {
    name: string;
    description: string option;
    arguments: argument list;
    handler: handler;
  }

  let create ~name ?description ?(arguments=[]) ~handler () = 
    { name; description; arguments; handler }

  let create_argument ~name ?description ?(required=false) () =
    { name; description; required }

  let to_json prompt =
    let assoc = [
      ("name", `String prompt.name);
    ] in
    let assoc = match prompt.description with
      | Some desc -> ("description", `String desc) :: assoc
      | None -> assoc
    in
    let assoc = if prompt.arguments <> [] then
      let args = List.map (fun (arg: argument) ->
        let arg_assoc = [
          ("name", `String arg.name);
        ] in
        let arg_assoc = match arg.description with
          | Some desc -> ("description", `String desc) :: arg_assoc
          | None -> arg_assoc
        in
        let arg_assoc = 
          if arg.required then
            ("required", `Bool true) :: arg_assoc
          else
            arg_assoc
        in
        `Assoc arg_assoc
      ) prompt.arguments in
      ("arguments", `List args) :: assoc
    else
      assoc
    in
    `Assoc assoc
end

(* Helper functions for creating common objects *)
(* Content type constructors *)
let make_text_content text =
  Text (TextContent.{ text; annotations = None })

let make_image_content data mime_type =
  Image (ImageContent.{ data; mime_type; annotations = None })

(* Create audio content using the AudioContent module *)
let make_audio_content data mime_type =
  Audio (AudioContent.{ data; mime_type; annotations = None })

let make_resource_text_content uri text mime_type =
  Resource (EmbeddedResource.{
    resource = `Text TextResourceContents.{ uri; text; mime_type };
    annotations = None;
  })

let make_resource_blob_content uri blob mime_type =
  Resource (EmbeddedResource.{
    resource = `Blob BlobResourceContents.{ uri; blob; mime_type };
    annotations = None;
  })

(* Tool result handling *)
type tool_content = 
  | TextContent of string
  | ImageContent of { data: string; mime_type: string }
  | AudioContent of { data: string; mime_type: string }
  | ResourceContent of { uri: string; data: string; is_blob: bool; mime_type: string option }

let create_tool_result contents ~is_error =
  let content_json = List.map (function
    | TextContent text -> 
        `Assoc [
          ("type", `String "text");
          ("text", `String text)
        ]
    | ImageContent { data; mime_type } ->
        `Assoc [
          ("type", `String "image");
          ("data", `String data);
          ("mimeType", `String mime_type)
        ]
    | AudioContent { data; mime_type } ->
        `Assoc [
          ("type", `String "audio");
          ("data", `String data);
          ("mimeType", `String mime_type)
        ]
    | ResourceContent { uri; data; is_blob; mime_type } ->
        let resource_data = if is_blob then
          `Assoc (
            [("uri", `String uri); 
             ("blob", `String data)] @
            (match mime_type with 
             | Some mime -> [("mimeType", `String mime)]
             | None -> [])
          )
        else
          `Assoc (
            [("uri", `String uri); 
             ("text", `String data)] @
            (match mime_type with 
             | Some mime -> [("mimeType", `String mime)]
             | None -> [])
          )
        in
        `Assoc [
          ("type", `String "resource");
          ("resource", resource_data)
        ]
  ) contents in
  
  `Assoc [
    ("content", `List content_json);
    ("isError", `Bool is_error)
  ]

(* Error types with standard JSON-RPC error codes *)
type error_code = 
  | ParseError (* -32700 *)
  | InvalidRequest (* -32600 *)
  | MethodNotFound (* -32601 *)
  | InvalidParams (* -32602 *)
  | InternalError (* -32603 *)
  | ResourceNotFound (* -32002 Custom code for MCP *)
  | AuthenticationRequired (* -32001 Custom code for MCP *)
  | CustomError of int

let error_code_to_int = function
  | ParseError -> -32700
  | InvalidRequest -> -32600
  | MethodNotFound -> -32601
  | InvalidParams -> -32602
  | InternalError -> -32603
  | ResourceNotFound -> -32002
  | AuthenticationRequired -> -32001
  | CustomError code -> code

let make_tool_schema properties required =
  let props = List.map (fun (name, schema_type, description) ->
    (name, `Assoc [
      ("type", `String schema_type);
      ("description", `String description)
    ])
  ) properties in
  let required_json = `List (List.map (fun name -> `String name) required) in
  `Assoc [
    ("type", `String "object");
    ("properties", `Assoc props);
    ("required", required_json)
  ]

(* Main server type *)
type server = {
  name: string;
  version: string;
  protocol_version: string;
  lifespan_context: (string * Json.t) list;
  mutable capabilities: Json.t;
  mutable tools: Tool.t list;
  mutable resources: Resource.t list;
  mutable prompts: Prompt.t list;
} 

let name { name; _ } = name
let version { version; _ } = version
let capabilities { capabilities; _ } = capabilities
let lifespan_context { lifespan_context; _ } = lifespan_context
let protocol_version { protocol_version; _ } = protocol_version
let tools { tools; _ } = tools
let resources { resources; _ } = resources
let prompts { prompts; _ } = prompts

(* Create a new server *)
let create_server ~name ?(version="0.1.0") ?(protocol_version="2024-11-05") () =
  { 
    name;
    version;
    protocol_version;
    capabilities = `Assoc [];
    tools = [];
    resources = [];
    prompts = [];
    lifespan_context = [];
  }

(* Default capabilities for the server *)
let default_capabilities ?(with_tools=true) ?(with_resources=false) ?(with_prompts=false) () =
  let caps = [] in
  let caps = 
    if with_tools then
      ("tools", `Assoc [
        ("listChanged", `Bool true)
      ]) :: caps
    else
      caps
  in
  let caps = 
    if with_resources then
      ("resources", `Assoc [
        ("listChanged", `Bool true);
        ("subscribe", `Bool false)
      ]) :: caps
    else if not with_resources then
      ("resources", `Assoc [
        ("listChanged", `Bool false);
        ("subscribe", `Bool false)
      ]) :: caps
    else
      caps
  in
  let caps = 
    if with_prompts then
      ("prompts", `Assoc [
        ("listChanged", `Bool true)
      ]) :: caps
    else if not with_prompts then
      ("prompts", `Assoc [
        ("listChanged", `Bool false)
      ]) :: caps
    else
      caps
  in
  `Assoc caps

(* Register a tool *)
let register_tool server tool =
  server.tools <- tool :: server.tools;
  tool

(* Create a rich tool result with multiple content types *)
let create_rich_tool_result ?(text=None) ?(image=None) ?(audio=None) ?(resource=None) ~is_error () =
  let contents = [] in
  
  (* Add text content if provided *)
  let contents = match text with
    | Some text -> (TextContent text) :: contents
    | None -> contents
  in
  
  (* Add image content if provided *)
  let contents = match image with
    | Some (data, mime_type) -> (ImageContent { data; mime_type }) :: contents
    | None -> contents
  in
  
  (* Add audio content if provided *)
  let contents = match audio with
    | Some (data, mime_type) -> (AudioContent { data; mime_type }) :: contents
    | None -> contents
  in
  
  (* Add resource content if provided *)
  let contents = match resource with
    | Some (uri, data, is_blob, mime_type) -> 
        (ResourceContent { uri; data; is_blob; mime_type }) :: contents
    | None -> contents
  in
  
  (* Create the final tool result *)
  create_tool_result (List.rev contents) ~is_error

(* Create and register a tool in one step *)
let add_tool server ~name ?description ?(schema_properties=[]) ?(schema_required=[]) handler =
  let input_schema = make_tool_schema schema_properties schema_required in
  let handler' ctx args =
    try
      Ok (handler args)
    with exn ->
      Error (Printexc.to_string exn)
  in
  let tool = Tool.create
    ~name
    ?description
    ~input_schema
    ~handler:handler'
    ()
  in
  register_tool server tool

(* Register a resource *)
let register_resource server resource =
  server.resources <- resource :: server.resources;
  resource

(* Create and register a resource in one step *)
let add_resource server ~uri_template ?description ?mime_type handler =
  let handler' _ctx params =
    try
      Ok (handler params)
    with exn ->
      Error (Printexc.to_string exn)
  in
  let resource = Resource.create
    ~uri_template
    ?description
    ?mime_type
    ~handler:handler'
    ()
  in
  register_resource server resource

(* Register a prompt *)
let register_prompt server prompt =
  server.prompts <- prompt :: server.prompts;
  prompt

(* Create and register a prompt in one step *)
let add_prompt server ~name ?description ?(arguments=[]) handler =
  let prompt_args = List.map (fun (name, desc, required) ->
    Prompt.create_argument ~name ?description:desc ~required ()
  ) arguments in
  let handler' _ctx args =
    try
      Ok (handler args)
    with exn ->
      Error (Printexc.to_string exn)
  in
  let prompt = Prompt.create
    ~name
    ?description
    ~arguments:prompt_args
    ~handler:handler'
    ()
  in
  register_prompt server prompt

(* Set server capabilities *)
let set_capabilities server capabilities =
  server.capabilities <- capabilities

(* Configure server with default capabilities based on registered components *)
let configure_server server ?with_tools ?with_resources ?with_prompts () =
  let with_tools = match with_tools with
    | Some b -> b
    | None -> server.tools <> []
  in
  let with_resources = match with_resources with
    | Some b -> b
    | None -> server.resources <> []
  in
  let with_prompts = match with_prompts with
    | Some b -> b
    | None -> server.prompts <> []
  in
  let capabilities = default_capabilities ~with_tools ~with_resources ~with_prompts () in
  set_capabilities server capabilities;
  server

(* MCP Protocol Message Helpers *)

(* Resources helpers *)
let create_resources_list_request ?cursor ?id () =
  Mcp_message.ResourcesList.create_request ?cursor ?id ()

let create_resources_list_response ~id ~resources ?next_cursor () =
  Mcp_message.ResourcesList.create_response ~id ~resources ?next_cursor ()

let create_resources_read_request ~uri ?id () =
  Mcp_message.ResourcesRead.create_request ~uri ?id ()

let create_resources_read_response ~id ~contents () =
  Mcp_message.ResourcesRead.create_response ~id ~contents ()

let create_resources_list_changed_notification () =
  Mcp_message.ListChanged.create_resources_notification ()

let create_resources_updated_notification ~uri () =
  Mcp_message.ResourceUpdated.create_notification ~uri ()

(* Tools helpers *)
let create_tools_list_request ?cursor ?id () =
  Mcp_message.ToolsList.create_request ?cursor ?id ()

let create_tools_list_response ~id ~tools ?next_cursor () =
  Mcp_message.ToolsList.create_response ~id ~tools ?next_cursor ()

let create_tools_call_request ~name ~arguments ?id () =
  Mcp_message.ToolsCall.create_request ~name ~arguments ?id ()

let create_tools_call_response ~id ~content ~is_error () =
  Mcp_message.ToolsCall.create_response ~id ~content ~is_error ()

let create_tools_list_changed_notification () =
  Mcp_message.ListChanged.create_tools_notification ()

(* Prompts helpers *)
let create_prompts_list_request ?cursor ?id () =
  Mcp_message.PromptsList.create_request ?cursor ?id ()

let create_prompts_list_response ~id ~prompts ?next_cursor () =
  Mcp_message.PromptsList.create_response ~id ~prompts ?next_cursor ()

let create_prompts_get_request ~name ~arguments ?id () =
  Mcp_message.PromptsGet.create_request ~name ~arguments ?id ()

let create_prompts_get_response ~id ?description ~messages () =
  Mcp_message.PromptsGet.create_response ~id ?description ~messages ()

let create_prompts_list_changed_notification () =
  Mcp_message.ListChanged.create_prompts_notification ()

(* Progress notification *)
let create_progress_notification ~progress ~total ~progress_token () =
  Mcp_message.Progress.create_notification ~progress ~total ~progress_token ()