cmdc

defmodule CMDC.TokenUsage do
  @moduledoc """
  公开的 token 用量 struct。

  CMDC 在 Provider 层把不同 LLM 厂商的 raw usage（OpenAI 风格 `prompt_tokens` /
  `completion_tokens`，Anthropic 风格 `input_tokens` / `output_tokens`，
  Google 风格 `prompt_token_count` / `candidates_token_count`）统一归一化为本 struct。

  事件协议：

  - `{:agent_end, messages, %CMDC.TokenUsage{}}` — 第三参数是 struct
  - `CMDC.status/1` 返回 `:token_usage` 字段为 struct，原 `:total_tokens` Int 字段保留

  ## 字段约定（OpenAI 行业事实标准）

  - `:prompt_tokens` — 输入 token 总数（Anthropic 的 input_tokens 也归到这里）
  - `:completion_tokens` — 输出 token 总数（Anthropic 的 output_tokens 也归到这里）
  - `:total_tokens` — 累计 token 总数（多数 Provider 自带；缺失时由 prompt + completion 补齐）
  - `:cost_usd` — 累计美元成本（基于 token 用量与模型定价计算；Provider 未配置定价时为 `nil`）
  - `:cached_tokens` — Anthropic prompt cache 命中的 token 数（Provider 不支持时为 `nil`）

  ## 使用

      iex> CMDC.TokenUsage.new(prompt_tokens: 100, completion_tokens: 50)
      %CMDC.TokenUsage{
        prompt_tokens: 100,
        completion_tokens: 50,
        total_tokens: 150,
        cost_usd: nil,
        cached_tokens: nil
      }

      # 从 Anthropic raw usage 归一化（total_tokens 缺失自动补齐）
      iex> u = CMDC.TokenUsage.from_raw(%{input_tokens: 80, output_tokens: 40})
      iex> {u.prompt_tokens, u.completion_tokens, u.total_tokens}
      {80, 40, 120}
  """

  @derive Jason.Encoder
  @enforce_keys [:prompt_tokens, :completion_tokens, :total_tokens]
  defstruct [
    :prompt_tokens,
    :completion_tokens,
    :total_tokens,
    cost_usd: nil,
    cached_tokens: nil
  ]

  @type t :: %__MODULE__{
          prompt_tokens: non_neg_integer(),
          completion_tokens: non_neg_integer(),
          total_tokens: non_neg_integer(),
          cost_usd: float() | nil,
          cached_tokens: non_neg_integer() | nil
        }

  # ==========================================================================
  # 构造
  # ==========================================================================

  @doc """
  从 keyword / map 构造 TokenUsage struct。

  缺省字段：`prompt_tokens` / `completion_tokens` / `total_tokens` 默认 0,
  `cost_usd` / `cached_tokens` 默认 `nil`。

  ## 示例

      iex> CMDC.TokenUsage.new()
      %CMDC.TokenUsage{prompt_tokens: 0, completion_tokens: 0, total_tokens: 0,
                       cost_usd: nil, cached_tokens: nil}

      iex> CMDC.TokenUsage.new(prompt_tokens: 10, completion_tokens: 5, cost_usd: 0.0001)
      %CMDC.TokenUsage{prompt_tokens: 10, completion_tokens: 5, total_tokens: 15,
                       cost_usd: 0.0001, cached_tokens: nil}
  """
  @spec new(map() | keyword()) :: t()
  def new(fields \\ []) do
    fields = Map.new(fields)

    prompt = fetch_int(fields, :prompt_tokens, 0)
    completion = fetch_int(fields, :completion_tokens, 0)
    total = fetch_int(fields, :total_tokens, prompt + completion)

    %__MODULE__{
      prompt_tokens: prompt,
      completion_tokens: completion,
      total_tokens: total,
      cost_usd: Map.get(fields, :cost_usd),
      cached_tokens: Map.get(fields, :cached_tokens)
    }
  end

  @doc """
  从任意 Provider raw usage map 构造（Anthropic / OpenAI / Google 兼容）。

  字段映射：

  - `prompt_tokens` ← `:prompt_tokens` / `:input_tokens` / `:prompt_token_count`
  - `completion_tokens` ← `:completion_tokens` / `:output_tokens` / `:candidates_token_count`
  - `total_tokens` ← `:total_tokens`（缺失时由 prompt + completion 补齐）
  - `cached_tokens` ← `:cache_read_input_tokens` / `:cached_tokens`

  支持 atom 或 string key 的混合输入。
  """
  @spec from_raw(map() | t()) :: t()
  def from_raw(%__MODULE__{} = already), do: already

  def from_raw(raw) when is_map(raw) do
    prompt = first_int(raw, [:prompt_tokens, :input_tokens, :prompt_token_count]) || 0

    completion =
      first_int(raw, [:completion_tokens, :output_tokens, :candidates_token_count]) || 0

    total = first_int(raw, [:total_tokens]) || prompt + completion
    cached = first_int(raw, [:cache_read_input_tokens, :cached_tokens])

    %__MODULE__{
      prompt_tokens: prompt,
      completion_tokens: completion,
      total_tokens: total,
      cost_usd: nil,
      cached_tokens: cached
    }
  end

  @doc """
  累加另一份 token usage（用于多轮 / 多 Provider 调用聚合）。

  - `prompt_tokens` / `completion_tokens` / `total_tokens` / `cached_tokens` 取加和
    （nil 视为 0）
  - `cost_usd` 也取加和（双方都是数字时累加；任意一方为 nil 时取另一方的值）

  ## 示例

      iex> a = CMDC.TokenUsage.new(prompt_tokens: 10, completion_tokens: 5)
      iex> b = CMDC.TokenUsage.new(prompt_tokens: 20, completion_tokens: 10)
      iex> CMDC.TokenUsage.add(a, b)
      %CMDC.TokenUsage{prompt_tokens: 30, completion_tokens: 15, total_tokens: 45,
                       cost_usd: nil, cached_tokens: nil}
  """
  @spec add(t(), t()) :: t()
  def add(%__MODULE__{} = a, %__MODULE__{} = b) do
    %__MODULE__{
      prompt_tokens: a.prompt_tokens + b.prompt_tokens,
      completion_tokens: a.completion_tokens + b.completion_tokens,
      total_tokens: a.total_tokens + b.total_tokens,
      cost_usd: add_optional(a.cost_usd, b.cost_usd),
      cached_tokens: add_optional_int(a.cached_tokens, b.cached_tokens)
    }
  end

  @doc """
  把 Agent State 内部 token map + cost_usd 转换为公开 struct。

  Agent 内部为兼容性保留扩展 map（含 `:context_window` / `:current_context_tokens`,
  这两个字段是上下文窗口元数据，不属于"用量"，不进入公开 struct）。
  """
  @spec from_state_token_map(map(), float() | nil, non_neg_integer() | nil) :: t()
  def from_state_token_map(token_map, cost_usd \\ nil, cached_tokens \\ nil) do
    %__MODULE__{
      prompt_tokens: Map.get(token_map, :prompt_tokens, 0),
      completion_tokens: Map.get(token_map, :completion_tokens, 0),
      total_tokens: Map.get(token_map, :total_tokens, 0),
      cost_usd: cost_usd,
      cached_tokens: cached_tokens
    }
  end

  # ==========================================================================
  # 私有辅助
  # ==========================================================================

  defp fetch_int(map, key, default) do
    case Map.get(map, key, default) do
      n when is_integer(n) -> n
      _ -> default
    end
  end

  defp first_int(usage, keys) do
    Enum.find_value(keys, fn key ->
      val = Map.get(usage, key) || Map.get(usage, to_string(key))
      if is_integer(val), do: val, else: nil
    end)
  end

  defp add_optional(nil, b), do: b
  defp add_optional(a, nil), do: a
  defp add_optional(a, b) when is_number(a) and is_number(b), do: a + b

  defp add_optional_int(nil, nil), do: nil
  defp add_optional_int(nil, b), do: b
  defp add_optional_int(a, nil), do: a
  defp add_optional_int(a, b) when is_integer(a) and is_integer(b), do: a + b
end