cmdc_gateway

defmodule CMDCGateway.Meter do
  @moduledoc """
  Usage Metering — per api_key 用量统计。

  v0.2 起在 `init/1` 中通过 `CMDC.EventBus.subscribe_all/0` **真正订阅** 所有
  session 的事件流，自动响应：

  - `{:agent_end, _messages, %CMDC.TokenUsage{}}` → 累计三类 token + cost + cached
  - `{:agent_end, _messages, raw_map}` → 同上（兼容老 Provider 未归一化的 raw）

  注意：`prompt_count` 在 Router 的 `POST /v1/sessions/:id/prompt` 端点同步 +1，
  不通过 EventBus 监听 `:prompt_received`（避免 Router 手动计数 + 自动监听双重计数）。

  api_key 通过 `CMDCGateway.SessionStore.get/1` 反查。

  ## 统计指标

  - `:prompt_count` — 总 prompt 请求数
  - `:total_prompt_tokens` — 累计输入 token 数
  - `:total_completion_tokens` — 累计输出 token 数
  - `:total_tokens` — 累计总 token 数
  - `:total_cost_usd` — 累计美元成本（v0.2 新增）
  - `:cached_tokens` — 累计 prompt cache 命中 token 数（v0.2 新增）
  - `:last_activity_at` — 最后活跃时间戳（毫秒）

  ## 查询

      CMDCGateway.Meter.get_usage("my-api-key")
      # => %{prompt_count: 42, total_tokens: 128_000, ...}
  """

  use GenServer
  require Logger

  alias CMDCGateway.SessionStore

  @table __MODULE__

  @type usage :: %{
          prompt_count: non_neg_integer(),
          total_prompt_tokens: non_neg_integer(),
          total_completion_tokens: non_neg_integer(),
          total_tokens: non_neg_integer(),
          total_cost_usd: float(),
          cached_tokens: non_neg_integer(),
          last_activity_at: integer()
        }

  # ==========================================================================
  # Public API
  # ==========================================================================

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

  @doc "获取某个 api_key 的累计用量。"
  @spec get_usage(String.t()) :: usage()
  def get_usage(api_key) do
    case :ets.lookup(@table, api_key) do
      [{^api_key, usage}] -> usage
      [] -> empty_usage()
    end
  end

  @doc "获取所有 api_key 的用量汇总。"
  @spec all_usage() :: [{String.t(), usage()}]
  def all_usage, do: :ets.tab2list(@table)

  @doc """
  记录一次 prompt 请求。

  v0.2 起 Meter 已自动通过 EventBus 监听 `:prompt_received` 事件；保留本函数
  仅用于 Router HTTP 层的同步计数（请求被接收即增加，不必等 Agent 真正开始处理）。
  """
  @spec record_prompt(String.t()) :: :ok
  def record_prompt(api_key) do
    now = System.system_time(:millisecond)
    update_usage(api_key, &%{&1 | prompt_count: &1.prompt_count + 1, last_activity_at: now})
    :ok
  end

  @doc """
  手动记录一次 token 用量（兼容老调用）。

  v0.2 起 Meter 已自动通过 EventBus 监听 `:agent_end`，本函数保留为外部
  自定义场景使用。
  """
  @spec record_tokens(String.t(), map() | CMDC.TokenUsage.t()) :: :ok
  def record_tokens(api_key, usage), do: do_record_tokens(api_key, usage)

  @doc "重置某个 api_key 的用量。"
  @spec reset(String.t()) :: :ok
  def reset(api_key) do
    :ets.delete(@table, api_key)
    :ok
  end

  @doc "重置所有用量数据。"
  @spec reset_all() :: :ok
  def reset_all do
    :ets.delete_all_objects(@table)
    :ok
  end

  # ==========================================================================
  # GenServer Callbacks
  # ==========================================================================

  @impl true
  def init(_opts) do
    :ets.new(@table, [
      :named_table,
      :set,
      :public,
      read_concurrency: true,
      write_concurrency: true
    ])

    case CMDC.EventBus.subscribe_all() do
      {:ok, _} ->
        Logger.info(
          "[CMDCGateway.Meter] EventBus.subscribe_all/0 succeeded — auto-metering active"
        )

      {:error, {:already_registered, _}} ->
        Logger.info("[CMDCGateway.Meter] EventBus.subscribe_all/0 already registered")
    end

    {:ok, %{}}
  end

  @impl true
  def handle_info({:cmdc_event, session_id, event}, state) do
    handle_event(session_id, event)
    {:noreply, state}
  end

  def handle_info(_msg, state), do: {:noreply, state}

  # ==========================================================================
  # Private — 事件处理
  # ==========================================================================

  defp handle_event(session_id, {:agent_end, _messages, %CMDC.TokenUsage{} = usage}) do
    with {:ok, %{api_key: api_key}} <- SessionStore.get(session_id) do
      do_record_tokens(api_key, usage)
    end
  end

  defp handle_event(session_id, {:agent_end, _messages, usage}) when is_map(usage) do
    with {:ok, %{api_key: api_key}} <- SessionStore.get(session_id) do
      do_record_tokens(api_key, usage)
    end
  end

  defp handle_event(_sid, _event), do: :ok

  # ==========================================================================
  # Private — 写入
  # ==========================================================================

  defp do_record_tokens(api_key, %CMDC.TokenUsage{} = u) do
    update_usage(api_key, fn existing ->
      now = System.system_time(:millisecond)

      %{
        existing
        | total_prompt_tokens: existing.total_prompt_tokens + u.prompt_tokens,
          total_completion_tokens: existing.total_completion_tokens + u.completion_tokens,
          total_tokens: existing.total_tokens + u.total_tokens,
          total_cost_usd: existing.total_cost_usd + (u.cost_usd || 0.0),
          cached_tokens: existing.cached_tokens + (u.cached_tokens || 0),
          last_activity_at: now
      }
    end)

    :ok
  end

  defp do_record_tokens(api_key, raw) when is_map(raw) do
    do_record_tokens(api_key, CMDC.TokenUsage.from_raw(raw))
  end

  defp update_usage(api_key, fun) do
    existing =
      case :ets.lookup(@table, api_key) do
        [{^api_key, u}] -> u
        [] -> empty_usage()
      end

    :ets.insert(@table, {api_key, fun.(existing)})
  end

  defp empty_usage do
    %{
      prompt_count: 0,
      total_prompt_tokens: 0,
      total_completion_tokens: 0,
      total_tokens: 0,
      total_cost_usd: 0.0,
      cached_tokens: 0,
      last_activity_at: 0
    }
  end
end