const crypto = require("node:crypto");
const { pool, withTransaction } = require("./db");

const STALE_LEASE_MINUTES = 30;
const DEFAULT_WAIT_TIMEOUT_MS = 25_000;
const MAX_WAIT_TIMEOUT_MS = 5 * 60 * 1000;

const providerWaitQueues = new Map();
const providerDrainPromises = new Map();

function normalizeWaitTimeoutMs(value) {
  const numeric = Number(value);
  if (!Number.isFinite(numeric) || numeric <= 0) {
    return 0;
  }

  return Math.min(Math.trunc(numeric), MAX_WAIT_TIMEOUT_MS);
}

function getProviderWaitQueue(provider) {
  const existingQueue = providerWaitQueues.get(provider);
  if (existingQueue) {
    return existingQueue;
  }

  const nextQueue = [];
  providerWaitQueues.set(provider, nextQueue);
  return nextQueue;
}

function compactProviderWaitQueue(provider) {
  const queue = providerWaitQueues.get(provider);
  if (!queue || queue.length === 0) {
    providerWaitQueues.delete(provider);
    return [];
  }

  const activeWaiters = queue.filter((waiter) => !waiter.settled);
  if (activeWaiters.length === 0) {
    providerWaitQueues.delete(provider);
    return [];
  }

  providerWaitQueues.set(provider, activeWaiters);
  return activeWaiters;
}

function removeWaiter(waiter) {
  const queue = providerWaitQueues.get(waiter.provider);
  if (!queue) {
    return;
  }

  const index = queue.indexOf(waiter);
  if (index >= 0) {
    queue.splice(index, 1);
  }

  compactProviderWaitQueue(waiter.provider);
}

function settleWaiter(waiter, outcome) {
  if (waiter.settled) {
    return;
  }

  waiter.settled = true;
  clearTimeout(waiter.timer);
  if (waiter.signal && waiter.abortHandler) {
    waiter.signal.removeEventListener("abort", waiter.abortHandler);
  }
  removeWaiter(waiter);

  if (outcome instanceof Error) {
    waiter.reject(outcome);
    return;
  }

  waiter.resolve(outcome);
}

function normalizeUserContext(user) {
  if (user && typeof user === "object") {
    return {
      userId: user.id,
      enterpriseId: user.enterpriseId ?? null,
      accountType: user.accountType ?? "personal",
    };
  }

  return {
    userId: user,
    enterpriseId: null,
    accountType: "personal",
  };
}

async function tryAcquireKey(provider, user, preauthResult) {
  const { userId, enterpriseId } = normalizeUserContext(user);
  const leaseToken = crypto.randomUUID();
  const estimatedCostCents = preauthResult?.estimatedCostCents || 0;

  return withTransaction(async (client) => {
    const { rows } = await client.query(
      `
      WITH candidate AS (
        SELECT id, api_key, label
        FROM api_keys
        WHERE provider = $1 AND enabled = 1 AND active_count < max_concurrency
        ORDER BY active_count ASC, id ASC
        FOR UPDATE SKIP LOCKED
        LIMIT 1
      ),
      updated AS (
        UPDATE api_keys
        SET active_count = active_count + 1,
            total_used = total_used + 1
        WHERE id = (SELECT id FROM candidate)
        RETURNING id, api_key, label
      )
      SELECT id, api_key, label FROM updated
    `,
      [provider],
    );

    const key = rows[0];
    if (!key) {
      return null;
    }

    await client.query(
      `
      INSERT INTO key_leases (key_id, user_id, lease_token, estimated_cost_cents, enterprise_id)
      VALUES ($1, $2, $3, $4, $5)
    `,
      [key.id, userId, leaseToken, estimatedCostCents, enterpriseId],
    );

    await client.query(
      `
      INSERT INTO usage_logs (user_id, enterprise_id, provider, key_id, action)
      VALUES ($1, $2, $3, $4, $5)
    `,
      [userId, enterpriseId, provider, key.id, "acquire"],
    );

    return {
      leaseToken,
      apiKey: key.api_key === "pool-slot" ? "" : key.api_key,
      provider,
      keyLabel: key.label || `Key #${key.id}`,
    };
  });
}

async function drainProviderWaitQueue(provider) {
  const existingDrain = providerDrainPromises.get(provider);
  if (existingDrain) {
    return existingDrain;
  }

  const drainPromise = (async () => {
    while (true) {
      const queue = compactProviderWaitQueue(provider);
      const nextWaiter = queue[0];
      if (!nextWaiter) {
        break;
      }

      let acquiredLease = null;
      try {
        acquiredLease = await tryAcquireKey(provider, nextWaiter.user, nextWaiter.preauthResult);
      } catch (error) {
        settleWaiter(nextWaiter, error instanceof Error ? error : new Error(String(error)));
        continue;
      }

      if (!acquiredLease) {
        break;
      }

      if (nextWaiter.settled || nextWaiter.signal?.aborted) {
        await releaseLeaseInternal(acquiredLease.leaseToken, nextWaiter.user, { skipDrain: true });
        continue;
      }

      settleWaiter(nextWaiter, acquiredLease);
    }
  })().finally(() => {
    providerDrainPromises.delete(provider);
  });

  providerDrainPromises.set(provider, drainPromise);
  await drainPromise;
}

function enqueueAcquireWaiter(provider, user, preauthResult, options = {}) {
  const waitTimeoutMs = normalizeWaitTimeoutMs(options.waitTimeoutMs ?? DEFAULT_WAIT_TIMEOUT_MS);
  if (waitTimeoutMs <= 0) {
    return Promise.resolve(null);
  }

  return new Promise((resolve, reject) => {
    if (options.signal?.aborted) {
      reject(new Error("Cancelled"));
      return;
    }

    const waiter = {
      provider,
      user,
      preauthResult,
      resolve,
      reject,
      settled: false,
      signal: options.signal,
      abortHandler: null,
      timer: null,
    };

    waiter.abortHandler = () => {
      settleWaiter(waiter, new Error("Cancelled"));
    };

    if (waiter.signal) {
      waiter.signal.addEventListener("abort", waiter.abortHandler, { once: true });
    }

    waiter.timer = setTimeout(() => {
      settleWaiter(waiter, null);
    }, waitTimeoutMs);

    getProviderWaitQueue(provider).push(waiter);
  });
}

async function cleanStaleLeases() {
  const cutoff = new Date(Date.now() - STALE_LEASE_MINUTES * 60 * 1000).toISOString();
  const { rows: stale } = await pool.query(
    `
    SELECT l.id, l.key_id, k.provider
    FROM key_leases l
    JOIN api_keys k ON k.id = l.key_id
    WHERE l.released_at IS NULL AND l.leased_at < $1
  `,
    [cutoff],
  );

  if (stale.length === 0) return 0;

  await withTransaction(async (client) => {
    for (const lease of stale) {
      await client.query("UPDATE key_leases SET released_at = NOW() WHERE id = $1", [lease.id]);
      await client.query(
        "UPDATE api_keys SET active_count = GREATEST(0, active_count - 1) WHERE id = $1",
        [lease.key_id],
      );
    }
  });

  const affectedProviders = Array.from(
    new Set(stale.map((lease) => lease.provider).filter(Boolean)),
  );
  await Promise.all(affectedProviders.map((provider) => drainProviderWaitQueue(provider)));

  return stale.length;
}

async function acquireKey(provider, user, preauthResult, options = {}) {
  await cleanStaleLeases();

  const immediateLease = await tryAcquireKey(provider, user, preauthResult);
  if (immediateLease) {
    return immediateLease;
  }

  return enqueueAcquireWaiter(provider, user, preauthResult, options);
}

async function releaseLeaseInternal(leaseToken, user, options = {}) {
  const { userId, enterpriseId } = normalizeUserContext(user);

  const releaseResult = await withTransaction(async (client) => {
    const { rows } = await client.query(
      `
      WITH candidate AS (
        SELECT l.id, l.key_id, k.provider
        FROM key_leases l
        JOIN api_keys k ON k.id = l.key_id
        WHERE l.lease_token = $1 AND l.released_at IS NULL
        FOR UPDATE SKIP LOCKED
        LIMIT 1
      ),
      released AS (
        UPDATE key_leases
        SET released_at = NOW()
        WHERE id = (SELECT id FROM candidate)
        RETURNING id, key_id
      )
      SELECT r.id, r.key_id, c.provider
      FROM released r
      JOIN candidate c ON c.key_id = r.key_id
    `,
      [leaseToken],
    );

    const lease = rows[0];
    if (!lease) {
      const { rows: existingRows } = await client.query(
        `
        SELECT l.key_id, l.released_at, k.provider
        FROM key_leases l
        JOIN api_keys k ON k.id = l.key_id
        WHERE l.lease_token = $1
        LIMIT 1
      `,
        [leaseToken],
      );

      const existing = existingRows[0];
      if (!existing) {
        return { released: false, notFound: true, alreadyReleased: false, provider: null };
      }

      return {
        released: false,
        notFound: false,
        alreadyReleased: Boolean(existing.released_at),
        provider: existing.provider || null,
      };
    }

    await client.query(
      "UPDATE api_keys SET active_count = GREATEST(0, active_count - 1) WHERE id = $1",
      [lease.key_id],
    );
    await client.query(
      `
      INSERT INTO usage_logs (user_id, enterprise_id, provider, key_id, action)
      VALUES ($1, $2, (SELECT provider FROM api_keys WHERE id = $3), $4, $5)
    `,
      [userId, enterpriseId, lease.key_id, lease.key_id, "release"],
    );

    return {
      released: true,
      notFound: false,
      alreadyReleased: false,
      provider: lease.provider || null,
    };
  });

  if (releaseResult.released && releaseResult.provider) {
    if (options.skipDrain) {
      return releaseResult;
    }
    await drainProviderWaitQueue(releaseResult.provider);
  }

  return releaseResult;
}

async function releaseKey(leaseToken, user) {
  return releaseLeaseInternal(leaseToken, user);
}

function getQueuedWaiterCount(provider) {
  return compactProviderWaitQueue(provider).length;
}

async function getKeyStatus(provider) {
  const { rows: keys } = await pool.query(
    `
    SELECT id, provider, label, max_concurrency, active_count, total_used, enabled
    FROM api_keys
    WHERE provider = $1
    ORDER BY id
  `,
    [provider],
  );

  const totalCapacity = keys.reduce((sum, key) => sum + (key.enabled ? key.max_concurrency : 0), 0);
  const totalActive = keys.reduce((sum, key) => sum + key.active_count, 0);

  return {
    provider,
    keys: keys.map((key) => ({
      id: key.id,
      label: key.label || `Key #${key.id}`,
      active: key.active_count,
      capacity: key.max_concurrency,
      totalUsed: key.total_used,
      enabled: !!key.enabled,
    })),
    totalActive,
    totalCapacity,
    queuedCount: getQueuedWaiterCount(provider),
    available: totalCapacity - totalActive,
  };
}

async function getAllStatus() {
  const { rows: providers } = await pool.query(
    "SELECT DISTINCT provider FROM api_keys WHERE enabled = 1",
  );
  const results = [];
  for (const providerRow of providers) {
    results.push(await getKeyStatus(providerRow.provider));
  }
  return results;
}

module.exports = {
  acquireKey,
  releaseKey,
  getKeyStatus,
  getAllStatus,
  cleanStaleLeases,
};