ex_claw

defmodule ExClaw.Usage.Forecaster do
  @moduledoc """
  Forecasts token usage based on recent consumption patterns.
  Simple linear projection.
  """

  @doc """
  Estimates how many more requests fit in the remaining budget.
  Based on average tokens per request.
  """
  def estimate_remaining(tracker_name \\ ExClaw.Usage.Tracker, budget_config) do
    records = ExClaw.Usage.Tracker.all_records(tracker_name)
    %{input_tokens: input, output_tokens: output} = ExClaw.Usage.Tracker.total(tracker_name)
    total_used = input + output
    remaining = max(budget_config.max_tokens - total_used, 0)

    case length(records) do
      0 ->
        :infinity

      n ->
        avg_per_request = total_used / n
        if avg_per_request > 0, do: trunc(remaining / avg_per_request), else: :infinity
    end
  end

  @doc """
  Calculates tokens per minute over recent usage.
  Returns `{:ok, tokens_per_minute}` or `:no_data`.
  """
  def burn_rate(tracker_name \\ ExClaw.Usage.Tracker) do
    records = ExClaw.Usage.Tracker.all_records(tracker_name)

    case records do
      [] ->
        :no_data

      [_single] ->
        :no_data

      records ->
        first_ts = List.first(records).recorded_at
        last_ts = List.last(records).recorded_at
        elapsed_ms = last_ts - first_ts

        if elapsed_ms <= 0 do
          :no_data
        else
          total_tokens =
            Enum.reduce(records, 0, fn r, acc ->
              acc + r.input_tokens + r.output_tokens
            end)

          elapsed_minutes = elapsed_ms / 60_000
          {:ok, total_tokens / elapsed_minutes}
        end
    end
  end

  @doc """
  Estimates minutes until budget exhaustion based on burn rate.
  Returns `{:ok, minutes}` or `:no_data` or `:within_budget`.
  """
  def time_to_exhaustion(tracker_name \\ ExClaw.Usage.Tracker, budget_config) do
    case burn_rate(tracker_name) do
      :no_data ->
        :no_data

      {:ok, rate} when rate > 0 ->
        %{input_tokens: input, output_tokens: output} = ExClaw.Usage.Tracker.total(tracker_name)
        remaining = budget_config.max_tokens - (input + output)

        if remaining <= 0 do
          {:ok, 0.0}
        else
          {:ok, remaining / rate}
        end

      {:ok, _} ->
        :no_data
    end
  end
end