gen_mcp

defmodule GenMCP.MCP do
  @moduledoc """
  Helpers for building MCP payload structs from Servers, Tools or repositories.

  Each constructor normalizes various type of inputs into the MCP structs
  expected by the wire protocol.
  """

  alias GenMCP.MCP
  alias GenMCP.Suite.Tool

  @type content_block_option ::
          {:text, binary()}
          | {:resource, %{uri: binary(), text: binary()}}
          | {:resource, %{uri: binary(), blob: binary()}}
          | {:link, %{name: binary(), uri: binary()}}
          | {:image, {binary(), binary()}}
          | {:audio, {binary(), binary()}}

  @type content_block ::
          content_block_option()
          | MCP.TextContent.t()
          | MCP.AudioContent.t()
          | MCP.ImageContent.t()
          | MCP.EmbeddedResource.t()
          | MCP.ResourceLink.t()

  defp require_key!(keywords, key, errmsg) do
    case :lists.keyfind(key, 1, keywords) do
      {^key, value} -> value
      false -> raise KeyError, key: key, term: keywords, message: errmsg
    end
  end

  defmacrop cur_fun do
    {function, arity} = __ENV__.function
    _mfa = Exception.format_mfa(__ENV__.module, function, arity)
  end

  @doc """
  Builds a `%#{inspect(MCP.InitializeResult)}{}` for the initialize handshake.

  Requires `:server_info` (usually a `%#{inspect(MCP.Implementation)}{}`) and
  uses `:capabilities` if provided. The `protocolVersion` field is fixed to
  `"2025-11-25"` to match the latest MCP spec supported here.

  ## Example

      MCP.intialize_result(
        server_info: %#{inspect(MCP.Implementation)}{
          name: "TestServer",
          version: "1.0.0"
        }
      )
  """
  @spec intialize_result(keyword()) :: MCP.InitializeResult.t()
  def intialize_result(opts) do
    %MCP.InitializeResult{
      capabilities: Keyword.get(opts, :capabilities, %{}),
      serverInfo: Keyword.fetch!(opts, :server_info),
      protocolVersion: "2025-11-25"
    }
  end

  @doc """
  Normalizes capability flags/maps into `%#{inspect(MCP.ServerCapabilities)}{}`.

  Passing `true` for a key yields an empty map, while maps are returned
  unchanged and other values keep the field `nil`. This allows flexible
  declaration of server capabilities during initialization.
  """
  @spec capabilities(keyword()) :: MCP.ServerCapabilities.t()
  def capabilities(opts) do
    attrs =
      Enum.flat_map(opts, fn
        {key, value} when is_map(value) -> [{key, value}]
        {key, true} -> [{key, %{}}]
        _ -> []
      end)

    struct!(MCP.ServerCapabilities, attrs)
  end

  @doc """
  Builds `%#{inspect(MCP.Implementation)}{}` entries used in initialize replies.

  `:name` and `:version` trigger `KeyError` when missing, while `:title` is
  optional.
  """
  @spec server_info(keyword()) :: MCP.Implementation.t()
  def server_info(opts) do
    %MCP.Implementation{
      name: require_key!(opts, :name, "option :name is required by #{cur_fun()}"),
      version: require_key!(opts, :version, "option :version is required by #{cur_fun()}"),
      title: Keyword.get(opts, :title)
    }
  end

  # TODO handle cursor

  @doc """
  Builds `%#{inspect(MCP.ListToolsResult)}{}` for the tools advertised by a
  server.

  Structs already shaped as `%#{inspect(MCP.Tool)}{}` are left untouched and
  other entries go through `GenMCP.Suite.Tool.describe/1`.

  Pagination for tools is not supported yet.

  ## Example

      MCP.list_tools_result([
        %#{inspect(MCP.Tool)}{name: "tool1", inputSchema: %{type: "object"}},
        %{name: "tool2", inputSchema: %{type: "object"}}
      ])
  """
  @spec list_tools_result([Tool.tool() | MCP.Tool.t()]) :: MCP.ListToolsResult.t()
  def list_tools_result(tools) do
    %MCP.ListToolsResult{
      tools:
        Enum.map(tools, fn
          %MCP.Tool{} = tool -> tool
          tool -> Tool.describe(tool)
        end)
    }
  end

  @doc """
  Normalizes content shortcuts into the MCP structs that tools and prompts
  expect.

  This helper is generally not used directly, but delegated to from
  `call_tool_result/1` or `get_prompt_result/1`.

  Supported shortcuts are
  - `{:text, binary()}`
  - `{:resource, %{uri: binary(), text: binary()}}`
  - `{:resource, %{uri: binary(), blob: binary()}}`
  - `{:link, %{name: binary(), uri: binary()}}`
  - `{:image, {mime, data}}`
  - `{:audio, {mime, data}}`

  ## Example

      MCP.content_block({:audio, {"audio/mp3", "base64data"}})
  """
  @spec content_block(content_block_option()) ::
          MCP.TextContent.t()
          | MCP.EmbeddedResource.t()
          | MCP.ResourceLink.t()
          | MCP.ImageContent.t()
          | MCP.AudioContent.t()
  def content_block(content_opts)

  def content_block({:text, text}) when is_binary(text) do
    %MCP.TextContent{text: text}
  end

  def content_block({:resource, %{text: text, uri: uri} = resource})
      when is_binary(text) and is_binary(uri) do
    %MCP.EmbeddedResource{resource: resource}
  end

  def content_block({:resource, %{blob: blob, uri: uri} = resource})
      when is_binary(blob) and is_binary(uri) do
    %MCP.EmbeddedResource{resource: resource}
  end

  def content_block({:link, %{name: name, uri: uri} = link})
      when is_binary(name) and is_binary(uri) do
    struct(MCP.ResourceLink, link)
  end

  def content_block({:image, {mime_type, data}}) when is_binary(mime_type) and is_binary(data) do
    %MCP.ImageContent{mimeType: mime_type, data: data}
  end

  def content_block({:audio, {mime_type, data}}) when is_binary(mime_type) and is_binary(data) do
    %MCP.AudioContent{mimeType: mime_type, data: data}
  end

  def content_block(other) do
    raise ArgumentError, "unsupported content block definition: #{inspect(other)}"
  end

  @doc """
  Builds `%#{inspect(MCP.CallToolResult)}{}` for results returned by tool
  implementations.

  The list may contain shortcuts such as `text: "foo"`, `image: {"mime", data}`
  or literal `%#{inspect(MCP)}.*` structs.

  ## Structured content

  Structured payloads can be provided as a naked map in the list, with the
  `data:` keyword shortcut, or with the `_data:` shortcut. Only one structured
  payload is allowed per result.

  - `data: map` and a naked map both set `structuredContent` and also append a
    JSON-encoded text content block mirroring the payload. This is convenient
    for clients that do not consume structured content directly.
  - `_data: map` sets `structuredContent` without adding the JSON text mirror,
    when the caller wants the structured payload to remain the only carrier
    of that data.

  ## Errors

  Errors can be reported by including `error: true` (sets `isError` to true) or
  `error: "message"` (adds a text content block with the message and sets
  `isError` to true).

  ## Example

      MCP.call_tool_result(text: "Hello, world!", error: true)

      MCP.call_tool_result(text: "Summary", data: %{rows: 3})

      MCP.call_tool_result(text: "Summary", _data: %{rows: 3})
  """
  @spec call_tool_result([
          content_block()
          | {:error, boolean() | binary() | nil}
          | {:data, map()}
          | {:_data, map()}
          | map()
        ]) ::
          MCP.CallToolResult.t()
  def call_tool_result(all_content) when is_list(all_content) do
    {content, {structured_content, error_or_nil?}} =
      Enum.flat_map_reduce(all_content, {nil, nil}, &flat_map_reduce_tool_result/2)

    %MCP.CallToolResult{
      content: content,
      structuredContent: structured_content,
      isError: error_or_nil?
    }
  end

  defp flat_map_reduce_tool_result({:error, error?}, {structured_content, error_or_nil?})
       when is_boolean(error?) do
    {[], {structured_content, error? || error_or_nil?}}
  end

  defp flat_map_reduce_tool_result({:error, errmsg}, {structured_content, _error_or_nil?})
       when is_binary(errmsg) do
    {[content_block({:text, errmsg})], {structured_content, true}}
  end

  defp flat_map_reduce_tool_result({:error, nil}, {structured_content, error_or_nil?}) do
    {[], {structured_content, error_or_nil?}}
  end

  defp flat_map_reduce_tool_result({:data, map}, acc) when is_map(map) do
    add_structured_content(map, acc, _mirror_text? = true)
  end

  defp flat_map_reduce_tool_result({:_data, map}, acc) when is_map(map) do
    add_structured_content(map, acc, _mirror_text? = false)
  end

  defp flat_map_reduce_tool_result(content, {structured_content, error_or_nil?})
       when is_struct(content, MCP.TextContent)
       when is_struct(content, MCP.AudioContent)
       when is_struct(content, MCP.ImageContent)
       when is_struct(content, MCP.EmbeddedResource)
       when is_struct(content, MCP.ResourceLink) do
    {[content], {structured_content, error_or_nil?}}
  end

  # Naked map sets structuredContent and mirrors it as JSON text content.
  defp flat_map_reduce_tool_result(map, acc) when is_map(map) do
    add_structured_content(map, acc, _mirror_text? = true)
  end

  defp flat_map_reduce_tool_result(elem, {structured_content, error_or_nil?}) do
    {[content_block(elem)], {structured_content, error_or_nil?}}
  end

  defp add_structured_content(map, {nil, error_or_nil?}, mirror_text?) do
    extra =
      if mirror_text? do
        [content_block({:text, JSV.Codec.encode!(map)})]
      else
        []
      end

    {extra, {map, error_or_nil?}}
  end

  defp add_structured_content(map, {existing, _error_or_nil?}, _mirror_text?) do
    raise ArgumentError,
          "cannot return multiple structured content, tried to add #{inspect(map)} with existing content: #{inspect(existing)}"
  end

  @doc """
  Wraps the provided `resources` list into
  `%#{inspect(MCP.ListResourcesResult)}{}` for `ListResources` responses.

  `next_cursor` arguments such as `"cursor-123"` are assigned to `nextCursor`,
  while the caller keeps ownership of the resource entries.
  """
  @spec list_resources_result([term()], term() | nil) :: MCP.ListResourcesResult.t()
  def list_resources_result(resources, next_cursor) do
    %MCP.ListResourcesResult{
      resources: resources,
      nextCursor: next_cursor
    }
  end

  @doc """
  Wraps template entries into `%#{inspect(MCP.ListResourceTemplatesResult)}{}`
  for resource template listings.

  Templates are returned as-is so callers may pass either structs or raw maps.

  Pagination is not supported for resource templates.
  """
  @spec list_resource_templates_result([term()]) :: MCP.ListResourceTemplatesResult.t()
  def list_resource_templates_result(templates) do
    %MCP.ListResourceTemplatesResult{
      resourceTemplates: templates
    }
  end

  @doc """
  Wraps resource content helpers into `%#{inspect(MCP.ReadResourceResult)}{}`.

  When given a keyword list (typically `:uri` with `:text` or `:blob`) it
  delegates to `resource_contents/1` and includes a single-element content list.

  Passing a list of maps returns those entries as the content list.

  It is not possible to mix keyword style items and maps in the list, as all
  keyword keys define the same resource content.

  ## Example

      MCP.read_resource_result(uri: "file:///doc.pdf", blob: "bytes")
  """
  @spec read_resource_result(keyword() | [map()]) :: MCP.ReadResourceResult.t()
  def read_resource_result([{k, _} | _] = opts) when is_atom(k) do
    true = Keyword.keyword?(opts)
    contents = resource_contents(opts)

    %MCP.ReadResourceResult{contents: [contents]}
  end

  def read_resource_result([%{} | _] = contents) do
    true = Enum.all?(contents, &is_map/1)

    %MCP.ReadResourceResult{contents: contents}
  end

  @doc """
  Builds either `%#{inspect(MCP.TextResourceContents)}{}` or `%#{inspect(MCP.BlobResourceContents)}{}` from keyword options.

  Requires `:uri` and either `:text` or `:blob`. Optional `:mime_type` and `:_meta` fields are forwarded to the resulting struct. Missing both `:text` and `:blob` raises `ArgumentError`, while omitting `:uri` raises `KeyError`.

  ## Example

      MCP.resource_contents(
        uri: "file:///doc.pdf",
        mime_type: "application/pdf",
        blob: "some-base-64"
      )
  """
  @spec resource_contents(keyword()) ::
          MCP.TextResourceContents.t() | MCP.BlobResourceContents.t()
  def resource_contents(opts) do
    uri = require_key!(opts, :uri, "option :uri is required by #{cur_fun()}")
    mime_type = Keyword.get(opts, :mime_type)
    meta = Keyword.get(opts, :_meta)

    cond do
      Keyword.has_key?(opts, :text) ->
        text = Keyword.fetch!(opts, :text)

        %MCP.TextResourceContents{
          _meta: meta,
          uri: uri,
          text: text,
          mimeType: mime_type
        }

      Keyword.has_key?(opts, :blob) ->
        blob = Keyword.fetch!(opts, :blob)

        %MCP.BlobResourceContents{
          _meta: meta,
          uri: uri,
          blob: blob,
          mimeType: mime_type
        }

      true ->
        raise ArgumentError, "resource_contents/1 requires either :text or :blob option"
    end
  end

  @doc """
  Wraps prompt entries into `%#{inspect(MCP.ListPromptsResult)}{}` for prompt listing responses.

  The provided `prompts` list is returned as-is and `nextCursor` mirrors the optional cursor token passed by the caller.
  """
  @spec list_prompts_result([term()], term() | nil) :: MCP.ListPromptsResult.t()
  def list_prompts_result(prompts, next_cursor) do
    %MCP.ListPromptsResult{
      prompts: prompts,
      nextCursor: next_cursor
    }
  end

  @doc """
  Builds `%#{inspect(MCP.GetPromptResult)}{}` from keyword helpers or explicit prompt entries.

  The helpers treat `text:` and `assistant:` keywords as alternating user/assistant messages, honor an optional `description:` and convert tuples or `%#{inspect(MCP.PromptMessage)}` structs via `content_block/1`. Unsupported shapes such as `:link` or `error:` raise `ArgumentError` so the caller can fix the prompt payload.

  ## Example

      MCP.get_prompt_result(
        description: "A helpful prompt",
        assistant: "You are a an expert for some reason…",
        text: user_input
      )
  """
  @spec get_prompt_result(keyword() | [term()]) :: MCP.GetPromptResult.t()
  def get_prompt_result(opts) do
    {description, opts} =
      case List.keytake(opts, :description, 0) do
        {{:description, description}, opts} -> {description, opts}
        nil -> {nil, opts}
      end

    messages = Enum.map(opts, &map_prompt_message/1)

    %MCP.GetPromptResult{messages: messages, description: description}
  end

  # Map with role and content keys allow to pass custom, invalid elements or
  # MCP.PromptMessage structs.
  defp map_prompt_message(%{role: role, content: _} = elem)
       when is_binary(role)
       when is_atom(role) do
    elem
  end

  defp map_prompt_message({:assistant, text}) when is_binary(text) do
    %MCP.PromptMessage{role: "assistant", content: content_block({:text, text})}
  end

  defp map_prompt_message({_, _} = elem) do
    case content_block(elem) do
      %MCP.ResourceLink{} ->
        raise ArgumentError, "unsupported ResourceLink content in prompt message"

      content ->
        %MCP.PromptMessage{role: "user", content: content}
    end
  end

  defp map_prompt_message(other) do
    raise ArgumentError, "unsupported content block definition for prompt: #{inspect(other)}"
  end
end