import type {
  PlayRunnerBudgetCharge,
  PlayRunnerBudgetChargeResult,
} from '../protocol';

export interface BudgetStateBackend {
  charge(input: {
    rootRunId: string;
    charges: readonly PlayRunnerBudgetCharge[];
  }): Promise<PlayRunnerBudgetChargeResult>;
}

export class BudgetStateLimitError extends Error {
  constructor(
    readonly key: string,
    readonly observed: number,
    readonly limit: number,
  ) {
    super(`Play execution ${key} budget exceeded (${observed}/${limit}).`);
    this.name = 'BudgetStateLimitError';
  }
}

export class InMemoryBudgetStateBackend implements BudgetStateBackend {
  private readonly counters = new Map<string, Record<string, number>>();

  async charge(input: {
    rootRunId: string;
    charges: readonly PlayRunnerBudgetCharge[];
  }): Promise<PlayRunnerBudgetChargeResult> {
    const current = this.counters.get(input.rootRunId) ?? {};
    const next = { ...current };
    for (const charge of input.charges) {
      const observed = (next[charge.key] ?? 0) + charge.amount;
      if (observed > charge.limit) {
        throw new BudgetStateLimitError(charge.key, observed, charge.limit);
      }
      next[charge.key] = observed;
    }
    this.counters.set(input.rootRunId, next);
    return { counters: { ...next } };
  }
}