const assert = require("node:assert/strict");
const { createRequire } = require("node:module");

const nodeRequire = createRequire(__filename);

function loadLimiterWithPool(pool) {
  const dbPath = nodeRequire.resolve("../src/db");
  const limiterPath = nodeRequire.resolve("../src/providerPollLimiter");
  const originalDbModule = nodeRequire.cache[dbPath];
  const originalLimiterModule = nodeRequire.cache[limiterPath];

  delete nodeRequire.cache[limiterPath];
  nodeRequire.cache[dbPath] = {
    id: dbPath,
    filename: dbPath,
    loaded: true,
    exports: { pool },
  };

  return {
    limiter: nodeRequire("../src/providerPollLimiter"),
    restore() {
      delete nodeRequire.cache[limiterPath];
      if (originalLimiterModule) nodeRequire.cache[limiterPath] = originalLimiterModule;
      if (originalDbModule) nodeRequire.cache[dbPath] = originalDbModule;
      else delete nodeRequire.cache[dbPath];
    },
  };
}

function createPool(options = {}) {
  const calls = [];
  return {
    calls,
    async query(sql, params = []) {
      calls.push({ sql, params });
      if (/CREATE TABLE IF NOT EXISTS generation_provider_poll_slots/i.test(sql)) return { rows: [] };
      if (/WITH candidate AS/i.test(sql)) {
        if (options.noAvailableSlot) return { rows: [] };
        return { rows: [{ scope: params[0], slot_no: 2 }] };
      }
      if (/DELETE FROM generation_provider_poll_slots/i.test(sql)) return { rows: [] };
      throw new Error(`Unexpected SQL: ${sql}`);
    },
  };
}

(async () => {
  const previousLimit = process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY;
  process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY = "3";

  const pool = createPool();
  const { limiter, restore } = loadLimiterWithPool(pool);
  try {
    const outcome = await limiter.withProviderPollSlot(101, async () => "polled");

    assert.equal(outcome.acquired, true);
    assert.equal(outcome.value, "polled");

    const acquireCall = pool.calls.find((call) => /WITH candidate AS/i.test(call.sql));
    assert.equal(acquireCall.params[1], 3);
    assert.equal(acquireCall.params[3], 101);

    const releaseCall = pool.calls.find((call) => /DELETE FROM generation_provider_poll_slots/i.test(call.sql));
    assert.equal(releaseCall.params[0], acquireCall.params[0]);
    assert.equal(releaseCall.params[1], 2);
    assert.equal(releaseCall.params[2], acquireCall.params[2]);
  } finally {
    if (previousLimit === undefined) delete process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY;
    else process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY = previousLimit;
    restore();
  }

  const saturatedPool = createPool({ noAvailableSlot: true });
  const { limiter: saturatedLimiter, restore: restoreSaturated } = loadLimiterWithPool(saturatedPool);
  try {
    let called = false;
    const outcome = await saturatedLimiter.withProviderPollSlot(202, async () => {
      called = true;
      return "should-not-run";
    });

    assert.equal(outcome.acquired, false);
    assert.equal(outcome.value, undefined);
    assert.equal(called, false);
    assert.equal(
      saturatedPool.calls.some((call) => /DELETE FROM generation_provider_poll_slots/i.test(call.sql)),
      false,
    );
  } finally {
    restoreSaturated();
  }
})().catch((error) => {
  console.error(error);
  process.exitCode = 1;
});