OmniRoute/src/lib/quota/saturationSignals.ts

/**
 * saturationSignals.ts — Read the current global saturation signal (0..1)
 * for a provider/connection/dimension combination.
 *
 * Strategy (per provider):
 *   codex            → codexQuotaFetcher (dual 5h + weekly window)
 *   bailian          → bailianQuotaFetcher (triple 5h + weekly + monthly window)
 *   anthropic/claude → REAL plan-window utilization from GET /api/oauth/usage
 *                      (the same path usage.ts already uses): window "5h" →
 *                      five_hour.utilization, "weekly" → seven_day.utilization.
 *                      Falls back to the per-minute REQUEST rate-limit headers
 *                      (anthropic-ratelimit-requests-*) only when no OAuth plan
 *                      window is available (e.g. API-key Claude connections).
 *   default          → getUsageForProvider (open-sse/services/usage.ts)
 *
 * Cache: in-memory Map, TTL = 30 seconds. The 30s TTL is what keeps the
 * NON-OFFICIAL, rate-limited oauth/usage endpoint from being polled per
 * request (it returns 429 under load) — never call it on the hot path without
 * this cache. usage.ts adds its own 429 cooldown on top.
 * Fail-open: on any error, return 0 (generous mode) and log pino.warn.
 * Hard Rule #12: no stack traces propagated to return values.
 *
 * Part of: Group B — Quota Sharing Engine (plan 22, frente F6).
 */

import { createLogger } from "@/shared/utils/logger";
import { updateAccountBuckets, type ClaudeUsageResult } from "./accountBuckets";
import type { QuotaUnit, QuotaWindow } from "./dimensions";

const log = createLogger("quota:saturation");

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

interface CacheEntry {
  value: number; // 0..1
  ts: number; // epoch ms
}

interface DimensionSpec {
  unit: QuotaUnit;
  window: QuotaWindow;
}

// ---------------------------------------------------------------------------
// In-memory cache (Map<cacheKey, CacheEntry>)
// ---------------------------------------------------------------------------

const CACHE_TTL_MS = 30_000; // 30 seconds

const _cache = new Map<string, CacheEntry>();

// ---------------------------------------------------------------------------
// Rate-limit header cache (populated by response handlers)
// ---------------------------------------------------------------------------

interface RateLimitHeaderEntry {
  limit: number;
  remaining: number;
  ts: number;
}

/**
 * TOKEN rate-limit header snapshot. Unlike the per-minute REQUEST headers, the
 * token headers ride on EVERY upstream response (success too), so they enable
 * proactive throttling before a 429. `resetAt` is the upstream reset normalized
 * to epoch ms (Anthropic RFC3339 → Date.parse; OpenAI duration → now + secs),
 * or null when the upstream sent no reset header.
 */
interface TokenHeaderEntry {
  limit: number;
  remaining: number;
  resetAt: number | null; // epoch ms, normalized; null when unknown
  ts: number;
}

const _rateLimitHeaders = new Map<string, RateLimitHeaderEntry>();
const _tokenHeaders = new Map<string, TokenHeaderEntry>();
const RL_HEADER_TTL_MS = 5 * 60 * 1000; // 5 minutes

/** Test-only: clear the rate-limit + token header caches between asserts. */
export function _clearRateLimitHeaders(): void {
  _rateLimitHeaders.clear();
  _tokenHeaders.clear();
}

/**
 * Parse an OpenAI rate-limit reset DURATION string into milliseconds.
 * OpenAI reports token/request resets as Go-style durations, e.g. "6m0s",
 * "1s", "1h30m15s", "1.5s". Returns null when unparseable.
 */
function parseDurationMs(raw: string): number | null {
  const s = raw.trim();
  if (!s) return null;
  // Bounded, non-overlapping segments to avoid ReDoS (PII learning #1).
  const re = /(\d+(?:\.\d+)?)(ms|h|m|s)/g;
  let total = 0;
  let matched = false;
  let m: RegExpExecArray | null;
  while ((m = re.exec(s)) !== null) {
    matched = true;
    const value = Number(m[1]);
    if (!Number.isFinite(value)) return null;
    switch (m[2]) {
      case "h":
        total += value * 3_600_000;
        break;
      case "m":
        total += value * 60_000;
        break;
      case "s":
        total += value * 1000;
        break;
      case "ms":
        total += value;
        break;
    }
  }
  return matched ? total : null;
}

/**
 * Normalize a token-reset header value to epoch ms.
 *   - Anthropic: RFC3339 timestamp ("2026-01-01T00:00:30Z") → Date.parse.
 *   - OpenAI: duration ("6m0s") → now + parsed ms.
 * Returns null when absent or unparseable.
 */
function normalizeTokenReset(raw: string | undefined, nowMs: number): number | null {
  if (!raw) return null;
  const s = raw.trim();
  if (!s) return null;
  // RFC3339 / ISO-8601 if it looks like a date (YYYY-MM-DD with a time sep).
  if (/\d{4}-\d{2}-\d{2}/.test(s) && /[T:]/.test(s)) {
    const t = Date.parse(s);
    if (Number.isFinite(t)) return t;
  }
  // Otherwise treat as an OpenAI-style duration relative to now.
  const durMs = parseDurationMs(s);
  return durMs === null ? null : nowMs + durMs;
}

/**
 * Pick the first present {limit, remaining[, reset]} triple from a list of
 * header-key candidates, in priority order. Returns null when none are usable
 * (missing keys, non-finite, or limit <= 0).
 */
function pickTokenTriple(
  headers: Record<string, string>,
  candidates: Array<{ limit: string; remaining: string; reset: string }>
): { limit: number; remaining: number; reset: string | undefined } | null {
  for (const c of candidates) {
    const limitStr = headers[c.limit];
    const remainingStr = headers[c.remaining];
    if (limitStr === undefined || remainingStr === undefined) continue;
    const limit = Number(limitStr);
    const remaining = Number(remainingStr);
    if (Number.isFinite(limit) && limit > 0 && Number.isFinite(remaining)) {
      return { limit, remaining, reset: headers[c.reset] };
    }
  }
  return null;
}

/**
 * Store rate-limit headers from an upstream response for saturation signal use.
 * Called by the response handler after a successful request.
 *
 * Captures two independent signals, both keyed `${provider}:${connectionId}`:
 *   - REQUEST headers (per-minute RPM burst) — legacy, anthropic fallback.
 *   - TOKEN headers (per-window TPM) — universal proactive saturation; present
 *     on EVERY response, so we can throttle before the 429.
 */
export function storeRateLimitHeaders(
  connectionId: string,
  provider: string,
  headers: Record<string, string>
): void {
  const key = `${provider}:${connectionId}`;

  // ── REQUEST headers (legacy path, unchanged) ──────────────────────────────
  // Anthropic: anthropic-ratelimit-requests-limit / anthropic-ratelimit-requests-remaining
  const limitStr =
    headers["anthropic-ratelimit-requests-limit"] ??
    headers["x-ratelimit-limit-requests"] ??
    headers["x-ratelimit-limit"];
  const remainingStr =
    headers["anthropic-ratelimit-requests-remaining"] ??
    headers["x-ratelimit-remaining-requests"] ??
    headers["x-ratelimit-remaining"];

  if (limitStr && remainingStr) {
    const limit = Number(limitStr);
    const remaining = Number(remainingStr);
    if (Number.isFinite(limit) && limit > 0 && Number.isFinite(remaining)) {
      _rateLimitHeaders.set(key, { limit, remaining, ts: Date.now() });
    }
  }

  // ── TOKEN headers (universal proactive saturation) ────────────────────────
  // Anthropic base tokens, then OpenAI x-ratelimit-*-tokens, then anthropic
  // input/output variants as a fallback. First usable triple wins.
  const tokenTriple = pickTokenTriple(headers, [
    {
      limit: "anthropic-ratelimit-tokens-limit",
      remaining: "anthropic-ratelimit-tokens-remaining",
      reset: "anthropic-ratelimit-tokens-reset",
    },
    {
      limit: "x-ratelimit-limit-tokens",
      remaining: "x-ratelimit-remaining-tokens",
      reset: "x-ratelimit-reset-tokens",
    },
    {
      limit: "anthropic-ratelimit-input-tokens-limit",
      remaining: "anthropic-ratelimit-input-tokens-remaining",
      reset: "anthropic-ratelimit-input-tokens-reset",
    },
    {
      limit: "anthropic-ratelimit-output-tokens-limit",
      remaining: "anthropic-ratelimit-output-tokens-remaining",
      reset: "anthropic-ratelimit-output-tokens-reset",
    },
  ]);

  if (tokenTriple) {
    const now = Date.now();
    _tokenHeaders.set(key, {
      limit: tokenTriple.limit,
      remaining: tokenTriple.remaining,
      resetAt: normalizeTokenReset(tokenTriple.reset, now),
      ts: now,
    });
  }
}

/**
 * Token-header saturation signal for a (provider, connectionId).
 * Returns `{ saturation, resetAt }` where saturation = 1 − remaining/limit
 * (clamped 0..1) and resetAt is the normalized epoch-ms reset (or null), or
 * null when no fresh token-header data exists.
 */
export function getTokenHeaderSaturation(
  provider: string,
  connectionId: string
): { saturation: number; resetAt: number | null } | null {
  const entry = _tokenHeaders.get(`${provider}:${connectionId}`);
  if (!entry || Date.now() - entry.ts > RL_HEADER_TTL_MS) return null;
  if (!(entry.limit > 0)) return null;
  const used = entry.limit - entry.remaining;
  const saturation = Math.min(1, Math.max(0, used / entry.limit));
  return { saturation, resetAt: entry.resetAt };
}

function cacheKey(connectionId: string, provider: string, dim: DimensionSpec): string {
  return `${provider}:${connectionId}:${dim.unit}:${dim.window}`;
}

// Exported for test reset
export function _clearSaturationCache(): void {
  _cache.clear();
}

// ---------------------------------------------------------------------------
// Provider-specific extractors
// ---------------------------------------------------------------------------

/**
 * Map QuotaWindow to the Codex window keys returned by the fetcher.
 */
function codexWindowKey(window: QuotaWindow): string {
  switch (window) {
    case "5h":
      return "session"; // CODEX_WINDOW_SESSION
    case "weekly":
      return "weekly"; // CODEX_WINDOW_WEEKLY
    default:
      return "session";
  }
}

async function fetchCodexSaturation(
  connectionId: string,
  dim: DimensionSpec,
  connection?: Record<string, unknown>
): Promise<number> {
  // Dynamic import — codexQuotaFetcher lives in open-sse workspace
  const mod = await import("@omniroute/open-sse/services/codexQuotaFetcher");
  // #6379: pass the loaded connection snapshot through so fetchCodexQuota can
  // read its accessToken/workspaceId even when this connection was never
  // registered via registerCodexConnection() (e.g. during headroom ranking,
  // which runs BEFORE any request is dispatched for the candidate). Without
  // this, fetchCodexQuota returns null for every candidate and saturation
  // fails open to 0 across the board — headroom then can't tell accounts
  // apart and keeps the original combo order.
  const quota = await mod.fetchCodexQuota(connectionId, connection);
  if (!quota) return 0;

  const winKey = codexWindowKey(dim.window);
  const windows = quota.windows as Record<string, { percentUsed: number } | undefined>;
  const win = windows[winKey];
  if (win && typeof win.percentUsed === "number") {
    return Math.min(1, Math.max(0, win.percentUsed));
  }
  // fallback to overall percentUsed
  return Math.min(1, Math.max(0, quota.percentUsed ?? 0));
}

async function fetchBailianSaturation(connectionId: string, dim: DimensionSpec): Promise<number> {
  const mod = await import("@omniroute/open-sse/services/bailianQuotaFetcher");
  const quota = await mod.fetchBailianQuota(connectionId);
  if (!quota) return 0;

  const q = quota as unknown as Record<string, unknown>;
  let pct = 0;
  switch (dim.window) {
    case "5h":
      pct = ((q.window5h as Record<string, unknown>)?.percentUsed as number) ?? 0;
      break;
    case "weekly":
      pct = ((q.windowWeekly as Record<string, unknown>)?.percentUsed as number) ?? 0;
      break;
    case "monthly":
      pct = ((q.windowMonthly as Record<string, unknown>)?.percentUsed as number) ?? 0;
      break;
    default:
      pct = (q.percentUsed as number) ?? 0;
  }
  return Math.min(1, Math.max(0, pct));
}

/**
 * Per-minute REQUEST rate-limit headers fallback. Used only when the OAuth
 * plan-window utilization is unavailable (e.g. API-key Claude connections that
 * have no /api/oauth/usage data). This signal reflects TPM/RPM bursts, NOT the
 * 5h/weekly plan window, so it is a weak last resort.
 */
function anthropicHeaderSaturation(connectionId: string): number {
  const entry = _rateLimitHeaders.get(`anthropic:${connectionId}`);
  if (!entry || Date.now() - entry.ts > RL_HEADER_TTL_MS) return 0;

  const used = entry.limit - entry.remaining;
  return Math.min(1, Math.max(0, used / entry.limit));
}

/**
 * Injectable seam (DB lookup + usage fetch) so the oauth/usage plan-window path
 * is unit-testable without touching the DB or the network. Defaults are wired
 * lazily to the real implementations inside fetchAnthropicSaturation.
 */
interface AnthropicSaturationDeps {
  /** Resolve a connection (with decrypted accessToken/authType) by id. */
  loadConnection: (connectionId: string) => Promise<Record<string, unknown> | null>;
  /** Fetch usage for the connection (delegates to getUsageForProvider). */
  fetchUsage: (conn: Record<string, unknown>) => Promise<unknown>;
}

let _anthropicDepsOverride: AnthropicSaturationDeps | null = null;

/** Test-only: inject ({loadConnection, fetchUsage}); pass null to restore. */
export function __setAnthropicSaturationDepsForTests(deps: AnthropicSaturationDeps | null): void {
  _anthropicDepsOverride = deps;
}

async function defaultAnthropicDeps(): Promise<AnthropicSaturationDeps> {
  const [localDbMod, usageMod] = await Promise.all([
    import("@/lib/db/readCache"),
    import("@omniroute/open-sse/services/usage"),
  ]);
  return {
    loadConnection: (connectionId) =>
      localDbMod.getCachedProviderConnectionById(connectionId) as Promise<Record<
        string,
        unknown
      > | null>,
    fetchUsage: (conn) =>
      usageMod.getUsageForProvider(conn as Parameters<typeof usageMod.getUsageForProvider>[0]),
  };
}

/**
 * Map a QuotaWindow to the Claude usage quota key produced by getClaudeUsage
 * (usage.ts). "session (5h)" carries five_hour.utilization and "weekly (7d)"
 * carries seven_day.utilization.
 */
function claudeUsageKeyForWindow(window: QuotaWindow): string | null {
  switch (window) {
    case "5h":
      return "session (5h)";
    case "weekly":
    case "monthly":
      // Anthropic exposes a 7-day plan window, not a monthly one — treat the
      // longer requested window as the weekly plan saturation.
      return "weekly (7d)";
    default:
      return null;
  }
}

/**
 * Extract the plan utilization (0..1) for the requested window from a
 * getClaudeUsage() result, or null when the OAuth plan window is unavailable
 * (e.g. legacy/admin API-key shape with no per-window quotas).
 */
function planUtilizationFromUsage(usage: unknown, window: QuotaWindow): number | null {
  if (!usage || typeof usage !== "object") return null;
  const quotas = (usage as Record<string, unknown>).quotas;
  if (!quotas || typeof quotas !== "object" || Array.isArray(quotas)) return null;

  const key = claudeUsageKeyForWindow(window);
  if (!key) return null;
  const entry = (quotas as Record<string, unknown>)[key];
  if (!entry || typeof entry !== "object") return null;

  // getClaudeUsage stores `used` = utilization (% used, 0..100).
  const used = (entry as Record<string, unknown>).used;
  if (typeof used !== "number" || !Number.isFinite(used)) return null;
  return Math.min(1, Math.max(0, used / 100));
}

async function fetchAnthropicSaturation(connectionId: string, dim: DimensionSpec): Promise<number> {
  // Try the REAL plan-window utilization first (5h / weekly), via the same
  // /api/oauth/usage path usage.ts already uses. This is the signal fairShare
  // actually needs for Claude Pro/Max — the per-minute request headers do not
  // reflect the plan window. Any failure here falls back to the header path,
  // and ultimately fails open (0).
  try {
    const deps = _anthropicDepsOverride ?? (await defaultAnthropicDeps());
    const conn = await deps.loadConnection(connectionId);
    // Only OAuth connections have plan-window usage; API-key Claude does not.
    const hasOauthToken =
      !!conn &&
      typeof conn.accessToken === "string" &&
      conn.accessToken.length > 0 &&
      (conn.authType === undefined || conn.authType === "oauth");
    if (hasOauthToken) {
      const usage = await deps.fetchUsage(conn as Record<string, unknown>);
      // Update the per-window saturating buckets (Phase 3 #3) off the request
      // hot path — this runs behind the 30s saturation cache. Fail-open: any
      // bucket error must never affect the primary 0..1 saturation signal.
      try {
        updateAccountBuckets(connectionId, usage as ClaudeUsageResult, Date.now());
      } catch {
        // intentionally swallowed — buckets are additive, never gate-breaking
      }
      const util = planUtilizationFromUsage(usage, dim.window);
      if (util !== null) return util;
    }
  } catch (err) {
    log.warn(
      { err: (err as Error)?.message, connectionId },
      "anthropic oauth/usage saturation failed — falling back to rate-limit headers"
    );
  }

  // Fallback: per-minute REQUEST rate-limit headers (weak, TPM/RPM only).
  return anthropicHeaderSaturation(connectionId);
}

/**
 * Injectable seam for the generic usage fetch so the token-header
 * complement/fallback is unit-testable without the open-sse usage service.
 * Defaults to getUsageForProvider; pass null to restore.
 */
type GenericUsageFetcher = (connectionId: string, provider: string) => Promise<unknown>;
let _genericUsageFetcherOverride: GenericUsageFetcher | null = null;

/** Test-only: inject the generic usage fetcher; pass null to restore. */
export function __setGenericUsageFetcherForTests(fetcher: GenericUsageFetcher | null): void {
  _genericUsageFetcherOverride = fetcher;
}

async function defaultGenericUsageFetch(connectionId: string, provider: string): Promise<unknown> {
  const mod = await import("@omniroute/open-sse/services/usage");
  const conn = { id: connectionId, provider } as Parameters<typeof mod.getUsageForProvider>[0];
  return mod.getUsageForProvider(conn);
}

async function fetchGenericSaturation(connectionId: string, provider: string): Promise<number> {
  // 1. Real usage percent is authoritative when present (a provider that
  //    actually reports utilization beats the burst-window token headers).
  try {
    const fetcher = _genericUsageFetcherOverride ?? defaultGenericUsageFetch;
    const result = await fetcher(connectionId, provider);
    if (result && typeof result === "object") {
      const obj = result as Record<string, unknown>;

      // Prefer the normalized quota shape (handles nested `quotas` map for
      // Antigravity / Claude / etc.). Fall back to legacy top-level fields.
      const { convertUsageToQuotaInfo } =
        await import("@omniroute/open-sse/services/genericQuotaFetcher");
      const quota = convertUsageToQuotaInfo(result);
      if (quota && Number.isFinite(quota.percentUsed)) {
        return Math.min(1, Math.max(0, quota.percentUsed));
      }

      const pct =
        typeof obj.percentUsed === "number"
          ? obj.percentUsed
          : typeof obj.used_percent === "number"
            ? obj.used_percent
            : null;
      if (pct !== null && Number.isFinite(pct)) {
        return Math.min(1, Math.max(0, pct));
      }
    }
  } catch {
    // fall through to the token-header complement
  }

  // 2. Complement/fallback: proactive TOKEN-header saturation (universal, rides
  //    on every response). Fail-open to 0 when no fresh token-header data.
  return getTokenHeaderSaturation(provider, connectionId)?.saturation ?? 0;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/**
 * Return the current global saturation signal (0..1) for a connection+dim.
 *
 * A value of 0 means "no saturation detected" (generous/borrowing mode allowed).
 * A value >= saturationThreshold triggers strict mode in fairShare.ts.
 *
 * Always fail-open: returns 0 on any error.
 */
export async function getSaturation(
  connectionId: string,
  provider: string,
  dim: DimensionSpec,
  connection?: Record<string, unknown>
): Promise<number> {
  const key = cacheKey(connectionId, provider, dim);
  const cached = _cache.get(key);
  if (cached && Date.now() - cached.ts < CACHE_TTL_MS) {
    return cached.value;
  }

  let value = 0;
  try {
    switch (provider) {
      case "codex":
        value = await fetchCodexSaturation(connectionId, dim, connection);
        break;
      case "bailian":
        value = await fetchBailianSaturation(connectionId, dim);
        break;
      case "anthropic":
      case "claude":
        value = await fetchAnthropicSaturation(connectionId, dim);
        break;
      default:
        value = await fetchGenericSaturation(connectionId, provider);
        break;
    }
  } catch (err) {
    log.warn(
      { err: (err as Error)?.message, connectionId, provider },
      "saturation fetch failed — failing open with 0"
    );
    value = 0;
  }

  _cache.set(key, { value, ts: Date.now() });
  return value;
}