syntropy

defmodule Syntropy.LLMClient.Usage do
  @moduledoc """
  Provider usage accounting for hosted model calls.

  Usage events are logged for Railway inspection and recorded in
  `UsageRecorder`. Completed run envelopes persist the summarized usage when
  Postgres run-envelope persistence is enabled.
  """

  require Logger

  alias Syntropy.LLMClient.Config
  alias Syntropy.LLMClient.UsageRecorder

  @deepseek_v4_pro_pricing %{
    input_usd_per_1m: 2.10,
    cached_input_usd_per_1m: 0.20,
    output_usd_per_1m: 4.40
  }

  @known_pricing %{
    "deepseek-ai/deepseek-v4-pro" => @deepseek_v4_pro_pricing
  }

  @type usage :: %{
          input_tokens: non_neg_integer(),
          cached_input_tokens: non_neg_integer(),
          output_tokens: non_neg_integer(),
          total_tokens: non_neg_integer()
        }

  @spec log_completion(Config.t(), map(), map(), non_neg_integer()) :: :ok
  def log_completion(%Config{} = config, agent, response_body, duration_ms)
      when is_map(response_body) do
    case usage_from_body(response_body) do
      nil ->
        :ok

      usage ->
        event =
          %{
            event: "syntropy_llm_usage",
            provider: Config.provider_name(config),
            model: config.model,
            profile: config.profile,
            agent_id: agent_value(agent, :id),
            runtime_id: agent_value(agent, :runtime_id),
            node_id: agent_value(agent, :node_id),
            perspective: agent_value(agent, :perspective),
            task_id: Process.get(:syntropy_task_id),
            stage: Process.get(:syntropy_llm_stage),
            duration_ms: duration_ms,
            usage: usage,
            pricing: pricing_for_model(config.model),
            estimated_cost_usd: estimated_cost_usd(config.model, usage)
          }

        :ok = UsageRecorder.record(event)
        Logger.info(Jason.encode!(event))
    end

    :ok
  end

  @spec usage_from_body(map()) :: usage() | nil
  def usage_from_body(%{"usage" => usage}) when is_map(usage) do
    input_tokens = integer_field(usage, ["prompt_tokens", "input_tokens"])
    output_tokens = integer_field(usage, ["completion_tokens", "output_tokens"])
    cached_input_tokens = cached_input_tokens(usage)
    total_tokens = integer_field(usage, ["total_tokens"]) || input_tokens + output_tokens

    %{
      input_tokens: input_tokens,
      cached_input_tokens: min(cached_input_tokens, input_tokens),
      output_tokens: output_tokens,
      total_tokens: total_tokens
    }
  end

  def usage_from_body(_body), do: nil

  @spec estimated_cost_usd(String.t() | nil, usage() | nil) :: float() | nil
  def estimated_cost_usd(_model, nil), do: nil

  def estimated_cost_usd(model, usage) do
    case pricing_for_model(model) do
      nil ->
        nil

      pricing ->
        billable_input_tokens =
          max(usage.input_tokens - usage.cached_input_tokens, 0)

        billable_input_tokens * pricing.input_usd_per_1m / 1_000_000 +
          usage.cached_input_tokens * pricing.cached_input_usd_per_1m / 1_000_000 +
          usage.output_tokens * pricing.output_usd_per_1m / 1_000_000
    end
  end

  @spec pricing_for_model(String.t() | nil) :: map() | nil
  def pricing_for_model(model) when is_binary(model) do
    @known_pricing[String.downcase(model)]
  end

  def pricing_for_model(_model), do: nil

  defp cached_input_tokens(usage) do
    cond do
      is_map(usage["prompt_tokens_details"]) ->
        integer_field(usage["prompt_tokens_details"], ["cached_tokens"])

      is_map(usage["input_tokens_details"]) ->
        integer_field(usage["input_tokens_details"], ["cached_tokens"])

      true ->
        0
    end
  end

  defp integer_field(map, keys) do
    keys
    |> Enum.find_value(fn key ->
      case Map.get(map, key) do
        value when is_integer(value) and value >= 0 -> value
        value when is_float(value) and value >= 0 -> trunc(value)
        _other -> nil
      end
    end)
    |> Kernel.||(0)
  end

  defp agent_value(agent, key) when is_map(agent) do
    Map.get(agent, key) || Map.get(agent, Atom.to_string(key))
  end

  defp agent_value(_agent, _key), do: nil
end