mcp_kit

defmodule MCPKit.Tool do
  @moduledoc """
  Behaviour and schema DSL for MCP tool modules.

  A tool module declares its JSON-schema-like input contract with `schema/1`
  and implements `execute/2` to return an `MCPKit.Response`.

  Example:

      defmodule MyApp.MCP.Tools.Ping do
        use MCPKit.Tool

        alias MCPKit.Response

        schema do
          field :message, :string, required: true
        end

        def execute(arguments, context) do
          {:reply, Response.tool() |> Response.structured(arguments), context}
        end
      end

  Supported schema nodes today:
  - `field/2`
  - `field/3`
  - `embeds_many/2`
  - `embeds_many/3`
  """

  @callback description() :: String.t() | nil
  @callback input_schema() :: map()
  @callback validate_arguments(map()) :: {:ok, map()} | {:error, String.t()}
  @callback execute(map(), map()) :: {:reply, MCPKit.Response.t(), map()}

  defmacro __using__(_opts) do
    quote do
      @behaviour MCPKit.Tool
      import MCPKit.Tool, only: [schema: 1]

      def description, do: MCPKit.Tool.module_description(@moduledoc)

      defoverridable description: 0
    end
  end

  defmacro schema(do: block) do
    spec = parse_schema_block(block)

    quote do
      @mcp_argument_spec unquote(Macro.escape(spec))

      def input_schema, do: MCPKit.Tool.to_json_schema(@mcp_argument_spec)

      def validate_arguments(arguments),
        do: MCPKit.Tool.validate_arguments(arguments, @mcp_argument_spec)
    end
  end

  def module_description({_, value}), do: module_description(value)
  def module_description(value) when is_binary(value), do: String.trim(value)
  def module_description(_value), do: nil

  def to_json_schema(spec) when is_list(spec) do
    properties =
      Map.new(spec, fn
        {:field, name, type, opts} ->
          {Atom.to_string(name), field_schema(type, opts)}

        {:embeds_many, name, opts, nested_spec} ->
          {Atom.to_string(name),
           %{
             "type" => "array",
             "items" => to_json_schema(nested_spec),
             "description" => Keyword.get(opts, :description)
           }
           |> reject_nil_values()}
      end)

    required =
      spec
      |> Enum.filter(&required_field?/1)
      |> Enum.map(fn {_, name, _, _} -> Atom.to_string(name) end)

    %{
      "type" => "object",
      "properties" => properties,
      "required" => required,
      "additionalProperties" => false
    }
  end

  def validate_arguments(arguments, spec) when is_map(arguments) do
    case validate_object(arguments, spec, []) do
      {:ok, normalized, []} -> {:ok, normalized}
      {:ok, _normalized, errors} -> {:error, Enum.join(errors, "; ")}
    end
  end

  def validate_arguments(_arguments, _spec), do: {:error, "arguments must be an object"}

  defp validate_object(arguments, spec, path) do
    allowed_keys = MapSet.new(Enum.map(spec, fn {_, name, _, _} -> Atom.to_string(name) end))

    {normalized, errors} =
      Enum.reduce(spec, {%{}, []}, fn schema, {acc, errs} ->
        {value, new_errors} = validate_entry(schema, arguments, path)

        case value do
          :missing -> {acc, errs ++ new_errors}
          {key, normalized_value} -> {Map.put(acc, key, normalized_value), errs ++ new_errors}
        end
      end)

    unknown_errors =
      arguments
      |> Map.keys()
      |> Enum.reject(&MapSet.member?(allowed_keys, &1))
      |> Enum.map(&(format_path(path, &1) <> " is not allowed"))

    {:ok, normalized, errors ++ unknown_errors}
  end

  defp validate_entry({:field, name, type, opts}, arguments, path) do
    key = Atom.to_string(name)
    required? = Keyword.get(opts, :required, false)

    case Map.fetch(arguments, key) do
      :error when required? ->
        {:missing, [format_path(path, key) <> " is required"]}

      :error ->
        {:missing, []}

      {:ok, value} ->
        case validate_scalar(value, type) do
          {:ok, normalized} -> {{name, normalized}, []}
          {:error, message} -> {:missing, [format_path(path, key) <> " " <> message]}
        end
    end
  end

  defp validate_entry({:embeds_many, name, opts, nested_spec}, arguments, path) do
    key = Atom.to_string(name)
    required? = Keyword.get(opts, :required, false)

    case Map.fetch(arguments, key) do
      :error when required? ->
        {:missing, [format_path(path, key) <> " is required"]}

      :error ->
        {:missing, []}

      {:ok, value} when is_list(value) ->
        {items, errors} =
          value
          |> Enum.with_index()
          |> Enum.map_reduce([], fn {item, index}, errs ->
            if is_map(item) do
              case validate_object(item, nested_spec, path ++ ["#{key}[#{index}]"]) do
                {:ok, normalized, item_errors} -> {normalized, errs ++ item_errors}
              end
            else
              {%{}, errs ++ [format_path(path, "#{key}[#{index}]") <> " must be an object"]}
            end
          end)

        if errors == [], do: {{name, items}, []}, else: {:missing, errors}

      {:ok, _value} ->
        {:missing, [format_path(path, key) <> " must be an array"]}
    end
  end

  defp validate_scalar(value, :string) when is_binary(value), do: {:ok, value}
  defp validate_scalar(value, :boolean) when is_boolean(value), do: {:ok, value}
  defp validate_scalar(value, :integer) when is_integer(value), do: {:ok, value}
  defp validate_scalar(value, :number) when is_number(value), do: {:ok, value}
  defp validate_scalar(value, :map) when is_map(value), do: {:ok, value}
  defp validate_scalar(_value, _type), do: {:error, "has an invalid type"}

  defp field_schema(type, opts) do
    %{
      "type" => json_type(type),
      "description" => Keyword.get(opts, :description)
    }
    |> reject_nil_values()
  end

  defp json_type(:string), do: "string"
  defp json_type(:boolean), do: "boolean"
  defp json_type(:integer), do: "integer"
  defp json_type(:number), do: "number"
  defp json_type(:map), do: "object"

  defp required_field?({_kind, _name, opts, _nested_spec}) when is_list(opts),
    do: Keyword.get(opts, :required, false)

  defp required_field?({_kind, _name, _type, opts}) when is_list(opts),
    do: Keyword.get(opts, :required, false)

  defp reject_nil_values(map) do
    Map.reject(map, fn {_key, value} -> is_nil(value) end)
  end

  defp format_path([], key), do: key
  defp format_path(path, key), do: Enum.join(path ++ [key], ".")

  defp parse_schema_block({:__block__, _, fields}), do: Enum.map(fields, &parse_schema_entry/1)
  defp parse_schema_block(field), do: [parse_schema_entry(field)]

  defp parse_schema_entry({:field, _meta, [name, type]}) do
    {:field, extract_atom(name), extract_atom(type), []}
  end

  defp parse_schema_entry({:field, _meta, [name, type, opts]}) do
    {:field, extract_atom(name), extract_atom(type), extract_keyword(opts)}
  end

  defp parse_schema_entry({:embeds_many, _meta, [name, [do: block]]}) do
    {:embeds_many, extract_atom(name), [], parse_schema_block(block)}
  end

  defp parse_schema_entry({:embeds_many, _meta, [name, opts, [do: block]]}) do
    {:embeds_many, extract_atom(name), extract_keyword(opts), parse_schema_block(block)}
  end

  defp extract_atom(atom) when is_atom(atom), do: atom

  defp extract_keyword(keyword) when is_list(keyword), do: keyword
end