Jyotisha/frontend/src/lib/request-rate-limit.ts

export type RequestRateLimitStore = Map<string, number[]>;

export type RequestRateLimitDecision =
  | { readonly ok: true }
  | { readonly ok: false; readonly retryAfterSeconds: number };

export type RequestRateLimitInput = Readonly<{
  key: string;
  limit: number;
  windowMs: number;
  now?: number;
  store?: RequestRateLimitStore;
}>;

const state = globalThis as typeof globalThis & {
  jyotishaRequestRateLimits?: RequestRateLimitStore;
};

export const requestRateLimits = {
  locationSearch: { limit: 20, windowMs: 60_000 },
  locationTimezone: { limit: 30, windowMs: 60_000 },
  dailyStarlanguage: { limit: 10, windowMs: 60_000 },
  ephemeris: { limit: 30, windowMs: 60_000 },
  synastry: { limit: 8, windowMs: 60 * 60_000 },
  sessionWrite: { limit: 40, windowMs: 60_000 },
  // Two windows each: a minute burst and an hour budget (product 2026-10-14).
  consultMinute: { limit: 10, windowMs: 60_000 },
  consultHour: { limit: 120, windowMs: 60 * 60_000 },
  rectificationTurnMinute: { limit: 5, windowMs: 60_000 },
  rectificationTurnHour: { limit: 40, windowMs: 60 * 60_000 },
  reportGenerateHour: { limit: 3, windowMs: 60 * 60_000 },
} as const;

export const userRateLimitMessages = {
  consult: "这一会儿问得有点密，先停一下。",
  rectificationTurn: "这一会儿校正问得有点密，先停一下。",
  reportGenerate: "这一小时的报告已经够多了，先停一下。",
} as const;

export type RequestRateLimitBucket = keyof typeof requestRateLimits;

function defaultStore(): RequestRateLimitStore {
  state.jyotishaRequestRateLimits ??= new Map();
  return state.jyotishaRequestRateLimits;
}

let seenMaxWindowMs = 0;
let consumeCalls = 0;

export function resetRequestRateLimitMaintenanceForTests(): void {
  seenMaxWindowMs = 0;
  consumeCalls = 0;
}

function sweepExpiredRateLimitKeys(
  store: RequestRateLimitStore,
  now: number,
  maxWindowMs: number,
): void {
  for (const [key, stamps] of store) {
    let newest = Number.NEGATIVE_INFINITY;
    for (const stamp of stamps) {
      if (stamp > newest) newest = stamp;
    }
    if (!Number.isFinite(newest) || newest <= now - maxWindowMs) store.delete(key);
  }
}

function evaluateRequestRateLimit(
  input: RequestRateLimitInput,
  record: boolean,
): RequestRateLimitDecision {
  const now = input.now ?? Date.now();
  const store = input.store ?? defaultStore();
  if (record) {
    if (input.windowMs > seenMaxWindowMs) seenMaxWindowMs = input.windowMs;
    consumeCalls += 1;
    if (consumeCalls % 1_000 === 0 || store.size > 5_000) {
      sweepExpiredRateLimitKeys(store, now, seenMaxWindowMs);
    }
  }
  const windowStart = now - input.windowMs;
  const recent = (store.get(input.key) ?? []).filter((timestamp) => timestamp > windowStart);
  if (recent.length >= input.limit) {
    const retryAfterSeconds = Math.max(1, Math.ceil((recent[0]! + input.windowMs - now) / 1000));
    if (record) store.set(input.key, recent);
    return { ok: false, retryAfterSeconds };
  }
  if (!record) return { ok: true };
  recent.push(now);
  store.set(input.key, recent);
  return { ok: true };
}

export function consumeRequestRateLimit(input: RequestRateLimitInput): RequestRateLimitDecision {
  return evaluateRequestRateLimit(input, true);
}

export function consumeUserRequestRateLimit(
  bucket: RequestRateLimitBucket,
  userId: string,
  now?: number,
  store?: RequestRateLimitStore,
): RequestRateLimitDecision {
  const spec = requestRateLimits[bucket];
  return consumeRequestRateLimit({
    key: `${bucket}:${userId}`,
    limit: spec.limit,
    windowMs: spec.windowMs,
    now,
    store,
  });
}

/**
 * Allow the request only when every bucket still has room. A refusal writes
 * nothing, so the bucket that still had room is not charged for the miss.
 */
export function consumeUserRateLimitSet(
  buckets: readonly RequestRateLimitBucket[],
  userId: string,
  now?: number,
  store?: RequestRateLimitStore,
): RequestRateLimitDecision {
  let retryAfterSeconds = 0;
  for (const bucket of buckets) {
    const spec = requestRateLimits[bucket];
    const decision = evaluateRequestRateLimit({
      key: `${bucket}:${userId}`,
      limit: spec.limit,
      windowMs: spec.windowMs,
      now,
      store,
    }, false);
    if (!decision.ok) retryAfterSeconds = Math.max(retryAfterSeconds, decision.retryAfterSeconds);
  }
  if (retryAfterSeconds > 0) return { ok: false, retryAfterSeconds };
  for (const bucket of buckets) consumeUserRequestRateLimit(bucket, userId, now, store);
  return { ok: true };
}

export function userRateLimitDeniedResponse(retryAfterSeconds: number, message: string) {
  return {
    status: 429 as const,
    headers: { "Retry-After": String(retryAfterSeconds) },
    body: {
      error: message,
      code: "rate_limited" as const,
      retryAfterSeconds,
    },
  };
}