defmodule Baton.Pricing do
  @moduledoc """
  Behaviour for computing the cost of an LLM call.

  Hardcoded prices go stale fast, so Baton does not ship authoritative
  prices. Instead it defines this behaviour and calls whichever module the host
  configures:

      config :baton, pricing: MyApp.LLMPricing

  Your implementation receives a usage map and returns a cost as `Decimal.t()`,
  or `nil` if the model is unknown (cost is then recorded as nil rather than a
  wrong number).

      defmodule MyApp.LLMPricing do
        @behaviour Baton.Pricing

        @impl true
        def cost(%{model: "claude-sonnet-4-" <> _, input_tokens: i, output_tokens: o}) do
          Decimal.add(
            Decimal.mult(Decimal.new(i), Decimal.from_float(3.0 / 1_000_000)),
            Decimal.mult(Decimal.new(o), Decimal.from_float(15.0 / 1_000_000))
          )
        end

        def cost(_usage), do: nil
      end

  `Baton.Pricing.Default` is provided as a starting point and for tests, but
  you should supply your own and keep it current.
  """

  @type usage :: %{
          required(:model) => String.t() | nil,
          optional(:input_tokens) => non_neg_integer(),
          optional(:output_tokens) => non_neg_integer(),
          optional(:cache_read_tokens) => non_neg_integer(),
          optional(:cache_write_tokens) => non_neg_integer()
        }

  @doc "Return the cost in USD for a usage map, or nil if it can't be priced."
  @callback cost(usage()) :: Decimal.t() | nil

  @doc """
  Compute cost via the configured pricing module.

  Returns `{:ok, Decimal.t()}`, or `{:error, :unpriced}` when the module
  returns nil (unknown model, missing data, etc.).
  """
  @spec cost(usage()) :: {:ok, Decimal.t()} | {:error, :unpriced}
  def cost(usage) do
    case Baton.Config.pricing().cost(usage) do
      %Decimal{} = d -> {:ok, d}
      nil -> {:error, :unpriced}
    end
  end
end

defmodule Baton.Pricing.Default do
  @moduledoc """
  A reference `Baton.Pricing` implementation with a small built-in price
  table (USD per million tokens).

  **Verify these against current provider pricing before relying on them for
  billing.** They are intended as a starting point and for tests. Override by
  configuring your own module: `config :baton, pricing: MyApp.LLMPricing`.
  """

  @behaviour Baton.Pricing

  # USD per million tokens. Update or replace with your own module.
  @prices %{
    "claude-opus-4-20250514" => %{input: 15.0, output: 75.0, cache_read: 1.5, cache_write: 18.75},
    "claude-sonnet-4-20250514" => %{input: 3.0, output: 15.0, cache_read: 0.3, cache_write: 3.75},
    "claude-haiku-4-20250514" => %{input: 0.8, output: 4.0, cache_read: 0.08, cache_write: 1.0},
    "gpt-4o" => %{input: 2.5, output: 10.0, cache_read: 1.25, cache_write: nil},
    "gpt-4o-mini" => %{input: 0.15, output: 0.6, cache_read: 0.075, cache_write: nil},
    "o3" => %{input: 10.0, output: 40.0, cache_read: 2.5, cache_write: nil},
    "o4-mini" => %{input: 1.1, output: 4.4, cache_read: 0.275, cache_write: nil}
  }

  @impl Baton.Pricing
  def cost(%{model: model} = usage) when is_binary(model) do
    case Map.get(@prices, model) do
      nil ->
        nil

      rates ->
        [
          line(usage, :input_tokens, rates.input),
          line(usage, :output_tokens, rates.output),
          line(usage, :cache_read_tokens, rates[:cache_read]),
          line(usage, :cache_write_tokens, rates[:cache_write])
        ]
        |> Enum.reject(&is_nil/1)
        |> Enum.reduce(Decimal.new(0), &Decimal.add/2)
    end
  end

  def cost(_usage), do: nil

  @doc "The built-in price table (for display/inspection)."
  def prices, do: @prices

  defp line(usage, key, rate) do
    count = Map.get(usage, key, 0) || 0

    cond do
      count == 0 -> Decimal.new(0)
      is_nil(rate) -> nil
      true -> Decimal.mult(Decimal.new(count), Decimal.from_float(rate / 1_000_000))
    end
  end
end