syntropy

defmodule Syntropy.LLMClient.UsageRecorder do
  @moduledoc """
  In-memory per-run LLM usage summaries.

  Provider usage is also emitted to logs, but this GenServer keeps in-flight
  run-level cost and token accounting available to `/api/runs`. Completed run
  envelopes include the summary so persistence can retain the final accounting.
  """

  use GenServer

  @history_limit 200

  @type usage_event :: map()
  @type summary :: map()

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

  @spec record(usage_event()) :: :ok
  def record(%{task_id: task_id} = event) when is_binary(task_id) and task_id != "" do
    if Process.whereis(__MODULE__) do
      GenServer.cast(__MODULE__, {:record, stringify_keys(event)})
    end

    :ok
  end

  def record(_event), do: :ok

  @spec summary(String.t() | nil) :: summary() | nil
  def summary(nil), do: nil

  def summary(task_id) when is_binary(task_id) do
    if Process.whereis(__MODULE__) do
      GenServer.call(__MODULE__, {:summary, task_id})
    end
  end

  @spec reset() :: :ok
  def reset do
    if Process.whereis(__MODULE__) do
      GenServer.call(__MODULE__, :reset)
    else
      :ok
    end
  end

  @impl true
  def init(_opts), do: {:ok, %{events_by_task: %{}, order: []}}

  @impl true
  def handle_cast({:record, %{"task_id" => task_id} = event}, state) do
    events = [event | Map.get(state.events_by_task, task_id, [])]

    order =
      [task_id | Enum.reject(state.order, &(&1 == task_id))]
      |> Enum.take(@history_limit)

    events_by_task =
      state.events_by_task
      |> Map.put(task_id, events)
      |> Map.take(order)

    {:noreply, %{state | events_by_task: events_by_task, order: order}}
  end

  @impl true
  def handle_call({:summary, task_id}, _from, state) do
    summary =
      state.events_by_task
      |> Map.get(task_id, [])
      |> Enum.reverse()
      |> summarize()

    {:reply, summary, state}
  end

  def handle_call(:reset, _from, _state) do
    {:reply, :ok, %{events_by_task: %{}, order: []}}
  end

  defp summarize([]), do: nil

  defp summarize(events) do
    totals =
      Enum.reduce(events, empty_totals(), fn event, acc ->
        usage = Map.get(event, "usage", %{})

        acc
        |> add(:input_tokens, int(usage["input_tokens"]))
        |> add(:cached_input_tokens, int(usage["cached_input_tokens"]))
        |> add(:output_tokens, int(usage["output_tokens"]))
        |> add(:total_tokens, int(usage["total_tokens"]))
        |> add_cost(event["estimated_cost_usd"])
      end)

    first = List.first(events)

    %{
      provider: first["provider"],
      model: first["model"],
      profile: first["profile"],
      call_count: length(events),
      input_tokens: totals.input_tokens,
      cached_input_tokens: totals.cached_input_tokens,
      output_tokens: totals.output_tokens,
      total_tokens: totals.total_tokens,
      estimated_cost_usd: totals.estimated_cost_usd,
      stages: stage_summaries(events),
      calls: Enum.map(events, &call_summary/1)
    }
  end

  defp empty_totals do
    %{
      input_tokens: 0,
      cached_input_tokens: 0,
      output_tokens: 0,
      total_tokens: 0,
      estimated_cost_usd: nil
    }
  end

  defp stage_summaries(events) do
    events
    |> Enum.group_by(&(&1["stage"] || "unknown"))
    |> Enum.map(fn {stage, stage_events} ->
      totals =
        Enum.reduce(stage_events, empty_totals(), fn event, acc ->
          usage = Map.get(event, "usage", %{})

          acc
          |> add(:input_tokens, int(usage["input_tokens"]))
          |> add(:cached_input_tokens, int(usage["cached_input_tokens"]))
          |> add(:output_tokens, int(usage["output_tokens"]))
          |> add(:total_tokens, int(usage["total_tokens"]))
          |> add_cost(event["estimated_cost_usd"])
        end)

      %{
        stage: stage,
        call_count: length(stage_events),
        input_tokens: totals.input_tokens,
        cached_input_tokens: totals.cached_input_tokens,
        output_tokens: totals.output_tokens,
        total_tokens: totals.total_tokens,
        estimated_cost_usd: totals.estimated_cost_usd
      }
    end)
    |> Enum.sort_by(& &1.stage)
  end

  defp call_summary(event) do
    usage = Map.get(event, "usage", %{})

    %{
      stage: event["stage"] || "unknown",
      agent_id: event["agent_id"],
      runtime_id: event["runtime_id"],
      perspective: event["perspective"],
      duration_ms: event["duration_ms"],
      input_tokens: int(usage["input_tokens"]),
      cached_input_tokens: int(usage["cached_input_tokens"]),
      output_tokens: int(usage["output_tokens"]),
      total_tokens: int(usage["total_tokens"]),
      estimated_cost_usd: event["estimated_cost_usd"]
    }
  end

  defp add(totals, key, value), do: Map.update!(totals, key, &(&1 + value))

  defp add_cost(totals, value) when is_number(value) do
    Map.update!(totals, :estimated_cost_usd, fn
      nil -> value
      current -> current + value
    end)
  end

  defp add_cost(totals, _value), do: totals

  defp int(value) when is_integer(value) and value >= 0, do: value
  defp int(value) when is_float(value) and value >= 0, do: trunc(value)
  defp int(_value), do: 0

  defp stringify_keys(%{} = value) do
    Map.new(value, fn {key, nested_value} -> {to_string(key), stringify_keys(nested_value)} end)
  end

  defp stringify_keys(value) when is_list(value), do: Enum.map(value, &stringify_keys/1)
  defp stringify_keys(value), do: value
end