// ---------------------------------------------------------------------------
// JWKS cache for JWT verification (session + MCP auth).
// ---------------------------------------------------------------------------
//
// Cloudflare Workers boot many short-lived isolates. `createRemoteJWKSet`'s
// default cooldown (30s) and cache max-age (10m) still results in many JWKS
// fetches per hour because each new isolate starts cold. Production p99 for
// `mcp.auth.jwt_verify` was 1.7s — almost entirely the JWKS fetch.
//
// Module-scope memory alone only helps while an isolate stays warm. When
// `workos.session.local_verify` started missing ~92% of the time (2026-08-18),
// every miss paid a fresh upstream fetch — p50 3s, and the tail crossed the 5s
// request timeout, turning a slow key-server into `AbortError` → 500 on every
// authenticated API route and 503 on `/mcp`. A transient key-server blip must
// never read as an auth failure, and a cold isolate must not have to go
// upstream at all.
//
// So this module layers three caches, fastest first:
//
//   1. Module-scope memory — free, but dies with the isolate.
//   2. A cross-isolate store (the Workers Cache API by default) — colo-local,
//      survives isolate recycling, so a cold isolate reads keys in ~1ms
//      instead of paying an upstream round trip.
//   3. The upstream JWKS endpoint.
//
// On top of that it is stale-while-revalidate: once past `ttlMs` a usable key
// set is served immediately and refreshed in the background, and if a refresh
// fails we keep serving the last good keys until `staleMaxMs`. Only a fully
// cold path (no memory, no store) ever blocks on the network.
//
// Serving stale keys is safe in a way that serving a stale *token* would not
// be: key sets rotate on the order of days, tokens are still signature- and
// expiry-checked against them, and a token whose `kid` is absent forces a real
// refresh before it is rejected (see `get`). `staleMaxMs` bounds how long a
// retired key can still be honoured.
//
// The returned function is a `JWTVerifyGetKey` and slots directly into
// `jose.jwtVerify`. It also exposes `forceRefresh()` so the verify path can
// invalidate the cache and retry on a bad signature.
// ---------------------------------------------------------------------------

import {
  createLocalJWKSet,
  type FlattenedJWSInput,
  type JSONWebKeySet,
  type JWTHeaderParameters,
  type JWTVerifyGetKey,
  type KeyLike,
} from "jose";
import { Schema } from "effect";
import { JWKSNoMatchingKey } from "jose/errors";

/**
 * A cross-isolate key-set store. Defaults to the Workers Cache API; tests and
 * non-Workers hosts pass their own, or `null` to stay memory-only.
 */
export interface JwksStore {
  readonly get: (url: URL) => Promise<StoredJwks | null>;
  readonly put: (url: URL, stored: StoredJwks) => Promise<void>;
}

export interface StoredJwks {
  readonly jwks: JSONWebKeySet;
  /** Epoch ms of the upstream fetch these keys came from. */
  readonly fetchedAt: number;
}

export interface CachedRemoteJWKSetOptions {
  /**
   * How long a successful fetch is considered fresh. Defaults to 1 hour —
   * AuthKit rotates JWKS roughly daily, and a forced refresh on verify
   * failure handles unscheduled rotations.
   */
  readonly ttlMs?: number;
  /**
   * How long past `ttlMs` a key set may still be served while refreshes are
   * failing. Defaults to 24h — long enough to ride out a key-server outage,
   * short enough to bound how long a retired key stays honoured.
   */
  readonly staleMaxMs?: number;
  /** Override the fetch implementation for tests. */
  readonly fetch?: typeof globalThis.fetch;
  /** HTTP request timeout. Defaults to 5s, matching jose. */
  readonly timeoutMs?: number;
  /**
   * Cross-isolate store. Defaults to the Workers Cache API when available,
   * `null` disables the layer (memory-only).
   */
  readonly store?: JwksStore | null;
}

export interface CachedRemoteJWKSet extends JWTVerifyGetKey {
  /** Drop the cached JWKS so the next call refetches. */
  readonly forceRefresh: () => void;
  /**
   * Inspect the current cache state (testing/diagnostics/span annotation).
   * `fetchCount`/`fetchFailureCount` are lifetime counters for this resolver
   * instance; callers snapshot them around a verify to tell a cache hit from
   * a live upstream fetch (the Aug 2026 latency regression was exactly this
   * cache silently missing on ~92% of verifies, invisible in traces).
   *
   * `blockingFetchCount` counts only fetches a verify actually WAITED on.
   * Under stale-while-revalidate `fetchCount` also moves for background
   * revalidation the caller never paid for, so latency attribution wants
   * `blockingFetchCount`. `storeHitCount` counts cold-isolate reads answered
   * by the cross-isolate store instead of going upstream.
   */
  readonly inspect: () => {
    fetchedAt: number | null;
    hasJwks: boolean;
    fetchCount: number;
    fetchFailureCount: number;
    blockingFetchCount: number;
    storeHitCount: number;
    lastFetchDurationMs: number | null;
    /** Wall-clock of the last cross-isolate store read (I/O). */
    lastStoreReadMs: number | null;
    /** Wall-clock of the last key resolution — WebCrypto `importKey`. */
    lastResolveMs: number | null;
  };
}

const DEFAULT_TTL_MS = 60 * 60 * 1000;
const DEFAULT_STALE_MAX_MS = 24 * 60 * 60 * 1000;
const DEFAULT_TIMEOUT_MS = 5000;

const JsonWebKey = Schema.Record(Schema.String, Schema.Unknown);
const JsonWebKeySetPayload = Schema.Struct({
  keys: Schema.Array(JsonWebKey),
});
const decodeJsonWebKeySetPayload = Schema.decodeUnknownPromise(JsonWebKeySetPayload);

const ErrorWithCode = Schema.Struct({ code: Schema.String });
const isErrorWithCode = Schema.is(ErrorWithCode);

const isJwksNoMatchingKey = (cause: unknown): boolean =>
  isErrorWithCode(cause) && cause.code === JWKSNoMatchingKey.code;

interface CacheEntry {
  jwks: JSONWebKeySet;
  fetchedAt: number;
  resolver: (protectedHeader: JWTHeaderParameters, token?: FlattenedJWSInput) => Promise<KeyLike>;
}

const entryFrom = (stored: StoredJwks): CacheEntry => ({
  jwks: stored.jwks,
  fetchedAt: stored.fetchedAt,
  resolver: createLocalJWKSet(stored.jwks),
});

/**
 * Cache upkeep (store writes, background revalidation) is best effort: it must
 * never fail or delay the verify that happened to trigger it.
 */
const ignoreFailure = async (work: Promise<unknown>): Promise<void> => {
  // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: best-effort cache upkeep must not surface as a verify failure
  try {
    await work;
  } catch {
    // Deliberately swallowed — the caller already has usable keys, or will
    // take the upstream path on its next miss.
  }
};

const fetchJwksOnce = async (
  url: URL,
  fetchImpl: typeof globalThis.fetch,
  timeoutMs: number,
): Promise<JSONWebKeySet> => {
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), timeoutMs);
  // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: fetch adapter must clear abort timer while preserving promise rejection behavior
  try {
    const response = await fetchImpl(url.toString(), {
      method: "GET",
      headers: { accept: "application/json" },
      signal: controller.signal,
    });

    if (!response.ok) {
      // oxlint-disable-next-line executor/no-try-catch-or-throw, executor/no-error-constructor -- boundary: fetch-backed JWT key resolver must reject with the existing Error cause shape
      throw new Error(`JWKS fetch failed: ${response.status} ${response.statusText}`);
    }

    const body = await response.json();
    // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: fetch JSON validation maps Schema failures to the existing malformed JWKS rejection
    try {
      await decodeJsonWebKeySetPayload(body);
      return body as JSONWebKeySet;
    } catch {
      // oxlint-disable-next-line executor/no-try-catch-or-throw, executor/no-error-constructor -- boundary: fetch JSON validation preserves the existing malformed JWKS rejection
      throw new Error("JWKS fetch returned malformed payload");
    }
  } finally {
    clearTimeout(timer);
  }
};

// ---------------------------------------------------------------------------
// Workers Cache API store
// ---------------------------------------------------------------------------
//
// Keyed by the JWKS URL itself. We store our own Response rather than the
// upstream one so the body is already-validated JSON and the cache headers are
// ours: `max-age` covers the stale window, because an entry past `ttlMs` is
// still useful to us as a stale fallback.

const STORE_HEADER_FETCHED_AT = "x-jwks-fetched-at";

/** `caches.default` is a Workers extension the standard lib type omits. */
type WorkersCacheStorage = CacheStorage & { readonly default?: Cache };

const workersCacheStore = (staleMaxMs: number): JwksStore | null => {
  if (typeof caches === "undefined") return null;
  const open = (): Cache | null => (caches as WorkersCacheStorage).default ?? null;

  return {
    get: async (url) => {
      const cache = open();
      if (!cache) return null;
      const hit = await cache.match(url.toString());
      if (!hit) return null;
      const fetchedAtHeader = hit.headers.get(STORE_HEADER_FETCHED_AT);
      const fetchedAt = fetchedAtHeader === null ? Number.NaN : Number(fetchedAtHeader);
      if (!Number.isFinite(fetchedAt)) return null;
      const body = await hit.json();
      // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: a corrupt cache entry must degrade to a miss, never fail the verify
      try {
        await decodeJsonWebKeySetPayload(body);
      } catch {
        return null;
      }
      return { jwks: body as JSONWebKeySet, fetchedAt };
    },
    put: async (url, stored) => {
      const cache = open();
      if (!cache) return;
      const maxAgeSeconds = Math.max(1, Math.floor(staleMaxMs / 1000));
      await cache.put(
        url.toString(),
        new Response(JSON.stringify(stored.jwks), {
          status: 200,
          headers: {
            "content-type": "application/json",
            "cache-control": `max-age=${maxAgeSeconds}`,
            [STORE_HEADER_FETCHED_AT]: String(stored.fetchedAt),
          },
        }),
      );
    },
  };
};

/**
 * Creates a cached, single-flight, force-refreshable JWKS resolver compatible
 * with `jose.jwtVerify`. Drop-in replacement for `createRemoteJWKSet` for the
 * auth paths — see module header for why we don't just use jose's built-in.
 */
export const createCachedRemoteJWKSet = (
  url: URL,
  options: CachedRemoteJWKSetOptions = {},
): CachedRemoteJWKSet => {
  const ttlMs = options.ttlMs ?? DEFAULT_TTL_MS;
  const staleMaxMs = Math.max(options.staleMaxMs ?? DEFAULT_STALE_MAX_MS, ttlMs);
  const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
  const store = options.store === undefined ? workersCacheStore(staleMaxMs) : options.store;
  // Capture the fetch impl lazily so consumers can swap globalThis.fetch
  // (tests do this) without us snapshotting a stale reference.
  const fetchImpl = (): typeof globalThis.fetch =>
    options.fetch ?? globalThis.fetch.bind(globalThis);

  let entry: CacheEntry | null = null;
  let inflight: Promise<CacheEntry> | null = null;
  let fetchCount = 0;
  let fetchFailureCount = 0;
  let blockingFetchCount = 0;
  let storeHitCount = 0;
  let lastFetchDurationMs: number | null = null;
  let lastStoreReadMs: number | null = null;
  let lastResolveMs: number | null = null;

  const isFresh = (candidate: CacheEntry): boolean => Date.now() - candidate.fetchedAt < ttlMs;
  const isUsable = (candidate: CacheEntry): boolean =>
    Date.now() - candidate.fetchedAt < staleMaxMs;

  const refresh = (): Promise<CacheEntry> => {
    if (inflight) return inflight;
    const startedAt = Date.now();
    fetchCount += 1;
    inflight = (async () => {
      const jwks = await fetchJwksOnce(url, fetchImpl(), timeoutMs);
      const next = entryFrom({ jwks, fetchedAt: Date.now() });
      entry = next;
      if (store) {
        await ignoreFailure(store.put(url, { jwks: next.jwks, fetchedAt: next.fetchedAt }));
      }
      return next;
    })()
      .then(
        (next) => next,
        (error: unknown) => {
          fetchFailureCount += 1;
          // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: counting a fetch failure must preserve the original rejection for jose
          throw error;
        },
      )
      .finally(() => {
        lastFetchDurationMs = Date.now() - startedAt;
        inflight = null;
      });
    return inflight;
  };

  /** Fire-and-forget revalidation behind a stale hit. */
  const refreshInBackground = (): void => {
    void ignoreFailure(refresh());
  };

  const loadFromStore = async (): Promise<CacheEntry | null> => {
    if (!store) return null;
    // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: the L2 store is an optimization; any failure degrades to an upstream fetch
    try {
      const storeReadStartedAt = Date.now();
      const stored = await store.get(url);
      lastStoreReadMs = Date.now() - storeReadStartedAt;
      if (!stored) return null;
      const candidate = entryFrom(stored);
      if (!isUsable(candidate)) return null;
      entry = candidate;
      storeHitCount += 1;
      return candidate;
    } catch {
      return null;
    }
  };

  /** A refresh the caller waits on — the only kind that costs it latency. */
  const refreshBlocking = (): Promise<CacheEntry> => {
    blockingFetchCount += 1;
    return refresh();
  };

  const ensureFresh = async (forceRefresh: boolean): Promise<CacheEntry> => {
    if (forceRefresh) return refreshBlocking();
    if (entry && isFresh(entry)) return entry;

    // Memory is stale but still usable: serve it now, revalidate behind it.
    if (entry && isUsable(entry)) {
      refreshInBackground();
      return entry;
    }

    // Cold isolate — the L2 store saves us the upstream round trip.
    const stored = await loadFromStore();
    if (stored) {
      if (!isFresh(stored)) refreshInBackground();
      return stored;
    }

    // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: a failed refresh must fall back to stale keys rather than fail the verify
    try {
      return await refreshBlocking();
    } catch (error) {
      // Upstream is slow or down. Last good keys beat failing every request.
      if (entry && isUsable(entry)) return entry;
      // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: nothing usable is cached, so the upstream failure is the real answer
      throw error;
    }
  };

  const get: JWTVerifyGetKey = async (protectedHeader, token) => {
    const current = await ensureFresh(false);
    // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: jose JWTVerifyGetKey retry path is defined by thrown resolver failures
    try {
      const resolveStartedAt = Date.now();
      const key = await current.resolver(protectedHeader, token);
      lastResolveMs = Date.now() - resolveStartedAt;
      return key;
    } catch (error) {
      // Likely cause: keys rotated upstream after our TTL window started, or
      // we answered from a stale entry. Refetch once and try again. Anything
      // still failing bubbles up so jose can classify it (we do not silently
      // swallow real failures).
      if (!isJwksNoMatchingKey(error)) {
        // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: jose JWTVerifyGetKey requires preserving upstream resolver rejection
        throw error;
      }
      const refreshed = await ensureFresh(true);
      return refreshed.resolver(protectedHeader, token);
    }
  };

  const result = get as CachedRemoteJWKSet;
  Object.defineProperty(result, "forceRefresh", {
    value: () => {
      entry = null;
    },
  });
  Object.defineProperty(result, "inspect", {
    value: () => ({
      fetchedAt: entry?.fetchedAt ?? null,
      hasJwks: entry !== null,
      fetchCount,
      fetchFailureCount,
      blockingFetchCount,
      storeHitCount,
      lastFetchDurationMs,
      lastStoreReadMs,
      lastResolveMs,
    }),
  });
  return result;
};