llm

defmodule LLM.Tool do
  @moduledoc """
  Tool definition behaviour and inline tool creation.

  Tools can be defined as modules implementing the behaviour, or inline
  using `new/1`.

  ## Module Tool

      defmodule MyApp.Tools.ReadFile do
        @behaviour LLM.Tool

        @impl true
        def name, do: "read_file"

        @impl true
        def description, do: "Read the contents of a file"

        @impl true
        def input_schema do
          %{
            "type" => "object",
            "properties" => %{
              "path" => %{"type" => "string", "description" => "File path"}
            },
            "required" => ["path"]
          }
        end

        @impl true
        def execute(%{"path" => path}, _context) do
          File.read(path)
        end
      end

  ## Inline Tool

      LLM.Tool.new(%{
        name: "shell",
        description: "Execute a shell command",
        input_schema: %{
          "type" => "object",
          "properties" => %{
            "command" => %{"type" => "string"}
          },
          "required" => ["command"]
        },
        run: fn %{"command" => cmd} ->
          {output, _} = System.cmd("sh", ["-c", cmd])
          {:ok, output}
        end
      })
  """

  @type t :: %__MODULE__{
          name: String.t(),
          description: String.t(),
          input_schema: map(),
          execute: (map(), map() -> {:ok, term()} | {:error, term()})
        }

  defstruct [:name, :description, :input_schema, :execute]

  @callback name() :: String.t()
  @callback description() :: String.t()
  @callback input_schema() :: map()
  @callback execute(map(), map()) :: {:ok, term()} | {:error, term()}

  @doc """
  Create an inline tool from a map.

  ## Options

  - `:name` - tool name (required)
  - `:description` - tool description (required)
  - `:input_schema` - JSON Schema for tool input (required)
  - `:run` - function to execute. Receives input map. Can be arity 1 or 2.
  """
  def new(%{name: name, description: desc, input_schema: schema, run: fun})
      when is_function(fun, 1) do
    %__MODULE__{
      name: name,
      description: desc,
      input_schema: schema,
      execute: fn input, _ctx -> fun.(input) end
    }
  end

  def new(%{name: name, description: desc, input_schema: schema, run: fun})
      when is_function(fun, 2) do
    %__MODULE__{
      name: name,
      description: desc,
      input_schema: schema,
      execute: fun
    }
  end

  def new(%{name: _, description: _, input_schema: _} = attrs) do
    raise ArgumentError,
          "Tool requires a :run function. Got: #{inspect(Map.keys(attrs))}"
  end

  def new(attrs) do
    raise ArgumentError,
          "Tool requires :name, :description, :input_schema, and :run. Got: #{inspect(attrs)}"
  end

  @doc """
  Convert a module implementing the `LLM.Tool` behaviour to a tool struct.
  """
  def from_module(module) when is_atom(module) do
    %__MODULE__{
      name: module.name(),
      description: module.description(),
      input_schema: module.input_schema(),
      execute: &module.execute/2
    }
  end

  @doc """
  Normalize a tool specification to a `LLM.Tool.t()` struct.

  Accepts:
  - `%LLM.Tool{}` — returned as-is
  - A module — converted via `from_module/1`
  - A `{module, opts}` tuple — module tool with extra context merged into execution
  """
  def normalize(%LLM.Tool{} = tool), do: tool

  def normalize(module) when is_atom(module) do
    from_module(module)
  end

  def normalize({module, extra_context}) when is_atom(module) do
    tool = from_module(module)

    %{
      tool
      | execute: fn input, ctx ->
          merged = Map.merge(ctx, if(is_map(extra_context), do: extra_context, else: %{}))
          tool.execute.(input, merged)
        end
    }
  end
end