// src/utils/pruneCalibration.ts
import type { PruneCalibrationState } from '@/types/graph';
import {
  PRUNING_EMA_ALPHA,
  PRUNING_INITIAL_CALIBRATION,
} from '@/common/constants';

/**
 * Creates an initial pruning calibration state.
 *
 * @param initialRatio - Starting calibration ratio (default: 1.0)
 * @returns Fresh calibration state
 */
export function createPruneCalibration(
  initialRatio?: number
): PruneCalibrationState {
  return {
    ratio: initialRatio ?? PRUNING_INITIAL_CALIBRATION,
    iterations: 0,
  };
}

/**
 * Updates the pruning calibration using Exponential Moving Average (EMA).
 *
 * Problem: Without calibration, the pruner's token estimates can diverge from
 * reality across iterations, causing either:
 * - Over-pruning (context cliff): Too many messages removed at once, losing critical tool results
 * - Under-pruning: Not enough messages removed, hitting hard token limits
 *
 * Solution: Track the ratio between actual token usage (from API response) and
 * estimated token usage (from our token counter). Apply EMA smoothing so the
 * calibration adjusts gradually, preventing oscillation.
 *
 * The calibration ratio is applied to maxTokens in the pruner:
 *   effectiveMaxTokens = maxTokens * calibrationRatio
 *
 * If actual > estimated → ratio decreases → prune more aggressively
 * If actual < estimated → ratio increases → prune less aggressively
 *
 * @param state - Current calibration state
 * @param actualTokens - Actual token count from API response (UsageMetadata)
 * @param estimatedTokens - Estimated token count from token counter
 * @param alpha - EMA smoothing factor (default: PRUNING_EMA_ALPHA)
 * @returns Updated calibration state (new object, does not mutate input)
 */
export function updatePruneCalibration(
  state: PruneCalibrationState,
  actualTokens: number,
  estimatedTokens: number,
  alpha: number = PRUNING_EMA_ALPHA
): PruneCalibrationState {
  // Guard against division by zero or invalid inputs
  if (estimatedTokens <= 0 || actualTokens <= 0) {
    return state;
  }

  // Raw ratio: how much our estimate differs from reality
  const observedRatio = estimatedTokens / actualTokens;

  // Clamp to prevent extreme adjustments from outlier readings
  // Range [0.5, 2.0] means we never more than double or halve the budget
  const clampedRatio = Math.max(0.5, Math.min(2.0, observedRatio));

  // Apply EMA: new_ratio = α * observed + (1 - α) * previous
  const newRatio = alpha * clampedRatio + (1 - alpha) * state.ratio;

  return {
    ratio: newRatio,
    iterations: state.iterations + 1,
  };
}

/**
 * Applies the calibration ratio to a max token budget.
 * The ratio adjusts the effective budget so pruning is more or less aggressive
 * based on observed vs. estimated token divergence.
 *
 * @param maxTokens - Raw max token budget
 * @param state - Current calibration state
 * @returns Adjusted max token budget
 */
export function applyCalibration(
  maxTokens: number,
  state: PruneCalibrationState
): number {
  if (state.iterations === 0) {
    // No calibration data yet — use raw budget
    return maxTokens;
  }
  return Math.floor(maxTokens * state.ratio);
}