cmdc

defmodule CMDC.Plugin.Builtin.CostGuard do
  @moduledoc """
  P1 成本硬限制插件 — Token 和成本预算控制。

  在每次 LLM 请求前（`before_request`）检查累计 token/cost 是否超预算。
  超限时立即 abort Agent，防止成本失控。

  ## 配置

      {CMDC.Plugin.Builtin.CostGuard,
        max_tokens: 100_000,     # 单次会话最大 token（nil 不限制）
        max_cost_usd: 1.0,       # 单次会话最大成本 USD（nil 不限制）
        warn_threshold: 0.8      # 用量达 80% 时 emit 告警（0.0-1.0）
      }

  ## 触发 Hook

  `{:before_request, messages}` — 发送 LLM 请求前触发。
  """

  @behaviour CMDC.Plugin

  require Logger

  @default_warn_threshold 0.8

  # ==========================================================================
  # Callbacks
  # ==========================================================================

  @impl true
  def init(opts) do
    state = %{
      max_tokens: Keyword.get(opts, :max_tokens),
      max_cost_usd: Keyword.get(opts, :max_cost_usd),
      warn_threshold: Keyword.get(opts, :warn_threshold, @default_warn_threshold),
      warned: false
    }

    {:ok, state}
  end

  @impl true
  def priority, do: 120

  @impl true
  def handle_event({:before_request, _messages}, state, ctx) do
    total_tokens = Map.get(ctx, :total_tokens, 0)
    cost_usd = Map.get(ctx, :cost_usd, 0.0)

    cond do
      token_exceeded?(state, total_tokens) ->
        Logger.error("[CostGuard] Token 超预算: #{total_tokens}/#{state.max_tokens}")
        {:abort, "Token 预算已耗尽（#{total_tokens}/#{state.max_tokens}）", state}

      cost_exceeded?(state, cost_usd) ->
        Logger.error("[CostGuard] 成本超预算: $#{cost_usd}/$#{state.max_cost_usd}")
        {:abort, "成本预算已耗尽（$#{cost_usd}/$#{state.max_cost_usd}）", state}

      not state.warned and near_limit?(state, total_tokens, cost_usd) ->
        Logger.warning("[CostGuard] 用量接近预算限制")

        {:emit, {:cost_warning, %{total_tokens: total_tokens, cost_usd: cost_usd}},
         %{state | warned: true}}

      true ->
        {:continue, state}
    end
  end

  def handle_event(_event, state, _ctx), do: {:continue, state}

  # ==========================================================================
  # Private
  # ==========================================================================

  defp token_exceeded?(%{max_tokens: nil}, _), do: false
  defp token_exceeded?(%{max_tokens: max}, current), do: current >= max

  defp cost_exceeded?(%{max_cost_usd: nil}, _), do: false
  defp cost_exceeded?(%{max_cost_usd: max}, current), do: current >= max

  defp near_limit?(state, tokens, cost) do
    threshold = state.warn_threshold

    token_near =
      state.max_tokens != nil and tokens >= state.max_tokens * threshold

    cost_near =
      state.max_cost_usd != nil and cost >= state.max_cost_usd * threshold

    token_near or cost_near
  end
end