import { describe, expect } from "vitest";
import { redisTest } from "@internal/testcontainers";
import { createRedisClient, type Redis } from "@internal/redis";
import { DRRScheduler } from "../schedulers/drr.js";
import { DefaultFairQueueKeyProducer } from "../keyProducer.js";
import type { FairQueueKeyProducer, SchedulerContext } from "../types.js";

describe("DRRScheduler", () => {
  let keys: FairQueueKeyProducer;

  describe("deficit management", () => {
    redisTest("should initialize deficit to 0 for new tenants", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const deficit = await scheduler.getDeficit("new-tenant");
      expect(deficit).toBe(0);

      await scheduler.close();
    });

    redisTest("should add quantum atomically with capping", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      // Setup: put queues in the master shard
      const masterKey = keys.masterQueueKey(0);
      const now = Date.now();

      await redis.zadd(masterKey, now, "tenant:t1:queue:q1");

      // Create context mock
      const context: SchedulerContext = {
        getCurrentConcurrency: async () => 0,
        getConcurrencyLimit: async () => 100,
        isAtCapacity: async () => false,
        getQueueDescriptor: (queueId) => ({
          id: queueId,
          tenantId: keys.extractTenantId(queueId),
          metadata: {},
        }),
      };

      // Run multiple iterations to accumulate deficit
      for (let i = 0; i < 15; i++) {
        await scheduler.selectQueues(masterKey, "consumer-1", context);
      }

      // Deficit should be capped at maxDeficit (50)
      const deficit = await scheduler.getDeficit("t1");
      expect(deficit).toBeLessThanOrEqual(50);

      await scheduler.close();
      await redis.quit();
    });

    redisTest("should decrement deficit when processing", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      // Manually set some deficit
      const deficitKey = `test:drr:deficit`;
      await redis.hset(deficitKey, "t1", "10");

      // Record processing
      await scheduler.recordProcessed("t1", "queue:q1");

      const deficit = await scheduler.getDeficit("t1");
      expect(deficit).toBe(9);

      await scheduler.close();
      await redis.quit();
    });

    redisTest("should not go below 0 on decrement", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const deficitKey = `test:drr:deficit`;
      await redis.hset(deficitKey, "t1", "0.5");

      await scheduler.recordProcessed("t1", "queue:q1");

      const deficit = await scheduler.getDeficit("t1");
      expect(deficit).toBe(0);

      await scheduler.close();
      await redis.quit();
    });

    redisTest(
      "should decrement deficit by count when using recordProcessedBatch",
      { timeout: 10000 },
      async ({ redisOptions }) => {
        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
        const redis = createRedisClient(redisOptions);

        const scheduler = new DRRScheduler({
          redis: redisOptions,
          keys,
          quantum: 5,
          maxDeficit: 50,
        });

        // Manually set deficit
        const deficitKey = `test:drr:deficit`;
        await redis.hset(deficitKey, "t1", "15");

        // Record batch processing of 7 messages
        await scheduler.recordProcessedBatch("t1", "queue:q1", 7);

        const deficit = await scheduler.getDeficit("t1");
        expect(deficit).toBe(8);

        await scheduler.close();
        await redis.quit();
      }
    );

    redisTest(
      "should not go below 0 when recordProcessedBatch decrements more than available",
      { timeout: 10000 },
      async ({ redisOptions }) => {
        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
        const redis = createRedisClient(redisOptions);

        const scheduler = new DRRScheduler({
          redis: redisOptions,
          keys,
          quantum: 5,
          maxDeficit: 50,
        });

        // Manually set deficit to 3
        const deficitKey = `test:drr:deficit`;
        await redis.hset(deficitKey, "t1", "3");

        // Record batch processing of 10 messages (more than deficit)
        await scheduler.recordProcessedBatch("t1", "queue:q1", 10);

        const deficit = await scheduler.getDeficit("t1");
        expect(deficit).toBe(0);

        await scheduler.close();
        await redis.quit();
      }
    );

    redisTest("should reset deficit for tenant", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const deficitKey = `test:drr:deficit`;
      await redis.hset(deficitKey, "t1", "25");

      await scheduler.resetDeficit("t1");

      const deficit = await scheduler.getDeficit("t1");
      expect(deficit).toBe(0);

      await scheduler.close();
      await redis.quit();
    });
  });

  describe("queue selection", () => {
    redisTest("should return queues grouped by tenant", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const masterKey = keys.masterQueueKey(0);
      const now = Date.now();

      // Add queues for different tenants (all timestamps in the past)
      await redis.zadd(
        masterKey,
        now - 200,
        "tenant:t1:queue:q1",
        now - 100,
        "tenant:t1:queue:q2",
        now - 50,
        "tenant:t2:queue:q1"
      );

      const context: SchedulerContext = {
        getCurrentConcurrency: async () => 0,
        getConcurrencyLimit: async () => 100,
        isAtCapacity: async () => false,
        getQueueDescriptor: (queueId) => ({
          id: queueId,
          tenantId: keys.extractTenantId(queueId),
          metadata: {},
        }),
      };

      const result = await scheduler.selectQueues(masterKey, "consumer-1", context);

      // Should have both tenants
      const tenantIds = result.map((r) => r.tenantId);
      expect(tenantIds).toContain("t1");
      expect(tenantIds).toContain("t2");

      // t1 should have 2 queues
      const t1 = result.find((r) => r.tenantId === "t1");
      expect(t1?.queues).toHaveLength(2);

      await scheduler.close();
      await redis.quit();
    });

    redisTest("should filter out tenants at capacity", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const masterKey = keys.masterQueueKey(0);
      const now = Date.now();

      await redis.zadd(masterKey, now - 100, "tenant:t1:queue:q1", now - 50, "tenant:t2:queue:q1");

      const context: SchedulerContext = {
        getCurrentConcurrency: async () => 0,
        getConcurrencyLimit: async () => 100,
        isAtCapacity: async (_, groupId) => groupId === "t1", // t1 at capacity
        getQueueDescriptor: (queueId) => ({
          id: queueId,
          tenantId: keys.extractTenantId(queueId),
          metadata: {},
        }),
      };

      const result = await scheduler.selectQueues(masterKey, "consumer-1", context);

      // Only t2 should be returned
      const tenantIds = result.map((r) => r.tenantId);
      expect(tenantIds).not.toContain("t1");
      expect(tenantIds).toContain("t2");

      await scheduler.close();
      await redis.quit();
    });

    redisTest("should skip tenants with insufficient deficit", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const masterKey = keys.masterQueueKey(0);
      const now = Date.now();

      await redis.zadd(masterKey, now - 100, "tenant:t1:queue:q1", now - 50, "tenant:t2:queue:q1");

      // Set t1 deficit to 0 (no credits)
      const deficitKey = `test:drr:deficit`;
      await redis.hset(deficitKey, "t1", "0");

      const context: SchedulerContext = {
        getCurrentConcurrency: async () => 0,
        getConcurrencyLimit: async () => 100,
        isAtCapacity: async () => false,
        getQueueDescriptor: (queueId) => ({
          id: queueId,
          tenantId: keys.extractTenantId(queueId),
          metadata: {},
        }),
      };

      // First call adds quantum to both tenants
      // t1: 0 + 5 = 5, t2: 0 + 5 = 5
      const result = await scheduler.selectQueues(masterKey, "consumer-1", context);

      // Both should be returned (both have deficit >= 1 after quantum added)
      const tenantIds = result.map((r) => r.tenantId);
      expect(tenantIds).toContain("t1");
      expect(tenantIds).toContain("t2");

      await scheduler.close();
      await redis.quit();
    });

    redisTest("should order tenants by deficit (highest first)", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const masterKey = keys.masterQueueKey(0);
      const now = Date.now();

      await redis.zadd(
        masterKey,
        now - 300,
        "tenant:t1:queue:q1",
        now - 200,
        "tenant:t2:queue:q1",
        now - 100,
        "tenant:t3:queue:q1"
      );

      // Set different deficits
      const deficitKey = `test:drr:deficit`;
      await redis.hset(deficitKey, "t1", "10");
      await redis.hset(deficitKey, "t2", "30");
      await redis.hset(deficitKey, "t3", "20");

      const context: SchedulerContext = {
        getCurrentConcurrency: async () => 0,
        getConcurrencyLimit: async () => 100,
        isAtCapacity: async () => false,
        getQueueDescriptor: (queueId) => ({
          id: queueId,
          tenantId: keys.extractTenantId(queueId),
          metadata: {},
        }),
      };

      const result = await scheduler.selectQueues(masterKey, "consumer-1", context);

      // Should be ordered by deficit: t2 (35), t3 (25), t1 (15)
      // (original + quantum of 5)
      expect(result[0]?.tenantId).toBe("t2");
      expect(result[1]?.tenantId).toBe("t3");
      expect(result[2]?.tenantId).toBe("t1");

      await scheduler.close();
      await redis.quit();
    });
  });

  describe("get all deficits", () => {
    redisTest("should return all tenant deficits", { timeout: 10000 }, async ({ redisOptions }) => {
      keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
      const redis = createRedisClient(redisOptions);

      const scheduler = new DRRScheduler({
        redis: redisOptions,
        keys,
        quantum: 5,
        maxDeficit: 50,
      });

      const deficitKey = `test:drr:deficit`;
      await redis.hset(deficitKey, "t1", "10");
      await redis.hset(deficitKey, "t2", "20");
      await redis.hset(deficitKey, "t3", "30");

      const deficits = await scheduler.getAllDeficits();

      expect(deficits.get("t1")).toBe(10);
      expect(deficits.get("t2")).toBe(20);
      expect(deficits.get("t3")).toBe(30);

      await scheduler.close();
      await redis.quit();
    });
  });
});
