open Jsonrpc

(* Standard error codes *)
module ErrorCode = struct
  let parse_error = -32700
  let invalid_request = -32600
  let method_not_found = -32601
  let invalid_params = -32602
  let internal_error = -32603
  let resource_not_found = -32002
  let server_error_start = -32000
  let server_error_end = -32099
end

(* Common types *)

module Role = struct
  type t = [ `User | `Assistant ]

  let to_string = function
    | `User -> "user"
    | `Assistant -> "assistant"

  let of_string = function
    | "user" -> `User
    | "assistant" -> `Assistant
    | s -> raise (Json.Of_json ("Unknown role: " ^ s, `String s))

  let yojson_of_t t = `String (to_string t)
  let t_of_yojson = function
    | `String s -> of_string s
    | j -> raise (Json.Of_json ("Expected string for Role", j))
end

module ProgressToken = struct
  type t = [ `String of string | `Int of int ]

  include (Id : Json.Jsonable.S with type t := t)
end

module RequestId = Id

module Cursor = struct
  type t = string

  let yojson_of_t t = `String t
  let t_of_yojson = function
    | `String s -> s
    | j -> raise (Json.Of_json ("Expected string for Cursor", j))
end

(* Annotations *)

module Annotated = struct
  type t = {
    annotations: annotation option;
  }
  and annotation = {
    audience: Role.t list option;
    priority: float option;
  }

  let yojson_of_annotation { audience; priority } =
    let assoc = [] in
    let assoc = match audience with
      | Some audience -> ("audience", `List (List.map Role.yojson_of_t audience)) :: assoc
      | None -> assoc
    in
    let assoc = match priority with
      | Some priority -> ("priority", `Float priority) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let annotation_of_yojson = function
    | `Assoc fields ->
      let audience = List.assoc_opt "audience" fields |> Option.map (function
        | `List items -> List.map Role.t_of_yojson items
        | j -> raise (Json.Of_json ("Expected list for audience", j))
      ) in
      let priority = List.assoc_opt "priority" fields |> Option.map (function
        | `Float f -> f
        | j -> raise (Json.Of_json ("Expected float for priority", j))
      ) in
      { audience; priority }
    | j -> raise (Json.Of_json ("Expected object for annotation", j))

  let yojson_of_t { annotations } =
    match annotations with
    | Some annotations -> `Assoc [ "annotations", yojson_of_annotation annotations ]
    | None -> `Assoc []

  let t_of_yojson = function
    | `Assoc fields ->
      let annotations = List.assoc_opt "annotations" fields |> Option.map annotation_of_yojson in
      { annotations }
    | j -> raise (Json.Of_json ("Expected object for Annotated", j))
end

(* Content types *)

module TextContent = struct
  type t = {
    text: string;
    annotations: Annotated.annotation option;
  }

  let yojson_of_t { text; annotations } =
    let assoc = [
      ("text", `String text);
      ("type", `String "text");
    ] in
    let assoc = match annotations with
      | Some annotations -> ("annotations", Annotated.yojson_of_annotation annotations) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let text = match List.assoc_opt "text" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'text' field", `Assoc fields))
      in
      let _ = match List.assoc_opt "type" fields with
        | Some (`String "text") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'type' field", `Assoc fields))
      in
      let annotations = List.assoc_opt "annotations" fields |> Option.map Annotated.annotation_of_yojson in
      { text; annotations }
    | j -> raise (Json.Of_json ("Expected object for TextContent", j))
end

module ImageContent = struct
  type t = {
    data: string;
    mime_type: string;
    annotations: Annotated.annotation option;
  }

  let yojson_of_t { data; mime_type; annotations } =
    let assoc = [
      ("data", `String data);
      ("mimeType", `String mime_type);
      ("type", `String "image");
    ] in
    let assoc = match annotations with
      | Some annotations -> ("annotations", Annotated.yojson_of_annotation annotations) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let data = match List.assoc_opt "data" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'data' field", `Assoc fields))
      in
      let mime_type = match List.assoc_opt "mimeType" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'mimeType' field", `Assoc fields))
      in
      let _ = match List.assoc_opt "type" fields with
        | Some (`String "image") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'type' field", `Assoc fields))
      in
      let annotations = List.assoc_opt "annotations" fields |> Option.map Annotated.annotation_of_yojson in
      { data; mime_type; annotations }
    | j -> raise (Json.Of_json ("Expected object for ImageContent", j))
end

module AudioContent = struct
  type t = {
    data: string;
    mime_type: string;
    annotations: Annotated.annotation option;
  }

  let yojson_of_t { data; mime_type; annotations } =
    let assoc = [
      ("data", `String data);
      ("mimeType", `String mime_type);
      ("type", `String "audio");
    ] in
    let assoc = match annotations with
      | Some annotations -> ("annotations", Annotated.yojson_of_annotation annotations) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let data = match List.assoc_opt "data" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'data' field", `Assoc fields))
      in
      let mime_type = match List.assoc_opt "mimeType" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'mimeType' field", `Assoc fields))
      in
      let _ = match List.assoc_opt "type" fields with
        | Some (`String "audio") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'type' field", `Assoc fields))
      in
      let annotations = List.assoc_opt "annotations" fields |> Option.map Annotated.annotation_of_yojson in
      { data; mime_type; annotations }
    | j -> raise (Json.Of_json ("Expected object for AudioContent", j))
end

module ResourceContents = struct
  type t = {
    uri: string;
    mime_type: string option;
  }

  let yojson_of_t { uri; mime_type } =
    let assoc = [
      ("uri", `String uri);
    ] in
    let assoc = match mime_type with
      | Some mime_type -> ("mimeType", `String mime_type) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let uri = match List.assoc_opt "uri" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'uri' field", `Assoc fields))
      in
      let mime_type = List.assoc_opt "mimeType" fields |> Option.map (function
        | `String s -> s
        | j -> raise (Json.Of_json ("Expected string for mimeType", j))
      ) in
      { uri; mime_type }
    | j -> raise (Json.Of_json ("Expected object for ResourceContents", j))
end

module TextResourceContents = struct
  type t = {
    uri: string;
    text: string;
    mime_type: string option;
  }

  let yojson_of_t { uri; text; mime_type } =
    let assoc = [
      ("uri", `String uri);
      ("text", `String text);
    ] in
    let assoc = match mime_type with
      | Some mime_type -> ("mimeType", `String mime_type) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let uri = match List.assoc_opt "uri" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'uri' field", `Assoc fields))
      in
      let text = match List.assoc_opt "text" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'text' field", `Assoc fields))
      in
      let mime_type = List.assoc_opt "mimeType" fields |> Option.map (function
        | `String s -> s
        | j -> raise (Json.Of_json ("Expected string for mimeType", j))
      ) in
      { uri; text; mime_type }
    | j -> raise (Json.Of_json ("Expected object for TextResourceContents", j))
end

module BlobResourceContents = struct
  type t = {
    uri: string;
    blob: string;
    mime_type: string option;
  }

  let yojson_of_t { uri; blob; mime_type } =
    let assoc = [
      ("uri", `String uri);
      ("blob", `String blob);
    ] in
    let assoc = match mime_type with
      | Some mime_type -> ("mimeType", `String mime_type) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let uri = match List.assoc_opt "uri" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'uri' field", `Assoc fields))
      in
      let blob = match List.assoc_opt "blob" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'blob' field", `Assoc fields))
      in
      let mime_type = List.assoc_opt "mimeType" fields |> Option.map (function
        | `String s -> s
        | j -> raise (Json.Of_json ("Expected string for mimeType", j))
      ) in
      { uri; blob; mime_type }
    | j -> raise (Json.Of_json ("Expected object for BlobResourceContents", j))
end

module EmbeddedResource = struct
  type t = {
    resource: [ `Text of TextResourceContents.t | `Blob of BlobResourceContents.t ];
    annotations: Annotated.annotation option;
  }

  let yojson_of_t { resource; annotations } =
    let resource_json = match resource with
      | `Text txt -> TextResourceContents.yojson_of_t txt
      | `Blob blob -> BlobResourceContents.yojson_of_t blob
    in
    let assoc = [
      ("resource", resource_json);
      ("type", `String "resource");
    ] in
    let assoc = match annotations with
      | Some annotations -> ("annotations", Annotated.yojson_of_annotation annotations) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let _ = match List.assoc_opt "type" fields with
        | Some (`String "resource") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'type' field", `Assoc fields))
      in
      let resource = match List.assoc_opt "resource" fields with
        | Some (`Assoc res_fields) ->
          if List.mem_assoc "text" res_fields then
            `Text (TextResourceContents.t_of_yojson (`Assoc res_fields))
          else if List.mem_assoc "blob" res_fields then
            `Blob (BlobResourceContents.t_of_yojson (`Assoc res_fields))
          else
            raise (Json.Of_json ("Invalid resource content", `Assoc res_fields))
        | _ -> raise (Json.Of_json ("Missing or invalid 'resource' field", `Assoc fields))
      in
      let annotations = List.assoc_opt "annotations" fields |> Option.map Annotated.annotation_of_yojson in
      { resource; annotations }
    | j -> raise (Json.Of_json ("Expected object for EmbeddedResource", j))
end

(** Tool definition *)
module Tool = struct
  type t = {
    name: string;
    description: string option;
    input_schema: Json.t;
  }

  let yojson_of_t { name; description; input_schema } =
    let assoc = [
      ("name", `String name);
      ("inputSchema", input_schema);
    ] in
    let assoc = match description with
      | Some desc -> ("description", `String desc) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let name = match List.assoc_opt "name" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'name' field", `Assoc fields))
      in
      let description = match List.assoc_opt "description" fields with
        | Some (`String s) -> Some s
        | _ -> None
      in
      let input_schema = match List.assoc_opt "inputSchema" fields with
        | Some json -> json
        | _ -> raise (Json.Of_json ("Missing 'inputSchema' field", `Assoc fields))
      in
      { name; description; input_schema }
    | j -> raise (Json.Of_json ("Expected object for Tool", j))
end

type content = 
  | Text of TextContent.t
  | Image of ImageContent.t
  | Audio of AudioContent.t
  | Resource of EmbeddedResource.t

let yojson_of_content = function
  | Text t -> TextContent.yojson_of_t t
  | Image i -> ImageContent.yojson_of_t i
  | Audio a -> AudioContent.yojson_of_t a
  | Resource r -> EmbeddedResource.yojson_of_t r

let content_of_yojson = function
  | `Assoc fields ->
    (match List.assoc_opt "type" fields with
     | Some (`String "text") -> Text (TextContent.t_of_yojson (`Assoc fields))
     | Some (`String "image") -> Image (ImageContent.t_of_yojson (`Assoc fields))
     | Some (`String "audio") -> Audio (AudioContent.t_of_yojson (`Assoc fields))
     | Some (`String "resource") -> Resource (EmbeddedResource.t_of_yojson (`Assoc fields))
     | _ -> raise (Json.Of_json ("Invalid or missing content type", `Assoc fields)))
  | j -> raise (Json.Of_json ("Expected object for content", j))

(** Tool result *)
module CallToolResult = struct
  type t = {
    content: content list;
    is_error: bool;
    meta: Json.t option;
  }

  let yojson_of_t { content; is_error; meta } =
    let assoc = [
      ("content", `List (List.map yojson_of_content content));
      ("isError", `Bool is_error);
    ] in
    let assoc = match meta with
      | Some meta_json -> ("_meta", meta_json) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let content = match List.assoc_opt "content" fields with
        | Some (`List items) -> List.map content_of_yojson items
        | _ -> raise (Json.Of_json ("Missing or invalid 'content' field", `Assoc fields))
      in
      let is_error = match List.assoc_opt "isError" fields with
        | Some (`Bool b) -> b
        | None -> false (* Default to false if not specified *)
        | _ -> raise (Json.Of_json ("Invalid 'isError' field", `Assoc fields))
      in
      let meta = List.assoc_opt "_meta" fields in
      { content; is_error; meta }
    | j -> raise (Json.Of_json ("Expected object for CallToolResult", j))
end

(** Resource definition *)
module Resource = struct
  type t = {
    name: string;
    uri: string;
    description: string option;
    mime_type: string option;
    size: int option;
    annotations: Annotated.annotation option;
  }

  let yojson_of_t { name; uri; description; mime_type; size; annotations } =
    let assoc = [
      ("name", `String name);
      ("uri", `String uri);
    ] in
    let assoc = match description with
      | Some desc -> ("description", `String desc) :: assoc
      | None -> assoc
    in
    let assoc = match mime_type with
      | Some mime -> ("mimeType", `String mime) :: assoc
      | None -> assoc
    in
    let assoc = match size with
      | Some s -> ("size", `Int s) :: assoc
      | None -> assoc
    in
    let assoc = match annotations with
      | Some ann -> ("annotations", Annotated.yojson_of_annotation ann) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let name = match List.assoc_opt "name" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'name' field", `Assoc fields))
      in
      let uri = match List.assoc_opt "uri" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'uri' field", `Assoc fields))
      in
      let description = match List.assoc_opt "description" fields with
        | Some (`String s) -> Some s
        | _ -> None
      in
      let mime_type = match List.assoc_opt "mimeType" fields with
        | Some (`String s) -> Some s
        | _ -> None
      in
      let size = match List.assoc_opt "size" fields with
        | Some (`Int s) -> Some s
        | _ -> None
      in
      let annotations = match List.assoc_opt "annotations" fields with
        | Some json -> Some (Annotated.annotation_of_yojson json)
        | _ -> None
      in
      { name; uri; description; mime_type; size; annotations }
    | j -> raise (Json.Of_json ("Expected object for Resource", j))
end

(** Resource Template definition *)
module ResourceTemplate = struct
  type t = {
    name: string;
    uri_template: string;
    description: string option;
    mime_type: string option;
    annotations: Annotated.annotation option;
  }

  let yojson_of_t { name; uri_template; description; mime_type; annotations } =
    let assoc = [
      ("name", `String name);
      ("uriTemplate", `String uri_template);
    ] in
    let assoc = match description with
      | Some desc -> ("description", `String desc) :: assoc
      | None -> assoc
    in
    let assoc = match mime_type with
      | Some mime -> ("mimeType", `String mime) :: assoc
      | None -> assoc
    in
    let assoc = match annotations with
      | Some ann -> ("annotations", Annotated.yojson_of_annotation ann) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let t_of_yojson = function
    | `Assoc fields ->
      let name = match List.assoc_opt "name" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'name' field", `Assoc fields))
      in
      let uri_template = match List.assoc_opt "uriTemplate" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'uriTemplate' field", `Assoc fields))
      in
      let description = match List.assoc_opt "description" fields with
        | Some (`String s) -> Some s
        | _ -> None
      in
      let mime_type = match List.assoc_opt "mimeType" fields with
        | Some (`String s) -> Some s
        | _ -> None
      in
      let annotations = match List.assoc_opt "annotations" fields with
        | Some json -> Some (Annotated.annotation_of_yojson json)
        | _ -> None
      in
      { name; uri_template; description; mime_type; annotations }
    | j -> raise (Json.Of_json ("Expected object for ResourceTemplate", j))
end

(** Resource Reference *)
module ResourceReference = struct
  type t = {
    uri: string;
  }

  let yojson_of_t { uri } =
    `Assoc [
      ("type", `String "ref/resource");
      ("uri", `String uri);
    ]

  let t_of_yojson = function
    | `Assoc fields ->
      let _ = match List.assoc_opt "type" fields with
        | Some (`String "ref/resource") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'type' field", `Assoc fields))
      in
      let uri = match List.assoc_opt "uri" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'uri' field", `Assoc fields))
      in
      { uri }
    | j -> raise (Json.Of_json ("Expected object for ResourceReference", j))
end

(** Prompt Reference *)
module PromptReference = struct
  type t = {
    name: string;
  }

  let yojson_of_t { name } =
    `Assoc [
      ("type", `String "ref/prompt");
      ("name", `String name);
    ]

  let t_of_yojson = function
    | `Assoc fields ->
      let _ = match List.assoc_opt "type" fields with
        | Some (`String "ref/prompt") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'type' field", `Assoc fields))
      in
      let name = match List.assoc_opt "name" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'name' field", `Assoc fields))
      in
      { name }
    | j -> raise (Json.Of_json ("Expected object for PromptReference", j))
end

(** Completion support *)
module Completion = struct
  
  module Argument = struct
    type t = {
      name: string;
      value: string;
    }

    let yojson_of_t { name; value } =
      `Assoc [
        ("name", `String name);
        ("value", `String value);
      ]

    let t_of_yojson = function
      | `Assoc fields ->
        let name = match List.assoc_opt "name" fields with
          | Some (`String s) -> s
          | _ -> raise (Json.Of_json ("Missing or invalid 'name' field", `Assoc fields))
        in
        let value = match List.assoc_opt "value" fields with
          | Some (`String s) -> s
          | _ -> raise (Json.Of_json ("Missing or invalid 'value' field", `Assoc fields))
        in
        { name; value }
      | j -> raise (Json.Of_json ("Expected object for Completion.Argument", j))
  end

  module Request = struct
    type reference = [ `Prompt of PromptReference.t | `Resource of ResourceReference.t ]

    type t = {
      argument: Argument.t;
      ref: reference;
    }

    let yojson_of_reference = function
      | `Prompt p -> PromptReference.yojson_of_t p
      | `Resource r -> ResourceReference.yojson_of_t r

    let reference_of_yojson = function
      | `Assoc fields ->
        (match List.assoc_opt "type" fields with
         | Some (`String "ref/prompt") -> `Prompt (PromptReference.t_of_yojson (`Assoc fields))
         | Some (`String "ref/resource") -> `Resource (ResourceReference.t_of_yojson (`Assoc fields))
         | _ -> raise (Json.Of_json ("Invalid or missing reference type", `Assoc fields)))
      | j -> raise (Json.Of_json ("Expected object for reference", j))

    let yojson_of_t { argument; ref } =
      `Assoc [
        ("argument", Argument.yojson_of_t argument);
        ("ref", yojson_of_reference ref);
      ]

    let t_of_yojson = function
      | `Assoc fields ->
        let argument = match List.assoc_opt "argument" fields with
          | Some json -> Argument.t_of_yojson json
          | _ -> raise (Json.Of_json ("Missing argument field", `Assoc fields))
        in
        let ref = match List.assoc_opt "ref" fields with
          | Some json -> reference_of_yojson json
          | _ -> raise (Json.Of_json ("Missing ref field", `Assoc fields))
        in
        { argument; ref }
      | j -> raise (Json.Of_json ("Expected object for Completion.Request", j))

    let create ~argument ~ref =
      { argument; ref }

    let to_params t =
      yojson_of_t t
  end

  module Result = struct
    type completion = {
      values: string list;
      has_more: bool option;
      total: int option;
    }

    type t = {
      completion: completion;
      meta: Json.t option;
    }

    let yojson_of_completion { values; has_more; total } =
      let assoc = [
        ("values", `List (List.map (fun s -> `String s) values));
      ] in
      let assoc = match has_more with
        | Some b -> ("hasMore", `Bool b) :: assoc
        | None -> assoc
      in
      let assoc = match total with
        | Some n -> ("total", `Int n) :: assoc
        | None -> assoc
      in
      `Assoc assoc

    let completion_of_yojson = function
      | `Assoc fields ->
        let values = match List.assoc_opt "values" fields with
          | Some (`List items) -> 
            List.map (function
              | `String s -> s
              | _ -> raise (Json.Of_json ("Expected string in values array", `List items))
            ) items
          | _ -> raise (Json.Of_json ("Missing or invalid 'values' field", `Assoc fields))
        in
        let has_more = match List.assoc_opt "hasMore" fields with
          | Some (`Bool b) -> Some b
          | None -> None
          | _ -> raise (Json.Of_json ("Invalid 'hasMore' field", `Assoc fields))
        in
        let total = match List.assoc_opt "total" fields with
          | Some (`Int n) -> Some n
          | None -> None
          | _ -> raise (Json.Of_json ("Invalid 'total' field", `Assoc fields))
        in
        { values; has_more; total }
      | j -> raise (Json.Of_json ("Expected object for completion", j))

    let yojson_of_t { completion; meta } =
      let assoc = [
        ("completion", yojson_of_completion completion);
      ] in
      let assoc = match meta with
        | Some meta_json -> ("_meta", meta_json) :: assoc
        | None -> assoc
      in
      `Assoc assoc

    let t_of_yojson = function
      | `Assoc fields ->
        let completion = match List.assoc_opt "completion" fields with
          | Some json -> completion_of_yojson json
          | _ -> raise (Json.Of_json ("Missing completion field", `Assoc fields))
        in
        let meta = List.assoc_opt "_meta" fields in
        { completion; meta }
      | j -> raise (Json.Of_json ("Expected object for Completion.Result", j))

    let create ~completion ?meta () =
      { completion; meta }

    let to_result t =
      yojson_of_t t
  end
end

(* Message types *)

module PromptMessage = struct
  type t = {
    role: Role.t;
    content: content;
  }

  let yojson_of_t { role; content } =
    `Assoc [
      ("role", Role.yojson_of_t role);
      ("content", yojson_of_content content);
    ]

  let t_of_yojson = function
    | `Assoc fields ->
      let role = match List.assoc_opt "role" fields with
        | Some json -> Role.t_of_yojson json
        | None -> raise (Json.Of_json ("Missing role field", `Assoc fields))
      in
      let content = match List.assoc_opt "content" fields with
        | Some json -> content_of_yojson json
        | None -> raise (Json.Of_json ("Missing content field", `Assoc fields))
      in
      { role; content }
    | j -> raise (Json.Of_json ("Expected object for PromptMessage", j))
end

module SamplingMessage = struct
  type t = {
    role: Role.t;
    content: [ `Text of TextContent.t | `Image of ImageContent.t ];
  }

  let yojson_of_t { role; content } =
    let content_json = match content with
      | `Text t -> TextContent.yojson_of_t t
      | `Image i -> ImageContent.yojson_of_t i
    in
    `Assoc [
      ("role", Role.yojson_of_t role);
      ("content", content_json);
    ]

  let t_of_yojson = function
    | `Assoc fields ->
      let role = match List.assoc_opt "role" fields with
        | Some json -> Role.t_of_yojson json
        | None -> raise (Json.Of_json ("Missing role field", `Assoc fields))
      in
      let content = match List.assoc_opt "content" fields with
        | Some (`Assoc content_fields) ->
          (match List.assoc_opt "type" content_fields with
           | Some (`String "text") -> `Text (TextContent.t_of_yojson (`Assoc content_fields))
           | Some (`String "image") -> `Image (ImageContent.t_of_yojson (`Assoc content_fields))
           | _ -> raise (Json.Of_json ("Invalid content type", `Assoc content_fields)))
        | _ -> raise (Json.Of_json ("Missing or invalid content field", `Assoc fields))
      in
      { role; content }
    | j -> raise (Json.Of_json ("Expected object for SamplingMessage", j))
end

(* Implementation info *)

module Implementation = struct
  type t = {
    name: string;
    version: string;
  }

  let yojson_of_t { name; version } =
    `Assoc [
      ("name", `String name);
      ("version", `String version);
    ]

  let t_of_yojson = function
    | `Assoc fields ->
      let name = match List.assoc_opt "name" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'name' field", `Assoc fields))
      in
      let version = match List.assoc_opt "version" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'version' field", `Assoc fields))
      in
      { name; version }
    | j -> raise (Json.Of_json ("Expected object for Implementation", j))
end

(* JSONRPC Message types *)


module JSONRPCMessage = struct
  type notification = {
    method_: string;
    params: Json.t option;
  }

  type request = {
    id: RequestId.t;
    method_: string;
    params: Json.t option;
    progress_token: ProgressToken.t option;
  }

  type response = {
    id: RequestId.t;
    result: Json.t;
  }

  type error = {
    id: RequestId.t;
    code: int;
    message: string;
    data: Json.t option;
  }

  type t =
    | Notification of notification
    | Request of request
    | Response of response
    | Error of error

  let yojson_of_notification (n: notification) =
    let assoc = [
      ("jsonrpc", `String "2.0");
      ("method", `String n.method_);
    ] in
    let assoc = match n.params with
      | Some params -> ("params", params) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let yojson_of_request (r: request) =
    let assoc = [
      ("jsonrpc", `String "2.0");
      ("id", Id.yojson_of_t r.id);
      ("method", `String r.method_);
    ] in
    let assoc = match r.params with
      | Some params ->
        let params_json = match params with
          | `Assoc fields ->
            let fields = match r.progress_token with
              | Some token ->
                let meta = `Assoc [ "progressToken", ProgressToken.yojson_of_t token ] in
                ("_meta", meta) :: fields
              | None -> fields
            in
            `Assoc fields
          | _ -> params
        in
        ("params", params_json) :: assoc
      | None -> assoc
    in
    `Assoc assoc

  let yojson_of_response (r: response) =
    `Assoc [
      ("jsonrpc", `String "2.0");
      ("id", Id.yojson_of_t r.id);
      ("result", r.result);
    ]

  let yojson_of_error (e: error) =
    let error_assoc = [
      ("code", `Int e.code);
      ("message", `String e.message);
    ] in
    let error_assoc = match e.data with
      | Some data -> ("data", data) :: error_assoc
      | None -> error_assoc
    in
    `Assoc [
      ("jsonrpc", `String "2.0");
      ("id", Id.yojson_of_t e.id);
      ("error", `Assoc error_assoc);
    ]

  let yojson_of_t = function
    | Notification n -> yojson_of_notification n
    | Request r -> yojson_of_request r
    | Response r -> yojson_of_response r
    | Error e -> yojson_of_error e

  let notification_of_yojson = function
    | `Assoc fields ->
      let method_ = match List.assoc_opt "method" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'method' field", `Assoc fields))
      in
      let params = List.assoc_opt "params" fields in
      { method_; params }
    | j -> raise (Json.Of_json ("Expected object for notification", j))

  let request_of_yojson = function
    | `Assoc fields ->
      let id = match List.assoc_opt "id" fields with
        | Some id_json -> Id.t_of_yojson id_json
        | _ -> raise (Json.Of_json ("Missing or invalid 'id' field", `Assoc fields))
      in
      let method_ = match List.assoc_opt "method" fields with
        | Some (`String s) -> s
        | _ -> raise (Json.Of_json ("Missing or invalid 'method' field", `Assoc fields))
      in
      let params = List.assoc_opt "params" fields in
      let progress_token = 
        match params with
        | Some (`Assoc param_fields) -> 
          (match List.assoc_opt "_meta" param_fields with
           | Some (`Assoc meta_fields) -> 
             (match List.assoc_opt "progressToken" meta_fields with
              | Some token_json -> Some (ProgressToken.t_of_yojson token_json)
              | None -> None)
           | _ -> None)
        | _ -> None
      in
      { id; method_; params; progress_token }
    | j -> raise (Json.Of_json ("Expected object for request", j))

  let response_of_yojson = function
    | `Assoc fields ->
      let id = match List.assoc_opt "id" fields with
        | Some id_json -> Id.t_of_yojson id_json
        | _ -> raise (Json.Of_json ("Missing or invalid 'id' field", `Assoc fields))
      in
      let result = match List.assoc_opt "result" fields with
        | Some result -> result
        | _ -> raise (Json.Of_json ("Missing 'result' field", `Assoc fields))
      in
      { id; result }
    | j -> raise (Json.Of_json ("Expected object for response", j))

  let error_of_yojson = function
    | `Assoc fields ->
      let id = match List.assoc_opt "id" fields with
        | Some id_json -> Id.t_of_yojson id_json
        | _ -> raise (Json.Of_json ("Missing or invalid 'id' field", `Assoc fields))
      in
      let error = match List.assoc_opt "error" fields with
        | Some (`Assoc error_fields) -> error_fields
        | _ -> raise (Json.Of_json ("Missing or invalid 'error' field", `Assoc fields))
      in
      let code = match List.assoc_opt "code" error with
        | Some (`Int code) -> code
        | _ -> raise (Json.Of_json ("Missing or invalid 'code' field in error", `Assoc error))
      in
      let message = match List.assoc_opt "message" error with
        | Some (`String msg) -> msg
        | _ -> raise (Json.Of_json ("Missing or invalid 'message' field in error", `Assoc error))
      in
      let data = List.assoc_opt "data" error in
      { id; code; message; data }
    | j -> raise (Json.Of_json ("Expected object for error", j))

  let t_of_yojson json =
    match json with
    | `Assoc fields ->
      let _jsonrpc = match List.assoc_opt "jsonrpc" fields with
        | Some (`String "2.0") -> ()
        | _ -> raise (Json.Of_json ("Missing or invalid 'jsonrpc' field", json))
      in
      if List.mem_assoc "method" fields then
        if List.mem_assoc "id" fields then
          Request (request_of_yojson json)
        else
          Notification (notification_of_yojson json)
      else if List.mem_assoc "result" fields then
        Response (response_of_yojson json)
      else if List.mem_assoc "error" fields then
        Error (error_of_yojson json)
      else
        raise (Json.Of_json ("Invalid JSONRPC message format", json))
    | j -> raise (Json.Of_json ("Expected object for JSONRPC message", j))

  let create_notification ?(params=None) ~method_ () =
    Notification { method_; params }

  let create_request ?(params=None) ?(progress_token=None) ~id ~method_ () =
    Request { id; method_; params; progress_token }

  let create_response ~id ~result =
    Response { id; result }

  let create_error ~id ~code ~message ?(data=None) () =
    Error { id; code; message; data }
end

(* MCP-specific request/response types *)

module Initialize = struct
  module Request = struct
    type t = {
      capabilities: Json.t; (* ClientCapabilities *)
      client_info: Implementation.t;
      protocol_version: string;
    }

    let yojson_of_t { capabilities; client_info; protocol_version } =
      `Assoc [
        ("capabilities", capabilities);
        ("clientInfo", Implementation.yojson_of_t client_info);
        ("protocolVersion", `String protocol_version);
      ]

    let t_of_yojson = function
      | `Assoc fields ->
        let capabilities = match List.assoc_opt "capabilities" fields with
          | Some json -> json
          | None -> raise (Json.Of_json ("Missing capabilities field", `Assoc fields))
        in
        let client_info = match List.assoc_opt "clientInfo" fields with
          | Some json -> Implementation.t_of_yojson json
          | None -> raise (Json.Of_json ("Missing clientInfo field", `Assoc fields))
        in
        let protocol_version = match List.assoc_opt "protocolVersion" fields with
          | Some (`String s) -> s
          | _ -> raise (Json.Of_json ("Missing or invalid protocolVersion field", `Assoc fields))
        in
        { capabilities; client_info; protocol_version }
      | j -> raise (Json.Of_json ("Expected object for InitializeRequest", j))

    let create ~capabilities ~client_info ~protocol_version =
      { capabilities; client_info; protocol_version }

    let to_jsonrpc ~id t =
      let params = yojson_of_t t in
      JSONRPCMessage.create_request ~id ~method_:"initialize" ~params:(Some params) ()
  end

  module Result = struct
    type t = {
      capabilities: Json.t; (* ServerCapabilities *)
      server_info: Implementation.t;
      protocol_version: string;
      instructions: string option;
      meta: Json.t option;
    }

    let yojson_of_t { capabilities; server_info; protocol_version; instructions; meta } =
      let assoc = [
        ("capabilities", capabilities);
        ("serverInfo", Implementation.yojson_of_t server_info);
        ("protocolVersion", `String protocol_version);
      ] in
      let assoc = match instructions with
        | Some instr -> ("instructions", `String instr) :: assoc
        | None -> assoc
      in
      let assoc = match meta with
        | Some meta -> ("_meta", meta) :: assoc
        | None -> assoc
      in
      `Assoc assoc

    let t_of_yojson = function
      | `Assoc fields ->
        let capabilities = match List.assoc_opt "capabilities" fields with
          | Some json -> json
          | None -> raise (Json.Of_json ("Missing capabilities field", `Assoc fields))
        in
        let server_info = match List.assoc_opt "serverInfo" fields with
          | Some json -> Implementation.t_of_yojson json
          | None -> raise (Json.Of_json ("Missing serverInfo field", `Assoc fields))
        in
        let protocol_version = match List.assoc_opt "protocolVersion" fields with
          | Some (`String s) -> s
          | _ -> raise (Json.Of_json ("Missing or invalid protocolVersion field", `Assoc fields))
        in
        let instructions = match List.assoc_opt "instructions" fields with
          | Some (`String s) -> Some s
          | _ -> None
        in
        let meta = List.assoc_opt "_meta" fields in
        { capabilities; server_info; protocol_version; instructions; meta }
      | j -> raise (Json.Of_json ("Expected object for InitializeResult", j))

    let create ~capabilities ~server_info ~protocol_version ?instructions ?meta () =
      { capabilities; server_info; protocol_version; instructions; meta }

    let to_jsonrpc ~id t =
      JSONRPCMessage.create_response ~id ~result:(yojson_of_t t)
  end
end

module Initialized = struct
  module Notification = struct
    type t = {
      meta: Json.t option;
    }

    let yojson_of_t { meta } =
      let assoc = [] in
      let assoc = match meta with
        | Some meta -> ("_meta", meta) :: assoc
        | None -> assoc
      in
      `Assoc assoc

    let t_of_yojson = function
      | `Assoc fields ->
        let meta = List.assoc_opt "_meta" fields in
        { meta }
      | j -> raise (Json.Of_json ("Expected object for InitializedNotification", j))

    let create ?meta () = { meta }

    let to_jsonrpc t =
      let params = match yojson_of_t t with
        | `Assoc [] -> None
        | json -> Some json
      in
      JSONRPCMessage.create_notification ~method_:"notifications/initialized" ~params ()
  end
end

(* Export the main interface for using the MCP protocol *)

let parse_message json =
  JSONRPCMessage.t_of_yojson json

let create_notification = JSONRPCMessage.create_notification
let create_request = JSONRPCMessage.create_request
let create_response = JSONRPCMessage.create_response
let create_error = JSONRPCMessage.create_error

(* Helper functions *)
let create_completion_request ~id ~argument ~ref =
  let params = Completion.Request.to_params { argument; ref } in
  create_request ~id ~method_:"completion/complete" ~params:(Some params) ()
  
let create_completion_response ~id ~values ?(has_more=None) ?(total=None) ?(meta=None) () =
  let completion = { Completion.Result.values; has_more; total } in
  let result = Completion.Result.to_result { completion; meta } in
  create_response ~id ~result