server/src/services/ratelimit.ts

// Sliding window rate limit tracker with SQLite persistence.

import { getDb } from '../db/index.js';

interface Window {
  timestamps: number[];
  tokenCount: number;
  tokenTimestamps: { ts: number; tokens: number }[];
}

// Key format: "platform:modelId:keyId:type" where type is rpm|rpd|tpm|tpd
const windows = new Map<string, Window>();
type RateLimitDb = ReturnType<typeof getDb>;
type UsageKind = 'request' | 'tokens';

function getWindow(key: string): Window {
  let w = windows.get(key);
  if (!w) {
    w = { timestamps: [], tokenCount: 0, tokenTimestamps: [] };
    windows.set(key, w);
  }
  return w;
}

function pruneTimestamps(timestamps: number[], windowMs: number, now: number): number[] {
  const cutoff = now - windowMs;
  return timestamps.filter(ts => ts > cutoff);
}

const MINUTE = 60 * 1000;
const DAY = 24 * 60 * MINUTE;

async function withDb<T>(fn: (db: RateLimitDb) => Promise<T>): Promise<T | undefined> {
  try {
    return await fn(getDb());
  } catch {
    return undefined;
  }
}

async function recordUsage(
  platform: string,
  modelId: string,
  keyId: number,
  kind: UsageKind,
  tokens: number,
  now: number,
): Promise<void> {
  await withDb(async db => {
    await db.run(`
      INSERT INTO rate_limit_usage (platform, model_id, key_id, kind, tokens, created_at_ms)
      VALUES (?, ?, ?, ?, ?, ?)
    `, [platform, modelId, keyId, kind, tokens, now]);
    await db.run('DELETE FROM rate_limit_usage WHERE created_at_ms <= ?', [now - DAY]);
  });
}

async function countPersistedRequests(
  platform: string,
  modelId: string,
  keyId: number,
  windowMs: number,
  now: number,
): Promise<number | undefined> {
  return await withDb(async db => {
    const row = await db.get<{ used: number }>(`
      SELECT COUNT(*) AS used
        FROM rate_limit_usage
       WHERE platform = ?
         AND model_id = ?
         AND key_id = ?
         AND kind = 'request'
         AND created_at_ms > ?
    `, [platform, modelId, keyId, now - windowMs]);
    return row?.used ?? 0;
  });
}

async function sumPersistedTokens(
  platform: string,
  modelId: string,
  keyId: number,
  windowMs: number,
  now: number,
): Promise<number | undefined> {
  return await withDb(async db => {
    const row = await db.get<{ used: number }>(`
      SELECT COALESCE(SUM(tokens), 0) AS used
        FROM rate_limit_usage
       WHERE platform = ?
         AND model_id = ?
         AND key_id = ?
         AND kind = 'tokens'
         AND created_at_ms > ?
    `, [platform, modelId, keyId, now - windowMs]);
    return row?.used ?? 0;
  });
}

function memoryRequestCount(key: string, windowMs: number, now: number): number {
  const w = getWindow(key);
  w.timestamps = pruneTimestamps(w.timestamps, windowMs, now);
  return w.timestamps.length;
}

function memoryTokenCount(key: string, windowMs: number, now: number): number {
  const w = getWindow(key);
  w.tokenTimestamps = w.tokenTimestamps.filter(t => t.ts > now - windowMs);
  return w.tokenTimestamps.reduce((sum, t) => sum + t.tokens, 0);
}

async function requestCount(
  platform: string,
  modelId: string,
  keyId: number,
  windowMs: number,
  now: number,
): Promise<number> {
  const persisted = await countPersistedRequests(platform, modelId, keyId, windowMs, now);
  if (persisted !== undefined) return persisted;
  const type = windowMs === MINUTE ? 'rpm' : 'rpd';
  return memoryRequestCount(`${platform}:${modelId}:${keyId}:${type}`, windowMs, now);
}

async function tokenCount(
  platform: string,
  modelId: string,
  keyId: number,
  windowMs: number,
  now: number,
): Promise<number> {
  const persisted = await sumPersistedTokens(platform, modelId, keyId, windowMs, now);
  if (persisted !== undefined) return persisted;
  const type = windowMs === MINUTE ? 'tpm' : 'tpd';
  return memoryTokenCount(`${platform}:${modelId}:${keyId}:${type}`, windowMs, now);
}

export async function canMakeRequest(
  platform: string,
  modelId: string,
  keyId: number,
  limits: { rpm: number | null; rpd: number | null; tpm: number | null; tpd: number | null },
): Promise<boolean> {
  const now = Date.now();

  if (limits.rpm !== null) {
    if ((await requestCount(platform, modelId, keyId, MINUTE, now)) >= limits.rpm) return false;
  }

  if (limits.rpd !== null) {
    if ((await requestCount(platform, modelId, keyId, DAY, now)) >= limits.rpd) return false;
  }

  return true;
}

export async function canUseTokens(
  platform: string,
  modelId: string,
  keyId: number,
  estimatedTokens: number,
  limits: { tpm: number | null; tpd: number | null },
): Promise<boolean> {
  const now = Date.now();

  if (limits.tpm !== null) {
    const used = await tokenCount(platform, modelId, keyId, MINUTE, now);
    if (used + estimatedTokens > limits.tpm) return false;
  }

  if (limits.tpd !== null) {
    const used = await tokenCount(platform, modelId, keyId, DAY, now);
    if (used + estimatedTokens > limits.tpd) return false;
  }

  return true;
}

// ── Provider-wide daily request caps (#162) ──
// Some providers enforce one daily REQUEST quota across the WHOLE account,
// shared by every model — not per model. OpenRouter's free tier is the classic
// case: ~1000 requests/day total (50/day if you've bought <10 credits) no
// matter how many different free models you spread them across. The
// per-(platform,model,key) rpd ledger can't see that, so without a provider-wide
// gate the router happily fires (models × rpd) requests and earns surprise 429s.
//
// Defaults below; override per provider with an env var, e.g.
//   PROVIDER_DAILY_REQUEST_CAP_OPENROUTER=50   (set 0 to disable the cap)
const DEFAULT_PROVIDER_DAILY_REQUEST_CAPS: Record<string, number> = {
  openrouter: 1000,
};

export function getProviderDailyRequestCap(platform: string): number | null {
  const raw = process.env[`PROVIDER_DAILY_REQUEST_CAP_${platform.toUpperCase()}`];
  if (raw !== undefined && raw.trim() !== '') {
    const n = Number(raw);
    if (Number.isFinite(n) && n >= 0) return n === 0 ? null : n;
  }
  return DEFAULT_PROVIDER_DAILY_REQUEST_CAPS[platform] ?? null;
}

async function countPersistedProviderRequests(
  platform: string,
  keyId: number,
  windowMs: number,
  now: number,
): Promise<number | undefined> {
  return await withDb(async db => {
    const row = await db.get<{ used: number }>(`
      SELECT COUNT(*) AS used
        FROM rate_limit_usage
       WHERE platform = ?
         AND key_id = ?
         AND kind = 'request'
         AND created_at_ms > ?
    `, [platform, keyId, now - windowMs]);
    return row?.used ?? 0;
  });
}

// Total requests today for a provider account+key, summed across every model.
export async function providerDailyRequestCount(platform: string, keyId: number, now = Date.now()): Promise<number> {
  const persisted = await countPersistedProviderRequests(platform, keyId, DAY, now);
  if (persisted !== undefined) return persisted;
  // DB-unavailable fallback: sum the per-model rpd windows for this platform+key.
  // Window key format is "platform:modelId:keyId:rpd" (modelId may contain ':').
  let total = 0;
  for (const [key, w] of windows) {
    if (key.startsWith(`${platform}:`) && key.endsWith(`:${keyId}:rpd`)) {
      total += pruneTimestamps(w.timestamps, DAY, now).length;
    }
  }
  return total;
}

// False when this provider account+key has hit its shared daily request cap, so
// the router skips every model on that provider for this key until UTC-ish reset.
export async function canUseProvider(platform: string, keyId: number, now = Date.now()): Promise<boolean> {
  const cap = getProviderDailyRequestCap(platform);
  if (cap === null) return true;
  return (await providerDailyRequestCount(platform, keyId, now)) < cap;
}

export async function recordRequest(platform: string, modelId: string, keyId: number): Promise<void> {
  const now = Date.now();

  const rpmKey = `${platform}:${modelId}:${keyId}:rpm`;
  getWindow(rpmKey).timestamps.push(now);

  const rpdKey = `${platform}:${modelId}:${keyId}:rpd`;
  getWindow(rpdKey).timestamps.push(now);

  await recordUsage(platform, modelId, keyId, 'request', 0, now);
}

export async function recordTokens(
  platform: string,
  modelId: string,
  keyId: number,
  tokens: number,
): Promise<void> {
  const now = Date.now();

  const tpmKey = `${platform}:${modelId}:${keyId}:tpm`;
  getWindow(tpmKey).tokenTimestamps.push({ ts: now, tokens });

  const tpdKey = `${platform}:${modelId}:${keyId}:tpd`;
  getWindow(tpdKey).tokenTimestamps.push({ ts: now, tokens });

  await recordUsage(platform, modelId, keyId, 'tokens', tokens, now);
}

// Cooldown: when a provider returns 429, block that model+key for a period
const cooldowns = new Map<string, number>(); // key -> expiry timestamp

// Escalating cooldown: track hits per key over a rolling 24h window so a
// daily-quota exhaustion (OpenRouter free: 50/day, Cohere free: 33/day, etc.)
// quarantines the key for the rest of the day instead of looping through
// the 2-minute cooldown 20 times per request and consuming every fallback slot.
// In-memory only — state resets on restart, which is fine (a clean restart
// will re-escalate on the next 429 if the quota is genuinely exhausted).
const cooldownHits = new Map<string, number[]>(); // key -> timestamps of recent cooldown set events
const HOUR = 60 * MINUTE;
const COOLDOWN_DURATIONS = [
  2 * MINUTE,   // 1st hit in 24h
  10 * MINUTE,  // 2nd
  HOUR,         // 3rd
  DAY,          // 4th and beyond
];

export function getNextCooldownDuration(platform: string, modelId: string, keyId: number): number {
  const key = `${platform}:${modelId}:${keyId}`;
  const now = Date.now();
  const hits = (cooldownHits.get(key) ?? []).filter(t => t > now - DAY);
  hits.push(now);
  cooldownHits.set(key, hits);
  const idx = Math.min(hits.length - 1, COOLDOWN_DURATIONS.length - 1);
  return COOLDOWN_DURATIONS[idx]!;
}

// Short cooldown for a transient (per-minute) 429 — recovers within ~one window.
const TRANSIENT_COOLDOWN_MS = 90 * 1000;

// Long cooldown for a 402 Payment Required (provider/key out of credits). Unlike
// a 429, this won't clear on the next minute/day window — it needs a top-up or
// billing reset. Bench the model+key for a full day so the router fails over to
// other providers instead of re-hammering a dead key every retry. Re-escalates
// on the next 402 after expiry if still unpaid; a restart re-benches on first hit.
export const PAYMENT_REQUIRED_COOLDOWN_MS = DAY;

// Decide how long to bench a model+key after an upstream 429. Escalate to the
// long quarantine (getNextCooldownDuration, up to 24h) ONLY when the model is
// genuinely at its DAILY limit (RPD or TPD) — that won't recover until the
// provider's daily reset, so a long bench avoids hammering a truly-dead key.
//
// A transient RPM/TPM 429 gets a short fixed cooldown and does NOT count toward
// escalation. This is the common case for providers with a tight per-minute
// token budget but a large daily quota — e.g. groq gpt-oss-120b has rpd=1000
// yet tpm=8000, so a single burst of large prompts 429s on TPM while the daily
// quota is barely touched. Without this split, those transient bursts escalated
// (2m → 10m → 1h → 24h) and quarantined a perfectly healthy provider for the
// rest of the day. Daily counters are persisted (countPersistedRequests /
// sumPersistedTokens), so this verdict is stable across restarts.
export async function getCooldownDurationForLimit(
  platform: string,
  modelId: string,
  keyId: number,
  limits: { rpd: number | null; tpd: number | null },
): Promise<number> {
  const now = Date.now();
  const rpdExhausted =
    limits.rpd !== null && (await requestCount(platform, modelId, keyId, DAY, now)) >= limits.rpd;
  const tpdExhausted =
    limits.tpd !== null && (await tokenCount(platform, modelId, keyId, DAY, now)) >= limits.tpd;
  if (rpdExhausted || tpdExhausted) {
    return getNextCooldownDuration(platform, modelId, keyId);
  }
  return TRANSIENT_COOLDOWN_MS;
}

async function persistedCooldownExpiry(
  platform: string,
  modelId: string,
  keyId: number,
): Promise<number | null | undefined> {
  return await withDb(async db => {
    const row = await db.get<{ expires_at_ms: number }>(`
      SELECT expires_at_ms
        FROM rate_limit_cooldowns
       WHERE platform = ?
         AND model_id = ?
         AND key_id = ?
    `, [platform, modelId, keyId]);
    return row?.expires_at_ms ?? null;
  });
}

async function persistCooldown(platform: string, modelId: string, keyId: number, expiresAtMs: number): Promise<void> {
  await withDb(async db => {
    await db.run(`
      INSERT INTO rate_limit_cooldowns (platform, model_id, key_id, expires_at_ms)
      VALUES (?, ?, ?, ?)
      ON CONFLICT(platform, model_id, key_id)
      DO UPDATE SET expires_at_ms = excluded.expires_at_ms
    `, [platform, modelId, keyId, expiresAtMs]);
  });
}

async function clearPersistedCooldown(platform: string, modelId: string, keyId: number): Promise<void> {
  await withDb(async db => {
    await db.run(`
      DELETE FROM rate_limit_cooldowns
       WHERE platform = ?
         AND model_id = ?
         AND key_id = ?
    `, [platform, modelId, keyId]);
  });
}

export async function setCooldown(platform: string, modelId: string, keyId: number, durationMs = 60_000): Promise<void> {
  const key = `${platform}:${modelId}:${keyId}:cooldown`;
  const expiresAtMs = Date.now() + durationMs;
  cooldowns.set(key, expiresAtMs);
  await persistCooldown(platform, modelId, keyId, expiresAtMs);
}

export async function isOnCooldown(platform: string, modelId: string, keyId: number): Promise<boolean> {
  const key = `${platform}:${modelId}:${keyId}:cooldown`;
  const now = Date.now();
  const persistedExpiry = await persistedCooldownExpiry(platform, modelId, keyId);
  if (persistedExpiry !== undefined && persistedExpiry !== null) {
    if (now > persistedExpiry) {
      cooldowns.delete(key);
      await clearPersistedCooldown(platform, modelId, keyId);
      return false;
    }
    cooldowns.set(key, persistedExpiry);
    return true;
  }

  const expiry = cooldowns.get(key);
  if (!expiry) return false;
  if (now > expiry) {
    cooldowns.delete(key);
    return false;
  }
  return true;
}

export async function getRateLimitStatus(
  platform: string,
  modelId: string,
  keyId: number,
  limits: { rpm: number | null; rpd: number | null; tpm: number | null; tpd: number | null },
): Promise<{
  rpm: { used: number; limit: number | null };
  rpd: { used: number; limit: number | null };
  tpm: { used: number; limit: number | null };
}> {
  const now = Date.now();

  return {
    rpm: { used: await requestCount(platform, modelId, keyId, MINUTE, now), limit: limits.rpm },
    rpd: { used: await requestCount(platform, modelId, keyId, DAY, now), limit: limits.rpd },
    tpm: { used: await tokenCount(platform, modelId, keyId, MINUTE, now), limit: limits.tpm },
  };
}