/**
 * Session token budget for eval cells.
 *
 * Port of pi-subagents' `RolloutBudget` (itself a port of Codex's
 * `RolloutBudget`) for the pi-codemode session. One in-memory budget
 * accumulates weighted spend across all cells in one session; the eval tool
 * rejects a new cell once the spend reaches the limit.
 *
 * Weighted spend per cell:
 *   outputTokens * samplingWeight + (inputTokens - cacheReadTokens) * prefillWeight
 * When provider token counts are absent, the cell text length (`textLength /
 * 4`) stands in for the token count and is weighted at the sampling rate.
 * A limit of zero means unlimited.
 */

export interface SessionBudgetConfig {
  /** Cumulative weighted spend limit in tokens. Zero means unlimited. */
  readonly limitTokens: number;
  /** Multiplier for the prefill (non-cached input) side. Defaults to 1. */
  readonly prefillWeight?: number;
  /** Multiplier for the sampling (output) side. Defaults to 1. */
  readonly samplingWeight?: number;
}

export interface SessionUsage {
  readonly cacheReadTokens?: number;
  readonly inputTokens?: number;
  readonly outputTokens?: number;
  readonly textLength?: number;
}

const FALLBACK_CHARS_PER_TOKEN = 4;

export class SessionBudget {
  readonly #limitTokens: number;
  readonly #prefillWeight: number;
  readonly #samplingWeight: number;
  #exhausted = false;
  #weightedTokensUsed = 0;

  constructor(config: SessionBudgetConfig) {
    if (!Number.isFinite(config.limitTokens) || config.limitTokens < 0) {
      throw new Error("budget limit must be a non-negative number");
    }
    const prefillWeight = config.prefillWeight ?? 1;
    const samplingWeight = config.samplingWeight ?? 1;
    if (!Number.isFinite(prefillWeight) || prefillWeight <= 0) {
      throw new Error("budget prefill weight must be a positive number");
    }
    if (!Number.isFinite(samplingWeight) || samplingWeight <= 0) {
      throw new Error("budget sampling weight must be a positive number");
    }
    this.#limitTokens = config.limitTokens;
    this.#prefillWeight = prefillWeight;
    this.#samplingWeight = samplingWeight;
  }

  /** True when the cumulative weighted spend reached the limit. */
  exhausted(): boolean {
    return this.#exhausted;
  }

  /** Weighted tokens left before the limit. Floors at zero. */
  remainingTokens(): number {
    return Math.max(0, Math.floor(this.#limitTokens - this.#weightedTokensUsed));
  }

  /** The cumulative weighted spend recorded so far. */
  weightedTokensUsed(): number {
    return this.#weightedTokensUsed;
  }

  /**
   * Accumulate weighted spend. Returns true once the limit is crossed
   * (including on later calls); a limit of 0 means unlimited.
   */
  recordUsage(usage: SessionUsage): boolean {
    if (this.#limitTokens <= 0) {
      return false;
    }
    this.#weightedTokensUsed += this.#weightedSpend(usage);
    if (this.#weightedTokensUsed >= this.#limitTokens) {
      this.#exhausted = true;
    }
    return this.#exhausted;
  }

  #weightedSpend(usage: SessionUsage): number {
    if (usage.inputTokens !== undefined || usage.outputTokens !== undefined) {
      const output = Math.max(0, usage.outputTokens ?? 0) * this.#samplingWeight;
      const nonCachedInput =
        Math.max(0, (usage.inputTokens ?? 0) - Math.max(0, usage.cacheReadTokens ?? 0));
      return output + nonCachedInput * this.#prefillWeight;
    }
    // No provider token counts: the text length is the only measurable side,
    // so the `textLength / 4` estimate is weighted at the sampling rate.
    return ((usage.textLength ?? 0) / FALLBACK_CHARS_PER_TOKEN) * this.#samplingWeight;
  }
}