ragex

defmodule Ragex.MCP.Telemetry do
  @moduledoc """
  Tracks MCP tool invocation patterns and latencies.

  Backed by an ETS table for near-zero overhead (~1us per write).
  Persists to `~/.ragex/telemetry/<project_hash>.etf` on shutdown
  and reloads on start.

  ## Recorded Metrics (per tool)

  - `:count` -- total invocations
  - `:total_time_us` -- cumulative execution time in microseconds
  - `:last_invoked` -- DateTime of last call
  - `:error_count` -- number of failed invocations

  ## Integration

  Wrap any tool call with `execute/3`:

      Telemetry.execute("semantic_search", fn -> Tools.call_tool(name, args) end)

  Or use the `:telemetry` library events:

  - `[:ragex, :tool, :start]` -- emitted before tool execution
  - `[:ragex, :tool, :stop]` -- emitted after successful execution
  - `[:ragex, :tool, :exception]` -- emitted on error
  """

  use GenServer
  require Logger

  @table :ragex_telemetry
  @persist_interval 60_000

  defmodule Stats do
    @moduledoc false
    defstruct count: 0,
              total_time_us: 0,
              error_count: 0,
              last_invoked: nil
  end

  # ── Client API ───────────────────────────────────────────────────────

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @doc """
  Execute a tool function and record telemetry.

  Returns the result of `fun.()` unchanged.
  """
  @spec execute(String.t(), (-> term())) :: term()
  def execute(tool_name, fun) when is_function(fun, 0) do
    start_time = System.monotonic_time(:microsecond)

    try do
      result = fun.()
      elapsed = System.monotonic_time(:microsecond) - start_time
      record_success(tool_name, elapsed)
      result
    rescue
      e ->
        elapsed = System.monotonic_time(:microsecond) - start_time
        record_error(tool_name, elapsed)
        reraise e, __STACKTRACE__
    end
  end

  @doc "Record a successful tool invocation."
  @spec record_success(String.t(), non_neg_integer()) :: :ok
  def record_success(tool_name, elapsed_us) do
    ensure_table()
    now = DateTime.utc_now()

    :ets.update_counter(@table, tool_name, [{2, 1}, {3, elapsed_us}], {tool_name, 0, 0, 0, nil})
    # Update last_invoked (position 5) -- can't use update_counter for non-integer
    case :ets.lookup(@table, tool_name) do
      [{^tool_name, count, total, errors, _last}] ->
        :ets.insert(@table, {tool_name, count, total, errors, now})

      _ ->
        :ets.insert(@table, {tool_name, 1, elapsed_us, 0, now})
    end

    :ok
  end

  @doc "Record a failed tool invocation."
  @spec record_error(String.t(), non_neg_integer()) :: :ok
  def record_error(tool_name, elapsed_us) do
    ensure_table()
    now = DateTime.utc_now()

    :ets.update_counter(
      @table,
      tool_name,
      [{2, 1}, {3, elapsed_us}, {4, 1}],
      {tool_name, 0, 0, 0, nil}
    )

    case :ets.lookup(@table, tool_name) do
      [{^tool_name, count, total, errors, _last}] ->
        :ets.insert(@table, {tool_name, count, total, errors, now})

      _ ->
        :ets.insert(@table, {tool_name, 1, elapsed_us, 1, now})
    end

    :ok
  end

  @doc """
  Get stats for all tools or a specific tool.

  ## Options
  - `:sort_by` -- `:count` | `:avg_time` | `:total_time` (default `:count`)
  - `:period` -- `:all` | `:today` | `:last_hour` (default `:all`)
  """
  @spec get_stats(keyword()) :: [map()]
  def get_stats(opts \\ []) do
    ensure_table()
    sort_by = Keyword.get(opts, :sort_by, :count)

    @table
    |> :ets.tab2list()
    |> Enum.map(fn {name, count, total_us, errors, last} ->
      avg = if count > 0, do: div(total_us, count), else: 0

      %{
        tool: name,
        count: count,
        total_time_us: total_us,
        avg_time_us: avg,
        error_count: errors,
        last_invoked: last
      }
    end)
    |> Enum.sort_by(fn stat ->
      case sort_by do
        :count -> -stat.count
        :avg_time -> -stat.avg_time_us
        :total_time -> -stat.total_time_us
        _ -> -stat.count
      end
    end)
  end

  @doc "Get stats for a single tool."
  @spec get_tool_stats(String.t()) :: map() | nil
  def get_tool_stats(tool_name) do
    ensure_table()

    case :ets.lookup(@table, tool_name) do
      [{^tool_name, count, total_us, errors, last}] ->
        avg = if count > 0, do: div(total_us, count), else: 0

        %{
          tool: tool_name,
          count: count,
          total_time_us: total_us,
          avg_time_us: avg,
          error_count: errors,
          last_invoked: last
        }

      [] ->
        nil
    end
  end

  @doc "Clear all telemetry data."
  @spec reset() :: :ok
  def reset do
    ensure_table()
    :ets.delete_all_objects(@table)
    :ok
  end

  @doc "Return total invocations across all tools."
  @spec total_invocations() :: non_neg_integer()
  def total_invocations do
    ensure_table()

    :ets.foldl(fn {_name, count, _total, _errors, _last}, acc -> acc + count end, 0, @table)
  end

  # ── Server callbacks ─────────────────────────────────────────────────

  @impl true
  def init(_opts) do
    ensure_table()
    load_persisted()
    schedule_persist()
    {:ok, %{}}
  end

  @impl true
  def handle_info(:persist, state) do
    persist()
    schedule_persist()
    {:noreply, state}
  end

  @impl true
  def terminate(_reason, _state) do
    persist()
    :ok
  end

  # ── Private ──────────────────────────────────────────────────────────

  defp ensure_table do
    if :ets.whereis(@table) == :undefined do
      :ets.new(@table, [
        :named_table,
        :set,
        :public,
        read_concurrency: true,
        write_concurrency: true
      ])
    end
  end

  defp schedule_persist do
    Process.send_after(self(), :persist, @persist_interval)
  end

  defp persist do
    path = persistence_path()
    File.mkdir_p!(Path.dirname(path))
    data = :ets.tab2list(@table)
    File.write!(path, :erlang.term_to_binary(data))
  rescue
    e -> Logger.warning("Failed to persist telemetry: #{Exception.message(e)}")
  end

  defp load_persisted do
    path = persistence_path()

    if File.exists?(path) do
      data = path |> File.read!() |> :erlang.binary_to_term()

      Enum.each(data, fn entry ->
        :ets.insert(@table, entry)
      end)

      Logger.debug("Loaded #{length(data)} telemetry entries")
    end
  rescue
    e -> Logger.warning("Failed to load telemetry: #{Exception.message(e)}")
  end

  defp persistence_path do
    hash =
      File.cwd!()
      |> then(&:crypto.hash(:sha256, &1))
      |> Base.encode16(case: :lower)
      |> String.slice(0, 12)

    Path.join([System.user_home!(), ".ragex", "telemetry", "#{hash}.etf"])
  end
end