codex_sdk

defmodule Codex.Tools.Registry do
  @moduledoc """
  ETS-backed registry that stores Codex tool definitions, tracks invocation telemetry,
  and provides lookup/invoke helpers used by the auto-run pipeline.
  """

  alias Codex.Telemetry
  alias Codex.Tools
  alias Codex.Tools.Handle

  @table :codex_tools_registry

  @doc """
  Registers a tool by name, storing the implementing module and metadata in ETS.
  Returns a `Codex.Tools.Handle` for later deregistration.
  """
  def register(%{name: name, module: module, metadata: metadata}) do
    ensure_table()

    if :ets.insert_new(@table, {name, module, metadata}) do
      {:ok, %Handle{name: name, module: module}}
    else
      {:error, {:already_registered, name}}
    end
  end

  @doc """
  Removes the supplied tool handle from the registry.
  """
  def deregister(%Handle{name: name}) do
    ensure_table()
    :ets.delete(@table, name)
    :ok
  end

  @doc """
  Looks up a tool by name and returns its metadata and module.
  """
  def lookup(name) when is_binary(name) do
    ensure_table()

    case :ets.lookup(@table, name) do
      [{^name, module, metadata}] ->
        {:ok, %{name: name, module: module, metadata: metadata}}

      [] ->
        lookup_alias(name)
    end
  end

  @doc """
  Invokes a registered tool by name, normalising arguments and emitting telemetry around
  the attempt. Returns the tool's response or an error tuple.
  """
  def invoke(name, args, context) when is_binary(name) do
    with {:ok, %{module: module, metadata: metadata} = info} <- lookup(name) do
      normalized_args = normalize_args(args)

      full_context =
        context
        |> Map.put(:tool, %{name: info.name, metadata: metadata})
        |> Map.put(:metadata, metadata)

      telemetry_meta =
        build_telemetry_metadata(
          info.name,
          module,
          metadata,
          normalized_args,
          full_context
        )
        |> Map.put(:originator, :sdk)
        |> Map.put(:span_token, make_ref())

      if tool_enabled?(metadata, full_context) do
        started = emit_tool_start(telemetry_meta)

        invoke_registered(
          info.name,
          module,
          metadata,
          normalized_args,
          full_context,
          telemetry_meta,
          started
        )
      else
        {:error, {:tool_disabled, name}}
      end
    end
  end

  defp safe_invoke(module, args, context) do
    module.invoke(args, context)
  rescue
    error -> {:error, {:tool_exception, module, error}}
  catch
    kind, reason -> {:error, {:tool_failure, module, {kind, reason}}}
  end

  defp emit_tool_start(telemetry_meta) do
    Telemetry.emit(
      [:codex, :tool, :start],
      %{system_time: System.system_time()},
      telemetry_meta
    )

    System.monotonic_time()
  end

  defp invoke_registered(name, module, metadata, args, context, telemetry_meta, started) do
    case safe_invoke(module, args, context) do
      {:ok, output} ->
        handle_success(name, telemetry_meta, started, output)

      {:error, reason} ->
        handle_invoke_error(name, metadata, reason, context, telemetry_meta, started)
    end
  end

  defp handle_invoke_error(name, metadata, reason, context, telemetry_meta, started) do
    case maybe_handle_error(metadata, reason, context) do
      {:ok, output} ->
        handle_success(name, telemetry_meta, started, output)

      {:error, handled_reason} = error ->
        handle_failure(name, telemetry_meta, started, handled_reason)
        error
    end
  end

  @doc """
  Clears the registry, deleting the underlying ETS table.
  """
  def reset! do
    case :ets.whereis(@table) do
      :undefined ->
        :ok

      _ ->
        try do
          :ets.delete(@table)
        rescue
          ArgumentError ->
            # Another process may have deleted the table between whereis/1 and delete/1.
            :ok
        end
    end

    ensure_table()
    :ok
  end

  defp ensure_table do
    case :ets.whereis(@table) do
      :undefined ->
        try do
          :ets.new(@table, [
            :named_table,
            :public,
            :set,
            read_concurrency: true,
            write_concurrency: true
          ])
        rescue
          ArgumentError ->
            # Another process created the table concurrently.
            :ok
        end

      _ ->
        :ok
    end
  end

  defp build_telemetry_metadata(tool, module, metadata, args, context) do
    event = Map.get(context, :event)
    thread = Map.get(context, :thread)
    warnings = sandbox_warnings(context)
    capabilities = capability_metadata(context)

    %{
      tool: tool,
      module: module,
      metadata: metadata,
      arguments: args,
      retry?: Map.get(context, :retry?, false),
      attempt: Map.get(context, :attempt),
      call_id: event && Map.get(event, :call_id),
      thread_id: thread && Map.get(thread, :thread_id),
      turn_id: event && Map.get(event, :turn_id)
    }
    |> maybe_put(:capabilities, capabilities)
    |> maybe_put(:sandbox_warnings, warnings)
  end

  defp duration_to_ms(duration_native) when is_integer(duration_native) do
    duration_native
    |> System.convert_time_unit(:native, :microsecond)
    |> div(1000)
  end

  defp sandbox_warnings(context) do
    Map.get(context, :sandbox_warnings) ||
      Map.get(context, "sandbox_warnings") ||
      Map.get(context, :warnings) ||
      Map.get(context, "warnings")
  end

  defp capability_metadata(context) do
    Map.get(context, :capabilities) ||
      Map.get(context, "capabilities")
  end

  defp normalize_args(args) when is_map(args), do: args
  defp normalize_args(args) when is_list(args), do: Map.new(args)

  defp normalize_args(args) when is_binary(args) do
    case Jason.decode(args) do
      {:ok, decoded} when is_map(decoded) -> decoded
      _ -> %{"input" => args}
    end
  rescue
    _ -> %{"input" => args}
  end

  defp normalize_args(nil), do: %{}

  defp normalize_args(other) do
    Map.new(other)
  rescue
    _ -> %{"input" => other}
  end

  defp tool_enabled?(metadata, context) do
    case metadata |> Map.get(:enabled?) || Map.get(metadata, "enabled?") do
      fun when is_function(fun, 2) -> fun.(context, metadata)
      fun when is_function(fun, 1) -> fun.(context)
      fun when is_function(fun, 0) -> fun.()
      value when is_boolean(value) -> value
      _ -> true
    end
  end

  defp lookup_alias(name) do
    name
    |> alias_candidates()
    |> Enum.reduce_while({:error, :not_found}, fn alias_name, _acc ->
      case :ets.lookup(@table, alias_name) do
        [{^alias_name, module, metadata}] ->
          {:halt, {:ok, %{name: alias_name, module: module, metadata: metadata}}}

        [] ->
          {:cont, {:error, :not_found}}
      end
    end)
  end

  defp alias_candidates("container.exec"), do: ["shell", "shell_command"]
  defp alias_candidates("local_shell"), do: ["shell", "shell_command"]
  defp alias_candidates("shell_command"), do: ["shell"]
  defp alias_candidates(_name), do: []

  defp maybe_handle_error(metadata, reason, context) do
    case metadata |> Map.get(:on_error) || Map.get(metadata, "on_error") do
      fun when is_function(fun, 3) ->
        safe_error_apply(fun, reason, context)

      fun when is_function(fun, 2) ->
        safe_error_apply(fun, reason, context)

      fun when is_function(fun, 1) ->
        safe_error_apply(fun, reason, context)

      _ ->
        {:error, reason}
    end
  end

  defp safe_error_apply(fun, reason, context) when is_function(fun, 3) do
    fun.(reason, context, context)
  rescue
    error -> {:error, {:error_handler_failed, error}}
  catch
    kind, payload -> {:error, {:error_handler_failed, {kind, payload}}}
  end

  defp safe_error_apply(fun, reason, context) when is_function(fun, 2) do
    fun.(reason, context)
  rescue
    error -> {:error, {:error_handler_failed, error}}
  catch
    kind, payload -> {:error, {:error_handler_failed, {kind, payload}}}
  end

  defp safe_error_apply(fun, reason, _context) when is_function(fun, 1) do
    fun.(reason)
  rescue
    error -> {:error, {:error_handler_failed, error}}
  catch
    kind, payload -> {:error, {:error_handler_failed, {kind, payload}}}
  end

  defp handle_success(name, telemetry_meta, started, output) do
    duration = System.monotonic_time() - started
    Tools.record_invocation(name, :success, duration_to_ms(duration))

    Telemetry.emit(
      [:codex, :tool, :success],
      %{duration: duration, system_time: System.system_time()},
      telemetry_meta
      |> Map.put(:output, output)
      |> Map.put(:result, :ok)
    )

    {:ok, output}
  end

  defp handle_failure(name, telemetry_meta, started, reason) do
    duration = System.monotonic_time() - started
    Tools.record_invocation(name, :failure, duration_to_ms(duration), reason)

    Telemetry.emit(
      [:codex, :tool, :failure],
      %{duration: duration, system_time: System.system_time()},
      telemetry_meta
      |> Map.put(:error, reason)
      |> Map.put(:result, :error)
    )
  end

  defp maybe_put(map, _key, nil), do: map
  defp maybe_put(map, _key, []), do: map
  defp maybe_put(map, key, value), do: Map.put(map, key, value)
end