type TokenLoader = () => Promise<string | null>;

/** Cache one Better Auth/Convex JWT per concrete session. */
export const createSessionTokenResolver = () => {
  let cachedSessionId: string | null = null;
  let cachedToken: string | null = null;
  let pending: {
    readonly sessionId: string;
    readonly promise: Promise<string | null>;
  } | null = null;

  return async (
    sessionId: string | null,
    loadToken: TokenLoader
  ): Promise<string | null> => {
    if (!sessionId) {
      cachedSessionId = null;
      cachedToken = null;
      pending = null;
      return null;
    }

    if (cachedSessionId !== sessionId) {
      cachedSessionId = sessionId;
      cachedToken = null;
      pending = null;
    }
    if (cachedToken) {
      return cachedToken;
    }
    if (pending?.sessionId === sessionId) {
      return pending.promise;
    }

    const requestedSessionId = sessionId;
    const promise = (async (): Promise<string | null> => {
      try {
        const token = await loadToken();
        if (cachedSessionId === requestedSessionId) {
          cachedToken = token;
        }
        return token;
      } catch {
        if (cachedSessionId === requestedSessionId) {
          cachedToken = null;
        }
        return null;
      }
    })();
    pending = { promise, sessionId: requestedSessionId };
    try {
      return await promise;
    } finally {
      if (pending?.promise === promise) {
        pending = null;
      }
    }
  };
};
