import { BestAgentError } from "@best-agent/foundation";
import type {
  ArtifactObligation,
  ProviderConfig,
  RequestBudgetSnapshot,
  RunPhaseName,
  SessionUsageOperation,
} from "@best-agent/protocol";
import type { AgentProfile } from "../../agent/profile-registry.js";
import type { ProviderToolOption } from "../../provider/model-provider.js";
import type {
  AdvancedPhaseLoopState,
  PendingCompactionFollowUpMarker,
  ProviderCorrectionTrace,
} from "../checkpoint-codec.js";
import type { RequiredToolConstraint } from "./tool-constraint-types.js";

export type ProviderTurnBudgetSnapshot = {
  enabled: boolean;
  max: number;
  used: number;
};

export type ProviderTurnInvokeContext = {
  activeSkillIds: string[];
  advancedPhase?: AdvancedPhaseLoopState;
  artifactObligations?: ArtifactObligation[];
  availableTools?: ProviderToolOption[];
  budget: ProviderTurnBudget;
  compactionFollowUpMarker?: PendingCompactionFollowUpMarker;
  correctionTrace?: ProviderCorrectionTrace;
  downgradeEligible?: boolean;
  escalationLevel?: number;
  fixedRouteConfigOverride?: boolean;
  objective?: string;
  phaseName?: RunPhaseName;
  profile: AgentProfile;
  profileName: string;
  requiredToolConstraint?: RequiredToolConstraint;
  routeConfigOverride?: ProviderConfig;
  requestBudgetSnapshot?: RequestBudgetSnapshot;
  supersedePendingProviderRequest?: boolean;
  toolStepCount?: number;
};

export class ProviderTurnBudget {
  #used: number;
  readonly enabled: boolean;
  readonly max: number;

  constructor(input: ProviderTurnBudgetSnapshot) {
    this.enabled = input.enabled;
    this.max = input.max;
    this.#used = input.used;
  }

  consume(reason: SessionUsageOperation): number {
    if (!this.enabled) {
      return this.#used;
    }
    const next = this.#used + 1;
    if (next > this.max) {
      throw new BestAgentError(
        "PROVIDER_TURN_BUDGET_EXCEEDED",
        `Provider exceeded the maximum ordinary turn decision budget of ${this.max}. reason=${reason}`,
      );
    }
    this.#used = next;
    return this.#used;
  }

  get used(): number {
    return this.#used;
  }

  snapshot(): ProviderTurnBudgetSnapshot {
    return {
      enabled: this.enabled,
      max: this.max,
      used: this.#used,
    };
  }
}
