/**
 * KHAELOR
 * File: src/context/budget.ts
 * Description: Token accounting from real ModelUsage only — pressure thresholds and config-driven context windows (ARCHITECTURE.md §6.2).
 *
 * Author: Simon-Pierre Boucher
 * Contact: contact@spboucher.ai
 */

import type { ModelUsage } from "../anthropic/index.js";

/**
 * Fallback context window when no config override matches the configured
 * model. Not a hard-coded model list — every value here is overridable via
 * `ContextWindowConfig` (CLAUDE.md §6: model identifiers evolve).
 */
export const DEFAULT_CONTEXT_WINDOW = 200_000;

/** Safety margin kept free below the usable window (ARCHITECTURE.md §6.2). */
export const DEFAULT_COMPACTION_BUFFER_TOKENS = 20_000;

/** Fraction of the usable window at which cheap pruning starts (prune-before-compress order). */
export const DEFAULT_PRUNE_THRESHOLD_RATIO = 0.85;

/** Config-driven context-window table — never a hard-coded model-id list. */
export interface ContextWindowConfig {
  /** Fallback window for models absent from `byModel`. Default 200_000. */
  defaultWindow?: number;
  /** Exact model id → context window override (from user/project config). */
  byModel?: Record<string, number>;
}

/** Resolve the context window for a configured model id. */
export function resolveContextWindow(model: string, config?: ContextWindowConfig): number {
  const override = config?.byModel?.[model];
  if (override !== undefined && Number.isFinite(override) && override > 0) return override;
  const fallback = config?.defaultWindow;
  if (fallback !== undefined && Number.isFinite(fallback) && fallback > 0) return fallback;
  return DEFAULT_CONTEXT_WINDOW;
}

/**
 * Cheap deterministic token ESTIMATE (~4 chars/token) — used only for
 * `/context` section stats and prune/cut selection, always labeled as an
 * estimate. Real accounting comes exclusively from API usage (Rule #4).
 */
export function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4);
}

/** Pressure levels: prune first (cheap, deterministic), compress only if still over budget. */
export type BudgetPressureLevel = "ok" | "prune" | "compact";

export interface ContextBudgetOptions {
  /** The configured main model id (window resolution key). */
  model: string;
  /** Output tokens reserved for the model's reply (config `maxOutputTokens`). */
  reservedOutputTokens: number;
  /** Config-driven window table + fallback. */
  contextWindow?: ContextWindowConfig;
  /** Safety margin below the window. Default 20_000. */
  compactionBufferTokens?: number;
  /** Fraction of the usable window at which pruning triggers. Default 0.85. */
  pruneThresholdRatio?: number;
}

/**
 * Token-pressure accounting fed exclusively by real API usage
 * (`ModelResponseCompleted.usage`) — never event counts, never estimates
 * (ARCHITECTURE.md §6.2, Absolute Rule #4).
 */
export class ContextBudget {
  readonly modelWindow: number;
  readonly reservedOutputTokens: number;
  readonly compactionBufferTokens: number;
  readonly pruneThresholdRatio: number;

  #lastTotalTokens = 0;
  #observedTurns = 0;

  constructor(options: ContextBudgetOptions) {
    this.modelWindow = resolveContextWindow(options.model, options.contextWindow);
    this.reservedOutputTokens = options.reservedOutputTokens;
    this.compactionBufferTokens =
      options.compactionBufferTokens ?? DEFAULT_COMPACTION_BUFFER_TOKENS;
    this.pruneThresholdRatio = options.pruneThresholdRatio ?? DEFAULT_PRUNE_THRESHOLD_RATIO;
  }

  /** usableWindow = modelWindow − reservedOutput − compactionBuffer (ARCHITECTURE.md §6.2). */
  get usableWindow(): number {
    return Math.max(1, this.modelWindow - this.reservedOutputTokens - this.compactionBufferTokens);
  }

  /**
   * Record real usage from the last completed response. The context footprint
   * of the next request tracks the full prompt the API just billed:
   * uncached input + cache reads + cache writes, plus the tokens the model
   * just produced (which become history).
   */
  onUsage(usage: ModelUsage): void {
    this.#lastTotalTokens =
      usage.inputTokens + usage.cacheReadTokens + usage.cacheWriteTokens + usage.outputTokens;
    this.#observedTurns += 1;
  }

  /** Last observed real token total (0 until the first response completes). */
  get lastTotalTokens(): number {
    return this.#lastTotalTokens;
  }

  /** Whether any real usage has been observed yet. */
  get hasObservedUsage(): boolean {
    return this.#observedTurns > 0;
  }

  /** Pressure as a fraction of the usable window (0 = empty, ≥1 = over budget). */
  pressure(): number {
    return this.#lastTotalTokens / this.usableWindow;
  }

  /** Current trigger level: ok → prune (cheap) → compact (LLM summarization). */
  level(): BudgetPressureLevel {
    const p = this.pressure();
    if (p >= 1) return "compact";
    if (p >= this.pruneThresholdRatio) return "prune";
    return "ok";
  }

  /** Proactive compaction trigger (ARCHITECTURE.md §6.2). */
  shouldCompact(): boolean {
    return this.level() === "compact";
  }
}