// src/utils/__tests__/pruneCalibration.test.ts
import {
  createPruneCalibration,
  updatePruneCalibration,
  applyCalibration,
} from '../pruneCalibration';
import {
  PRUNING_INITIAL_CALIBRATION,
  PRUNING_EMA_ALPHA,
} from '@/common/constants';

describe('pruneCalibration', () => {
  describe('createPruneCalibration', () => {
    it('creates initial state with default ratio', () => {
      const state = createPruneCalibration();
      expect(state.ratio).toBe(PRUNING_INITIAL_CALIBRATION);
      expect(state.iterations).toBe(0);
    });

    it('accepts custom initial ratio', () => {
      const state = createPruneCalibration(0.85);
      expect(state.ratio).toBe(0.85);
      expect(state.iterations).toBe(0);
    });
  });

  describe('updatePruneCalibration', () => {
    it('adjusts ratio when actual > estimated (over-counting)', () => {
      const state = createPruneCalibration();
      // Actual: 1000 tokens, estimated: 1500 tokens (our counter over-estimates)
      // observedRatio = 1500/1000 = 1.5
      // newRatio = 0.3 * 1.5 + 0.7 * 1.0 = 0.45 + 0.7 = 1.15
      const updated = updatePruneCalibration(state, 1000, 1500);
      expect(updated.ratio).toBeCloseTo(1.15, 2);
      expect(updated.iterations).toBe(1);
    });

    it('adjusts ratio when actual < estimated (under-counting)', () => {
      const state = createPruneCalibration();
      // Actual: 2000 tokens, estimated: 1000 tokens (our counter under-estimates)
      // observedRatio = 1000/2000 = 0.5
      // newRatio = 0.3 * 0.5 + 0.7 * 1.0 = 0.15 + 0.7 = 0.85
      const updated = updatePruneCalibration(state, 2000, 1000);
      expect(updated.ratio).toBeCloseTo(0.85, 2);
      expect(updated.iterations).toBe(1);
    });

    it('converges with consistent readings', () => {
      let state = createPruneCalibration();

      // Simulate 10 iterations where actual is consistently 1.5x estimated
      for (let i = 0; i < 10; i++) {
        state = updatePruneCalibration(state, 1500, 1000);
      }

      // Should converge toward ~0.667 (estimated/actual = 1000/1500)
      expect(state.ratio).toBeCloseTo(0.667, 1);
      expect(state.iterations).toBe(10);
    });

    it('clamps extreme ratios to prevent wild adjustments', () => {
      const state = createPruneCalibration();

      // Extreme case: estimated 10x actual (should be clamped to 2.0)
      const updated = updatePruneCalibration(state, 100, 10000);
      // Clamped observedRatio = 2.0
      // newRatio = 0.3 * 2.0 + 0.7 * 1.0 = 0.6 + 0.7 = 1.3
      expect(updated.ratio).toBeCloseTo(1.3, 2);
    });

    it('does not update with invalid inputs', () => {
      const state = createPruneCalibration();

      expect(updatePruneCalibration(state, 0, 1000)).toBe(state);
      expect(updatePruneCalibration(state, 1000, 0)).toBe(state);
      expect(updatePruneCalibration(state, -1, 1000)).toBe(state);
    });

    it('does not mutate input state', () => {
      const state = createPruneCalibration();
      const original = { ...state };

      updatePruneCalibration(state, 1000, 1500);
      expect(state).toEqual(original);
    });

    it('accepts custom alpha', () => {
      const state = createPruneCalibration();
      // With alpha=1.0, fully adapts to new reading
      const updated = updatePruneCalibration(state, 1000, 1500, 1.0);
      // observedRatio = 1.5, clamped to 1.5
      // newRatio = 1.0 * 1.5 + 0.0 * 1.0 = 1.5
      expect(updated.ratio).toBeCloseTo(1.5, 2);
    });
  });

  describe('applyCalibration', () => {
    it('returns raw budget when no iterations have occurred', () => {
      const state = createPruneCalibration();
      expect(applyCalibration(10000, state)).toBe(10000);
    });

    it('adjusts budget after calibration', () => {
      let state = createPruneCalibration();
      state = updatePruneCalibration(state, 1000, 1500);
      // ratio ≈ 1.15, so budget is increased (our counter over-estimates)
      const adjusted = applyCalibration(10000, state);
      expect(adjusted).toBeCloseTo(11500, -2);
    });

    it('decreases budget when under-counting', () => {
      let state = createPruneCalibration();
      state = updatePruneCalibration(state, 2000, 1000);
      // ratio ≈ 0.85, so budget is decreased (our counter under-estimates)
      const adjusted = applyCalibration(10000, state);
      expect(adjusted).toBeLessThan(10000);
    });

    it('returns floor of the adjusted value', () => {
      let state = createPruneCalibration();
      state = updatePruneCalibration(state, 1000, 1500);
      const adjusted = applyCalibration(10001, state);
      expect(Number.isInteger(adjusted)).toBe(true);
    });
  });

  describe('multi-iteration convergence', () => {
    it('smoothly transitions when accuracy changes', () => {
      let state = createPruneCalibration();

      // First 5 iterations: estimated is 1.5x actual
      for (let i = 0; i < 5; i++) {
        state = updatePruneCalibration(state, 1000, 1500);
      }
      const ratio5 = state.ratio;

      // Next 5 iterations: estimated matches actual
      for (let i = 0; i < 5; i++) {
        state = updatePruneCalibration(state, 1000, 1000);
      }
      const ratio10 = state.ratio;

      // Ratio should move toward 1.0 but still carry some history
      expect(ratio10).toBeLessThan(ratio5);
      expect(ratio10).toBeGreaterThan(0.9);
    });
  });
});