rotation.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";

import {
  HealthScoreTracker,
  TokenBucketTracker,
  addJitter,
  randomDelay,
  sortByLruWithHealth,
  selectHybridAccount,
  type AccountWithMetrics,
} from "./rotation";

describe("HealthScoreTracker", () => {
  beforeEach(() => {
    vi.useRealTimers();
  });

  describe("initial state", () => {
    it("returns initial score for unknown account", () => {
      const tracker = new HealthScoreTracker();
      expect(tracker.getScore(0)).toBe(70);
    });

    it("uses custom initial score from config", () => {
      const tracker = new HealthScoreTracker({ initial: 50 });
      expect(tracker.getScore(0)).toBe(50);
    });

    it("isUsable returns true for new accounts", () => {
      const tracker = new HealthScoreTracker();
      expect(tracker.isUsable(0)).toBe(true);
    });

    it("getConsecutiveFailures returns 0 for unknown account", () => {
      const tracker = new HealthScoreTracker();
      expect(tracker.getConsecutiveFailures(0)).toBe(0);
    });
  });

  describe("recordSuccess", () => {
    it("increases score by success reward", () => {
      const tracker = new HealthScoreTracker({ initial: 70, successReward: 5 });
      tracker.recordSuccess(0);
      expect(tracker.getScore(0)).toBe(75);
    });

    it("caps score at maxScore", () => {
      const tracker = new HealthScoreTracker({ initial: 98, successReward: 5, maxScore: 100 });
      tracker.recordSuccess(0);
      expect(tracker.getScore(0)).toBe(100);
    });

    it("resets consecutive failures", () => {
      const tracker = new HealthScoreTracker();
      tracker.recordRateLimit(0);
      tracker.recordRateLimit(0);
      expect(tracker.getConsecutiveFailures(0)).toBe(2);
      
      tracker.recordSuccess(0);
      expect(tracker.getConsecutiveFailures(0)).toBe(0);
    });
  });

  describe("recordRateLimit", () => {
    it("decreases score by rate limit penalty", () => {
      const tracker = new HealthScoreTracker({ initial: 70, rateLimitPenalty: -10 });
      tracker.recordRateLimit(0);
      expect(tracker.getScore(0)).toBe(60);
    });

    it("does not go below 0", () => {
      const tracker = new HealthScoreTracker({ initial: 5, rateLimitPenalty: -10 });
      tracker.recordRateLimit(0);
      expect(tracker.getScore(0)).toBe(0);
    });

    it("increments consecutive failures", () => {
      const tracker = new HealthScoreTracker();
      tracker.recordRateLimit(0);
      expect(tracker.getConsecutiveFailures(0)).toBe(1);
      
      tracker.recordRateLimit(0);
      expect(tracker.getConsecutiveFailures(0)).toBe(2);
    });
  });

  describe("recordFailure", () => {
    it("decreases score by failure penalty", () => {
      const tracker = new HealthScoreTracker({ initial: 70, failurePenalty: -20 });
      tracker.recordFailure(0);
      expect(tracker.getScore(0)).toBe(50);
    });

    it("does not go below 0", () => {
      const tracker = new HealthScoreTracker({ initial: 10, failurePenalty: -20 });
      tracker.recordFailure(0);
      expect(tracker.getScore(0)).toBe(0);
    });

    it("increments consecutive failures", () => {
      const tracker = new HealthScoreTracker();
      tracker.recordFailure(0);
      expect(tracker.getConsecutiveFailures(0)).toBe(1);
    });
  });

  describe("isUsable", () => {
    it("returns true when score >= minUsable", () => {
      const tracker = new HealthScoreTracker({ initial: 50, minUsable: 50 });
      expect(tracker.isUsable(0)).toBe(true);
    });

    it("returns false when score < minUsable", () => {
      const tracker = new HealthScoreTracker({ initial: 49, minUsable: 50 });
      expect(tracker.isUsable(0)).toBe(false);
    });

    it("becomes unusable after multiple failures", () => {
      const tracker = new HealthScoreTracker({ initial: 70, failurePenalty: -20, minUsable: 50 });
      tracker.recordFailure(0);
      expect(tracker.isUsable(0)).toBe(true);
      
      tracker.recordFailure(0);
      expect(tracker.isUsable(0)).toBe(false);
    });
  });

  describe("time-based recovery", () => {
    it("recovers points over time", () => {
      let mockTime = 0;
      vi.spyOn(Date, 'now').mockImplementation(() => mockTime);

      const tracker = new HealthScoreTracker({ 
        initial: 70, 
        failurePenalty: -20, 
        recoveryRatePerHour: 10 
      });
      
      tracker.recordFailure(0);
      expect(tracker.getScore(0)).toBe(50);

      mockTime = 2 * 60 * 60 * 1000;
      expect(tracker.getScore(0)).toBe(70);

      vi.restoreAllMocks();
    });

    it("caps recovery at maxScore", () => {
      let mockTime = 0;
      vi.spyOn(Date, 'now').mockImplementation(() => mockTime);

      const tracker = new HealthScoreTracker({ 
        initial: 90, 
        successReward: 5,
        recoveryRatePerHour: 20,
        maxScore: 100 
      });
      
      tracker.recordSuccess(0);
      expect(tracker.getScore(0)).toBe(95);
      
      mockTime = 60 * 60 * 1000;
      expect(tracker.getScore(0)).toBe(100);

      vi.restoreAllMocks();
    });

    it("floors recovered points (no partial points)", () => {
      let mockTime = 0;
      vi.spyOn(Date, 'now').mockImplementation(() => mockTime);

      const tracker = new HealthScoreTracker({ 
        initial: 70, 
        failurePenalty: -10, 
        recoveryRatePerHour: 2 
      });
      
      tracker.recordFailure(0);
      expect(tracker.getScore(0)).toBe(60);

      mockTime = 20 * 60 * 1000;
      expect(tracker.getScore(0)).toBe(60);

      mockTime = 30 * 60 * 1000;
      expect(tracker.getScore(0)).toBe(61);

      vi.restoreAllMocks();
    });
  });

  describe("reset", () => {
    it("clears health state for account", () => {
      const tracker = new HealthScoreTracker({ initial: 70 });
      tracker.recordSuccess(0);
      tracker.reset(0);
      
      expect(tracker.getScore(0)).toBe(70);
      expect(tracker.getConsecutiveFailures(0)).toBe(0);
    });
  });

  describe("getSnapshot", () => {
    it("returns current state of all tracked accounts", () => {
      const tracker = new HealthScoreTracker({ initial: 70, failurePenalty: -10 });
      tracker.recordSuccess(0);
      tracker.recordFailure(1);
      tracker.recordFailure(1);
      
      const snapshot = tracker.getSnapshot();
      expect(snapshot.get(0)?.score).toBe(71);
      expect(snapshot.get(0)?.consecutiveFailures).toBe(0);
      expect(snapshot.get(1)?.score).toBe(50);
      expect(snapshot.get(1)?.consecutiveFailures).toBe(2);
    });
  });
});

describe("TokenBucketTracker", () => {
  beforeEach(() => {
    vi.useRealTimers();
  });

  describe("initial state", () => {
    it("returns initial tokens for unknown account", () => {
      const tracker = new TokenBucketTracker();
      expect(tracker.getTokens(0)).toBe(50);
    });

    it("uses custom initial tokens from config", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 30 });
      expect(tracker.getTokens(0)).toBe(30);
    });

    it("hasTokens returns true for new accounts", () => {
      const tracker = new TokenBucketTracker();
      expect(tracker.hasTokens(0)).toBe(true);
    });

    it("getMaxTokens returns configured max tokens", () => {
      const tracker = new TokenBucketTracker({ maxTokens: 100 });
      expect(tracker.getMaxTokens()).toBe(100);
    });

    it("getMaxTokens returns default when not configured", () => {
      const tracker = new TokenBucketTracker();
      expect(tracker.getMaxTokens()).toBe(50);
    });
  });

  describe("consume", () => {
    it("reduces token balance", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 50 });
      expect(tracker.consume(0, 1)).toBe(true);
      // Use toBeCloseTo to handle floating point from micro-regeneration between calls
      expect(tracker.getTokens(0)).toBeCloseTo(49, 2);
    });

    it("returns false when insufficient tokens", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 5 });
      expect(tracker.consume(0, 10)).toBe(false);
      expect(tracker.getTokens(0)).toBe(5);
    });

    it("allows consuming exact remaining tokens", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 10 });
      expect(tracker.consume(0, 10)).toBe(true);
      // Use toBeCloseTo to handle floating point from micro-regeneration between calls
      expect(tracker.getTokens(0)).toBeCloseTo(0, 2);
    });

    it("handles multiple consumes", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 50 });
      tracker.consume(0, 10);
      tracker.consume(0, 10);
      tracker.consume(0, 10);
      expect(tracker.getTokens(0)).toBe(20);
    });
  });

  describe("hasTokens", () => {
    it("returns true when enough tokens", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 50 });
      expect(tracker.hasTokens(0, 50)).toBe(true);
    });

    it("returns false when insufficient tokens", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 10 });
      expect(tracker.hasTokens(0, 11)).toBe(false);
    });

    it("defaults to cost of 1", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 1 });
      expect(tracker.hasTokens(0)).toBe(true);
      
      tracker.consume(0, 1);
      expect(tracker.hasTokens(0)).toBe(false);
    });
  });

  describe("refund", () => {
    it("adds tokens back", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 50 });
      tracker.consume(0, 10);
      expect(tracker.getTokens(0)).toBeCloseTo(40, 2);
      
      tracker.refund(0, 5);
      expect(tracker.getTokens(0)).toBeCloseTo(45, 2);
    });

    it("caps at maxTokens", () => {
      const tracker = new TokenBucketTracker({ initialTokens: 50, maxTokens: 50 });
      tracker.refund(0, 10);
      expect(tracker.getTokens(0)).toBe(50);
    });
  });

  describe("token regeneration", () => {
    it("regenerates tokens over time", () => {
      let mockTime = 0;
      vi.spyOn(Date, 'now').mockImplementation(() => mockTime);

      const tracker = new TokenBucketTracker({ 
        initialTokens: 50, 
        maxTokens: 50,
        regenerationRatePerMinute: 6 
      });
      
      tracker.consume(0, 30);
      expect(tracker.getTokens(0)).toBe(20);

      mockTime = 5 * 60 * 1000;
      expect(tracker.getTokens(0)).toBe(50);

      vi.restoreAllMocks();
    });

    it("caps regeneration at maxTokens", () => {
      let mockTime = 0;
      vi.spyOn(Date, 'now').mockImplementation(() => mockTime);

      const tracker = new TokenBucketTracker({ 
        initialTokens: 40, 
        maxTokens: 50,
        regenerationRatePerMinute: 6 
      });
      
      tracker.consume(0, 1);
      
      mockTime = 10 * 60 * 1000;
      expect(tracker.getTokens(0)).toBe(50);

      vi.restoreAllMocks();
    });
  });
});

describe("addJitter", () => {
  it("returns value within jitter range", () => {
    const base = 1000;
    const jitterFactor = 0.3;
    
    for (let i = 0; i < 100; i++) {
      const result = addJitter(base, jitterFactor);
      expect(result).toBeGreaterThanOrEqual(base * (1 - jitterFactor));
      expect(result).toBeLessThanOrEqual(base * (1 + jitterFactor));
    }
  });

  it("uses default jitter factor of 0.3", () => {
    const base = 1000;
    
    for (let i = 0; i < 100; i++) {
      const result = addJitter(base);
      expect(result).toBeGreaterThanOrEqual(700);
      expect(result).toBeLessThanOrEqual(1300);
    }
  });

  it("never returns negative values", () => {
    for (let i = 0; i < 100; i++) {
      const result = addJitter(10, 0.9);
      expect(result).toBeGreaterThanOrEqual(0);
    }
  });

  it("returns rounded values", () => {
    for (let i = 0; i < 100; i++) {
      const result = addJitter(1000);
      expect(Number.isInteger(result)).toBe(true);
    }
  });
});

describe("randomDelay", () => {
  it("returns value within min-max range", () => {
    for (let i = 0; i < 100; i++) {
      const result = randomDelay(100, 500);
      expect(result).toBeGreaterThanOrEqual(100);
      expect(result).toBeLessThanOrEqual(500);
    }
  });

  it("returns rounded values", () => {
    for (let i = 0; i < 100; i++) {
      const result = randomDelay(100, 500);
      expect(Number.isInteger(result)).toBe(true);
    }
  });

  it("handles min === max", () => {
    const result = randomDelay(100, 100);
    expect(result).toBe(100);
  });
});

describe("sortByLruWithHealth", () => {
  it("filters out rate-limited accounts", () => {
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 70, isRateLimited: true, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = sortByLruWithHealth(accounts);
    expect(result).toHaveLength(1);
    expect(result[0]?.index).toBe(1);
  });

  it("filters out cooling down accounts", () => {
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: true },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = sortByLruWithHealth(accounts);
    expect(result).toHaveLength(1);
    expect(result[0]?.index).toBe(1);
  });

  it("filters out unhealthy accounts", () => {
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 40, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = sortByLruWithHealth(accounts, 50);
    expect(result).toHaveLength(1);
    expect(result[0]?.index).toBe(1);
  });

  it("sorts by lastUsed ascending (oldest first)", () => {
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 1000, healthScore: 70, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 500, healthScore: 70, isRateLimited: false, isCoolingDown: false },
      { index: 2, lastUsed: 2000, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = sortByLruWithHealth(accounts);
    expect(result.map(a => a.index)).toEqual([1, 0, 2]);
  });

  it("uses health score as tiebreaker", () => {
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 1000, healthScore: 60, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 1000, healthScore: 80, isRateLimited: false, isCoolingDown: false },
      { index: 2, lastUsed: 1000, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = sortByLruWithHealth(accounts);
    expect(result.map(a => a.index)).toEqual([1, 2, 0]);
  });

  it("returns empty array when all accounts filtered out", () => {
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 30, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: true, isCoolingDown: false },
    ];

    const result = sortByLruWithHealth(accounts, 50);
    expect(result).toHaveLength(0);
  });
});

describe("selectHybridAccount", () => {
  it("returns null when no accounts available", () => {
    const tokenTracker = new TokenBucketTracker();
    const result = selectHybridAccount([], tokenTracker);
    expect(result).toBeNull();
  });

  it("returns null when all accounts filtered out by health", () => {
    const tokenTracker = new TokenBucketTracker();
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 30, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker, 50);
    expect(result).toBeNull();
  });

  it("returns the best candidate by score", () => {
    const tokenTracker = new TokenBucketTracker();
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 1000, healthScore: 70, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 500, healthScore: 70, isRateLimited: false, isCoolingDown: false },
      { index: 2, lastUsed: 2000, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker);
    expect([0, 1, 2]).toContain(result);
  });

  it("filters out rate-limited accounts", () => {
    const tokenTracker = new TokenBucketTracker();
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 70, isRateLimited: true, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker);
    expect(result).toBe(1);
  });

  it("filters out accounts without tokens", () => {
    const tokenTracker = new TokenBucketTracker({ initialTokens: 1 });
    tokenTracker.consume(0, 1);
    
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker);
    expect(result).toBe(1);
  });

  it("filters out unhealthy accounts", () => {
    const tokenTracker = new TokenBucketTracker();
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 40, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker, 50);
    expect(result).toBe(1);
  });

  it("returns null when all accounts have no tokens", () => {
    const tokenTracker = new TokenBucketTracker({ initialTokens: 0 });
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 70, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker);
    expect(result).toBeNull();
  });

  it("selects only available candidate when one account is filtered", () => {
    const tokenTracker = new TokenBucketTracker({ initialTokens: 50 });
    
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 0, healthScore: 40, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 0, healthScore: 100, isRateLimited: false, isCoolingDown: false },
    ];

    const result = selectHybridAccount(accounts, tokenTracker, 50);
    expect(result).toBe(1);
  });

  it("returns a valid account index", () => {
    const tokenTracker = new TokenBucketTracker();
    const accounts: AccountWithMetrics[] = [
      { index: 0, lastUsed: 1000, healthScore: 70, isRateLimited: false, isCoolingDown: false },
      { index: 1, lastUsed: 500, healthScore: 80, isRateLimited: false, isCoolingDown: false },
      { index: 2, lastUsed: 2000, healthScore: 60, isRateLimited: false, isCoolingDown: false },
    ];

    for (let i = 0; i < 10; i++) {
      const result = selectHybridAccount(accounts, tokenTracker);
      expect([0, 1, 2]).toContain(result);
    }
  });
});