waseel/lib/session.tsx

import * as SecureStore from "expo-secure-store";
import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useState,
  type ReactNode,
} from "react";

import { setAuthToken, clearAuthToken, setUnauthorizedHandler } from "./fetch";
import { stopBackgroundTracking } from "./location-task";
import { releaseCurrentPush } from "./notifications";
import { TOKEN_KEY } from "./token-store";

const USER_KEY = "waseel_auth_user";
const REMEMBERED_EMAIL_KEY = "waseel_remembered_email";
const DEFAULT_TOKEN_TTL_SECONDS = 30 * 24 * 60 * 60;

/**
 * The address to prefill on the sign-in screen, or null when the last sign-in
 * cleared "remember me". The token itself is stored either way — what
 * "remember me" changes is how long the server makes it live.
 */
export const getRememberedEmail = (): Promise<string | null> =>
  SecureStore.getItemAsync(REMEMBERED_EMAIL_KEY);

export const rememberEmail = async (email: string | null): Promise<void> => {
  if (email) {
    await SecureStore.setItemAsync(REMEMBERED_EMAIL_KEY, email);
  } else {
    await SecureStore.deleteItemAsync(REMEMBERED_EMAIL_KEY);
  }
};

export type SessionUser = {
  id: string;
  name: string;
  email: string;
  role?: string | null;
  avatarUrl?: string | null;
};

type AuthResult = { token: string; user: SessionUser };

type SessionContextValue = {
  isLoaded: boolean;
  isSignedIn: boolean;
  userId: string | null;
  user: SessionUser | null;
  setSession: (result: AuthResult) => Promise<void>;
  setUserRole: (role: string) => void;
  signOut: () => Promise<void>;
};

const SessionContext = createContext<SessionContextValue | null>(null);

const BASE64_ALPHABET =
  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

/**
 * Decodes a base64url segment without `atob`.
 *
 * Neither React Native 0.74 nor Expo SDK 51 defines a global `atob`, so the
 * previous implementation threw on every call — which silently expired the
 * stored session on each launch (see decodeJwtExp). Skipping unknown
 * characters also makes the missing "=" padding in a JWT a non-issue.
 */
const decodeBase64Url = (segment: string): string => {
  const normalized = segment.replace(/-/g, "+").replace(/_/g, "/");

  let output = "";
  let buffer = 0;
  let bits = 0;

  for (const character of normalized) {
    const value = BASE64_ALPHABET.indexOf(character);

    if (value === -1) continue;

    buffer = (buffer << 6) | value;
    bits += 6;

    if (bits >= 8) {
      bits -= 8;
      output += String.fromCharCode((buffer >> bits) & 0xff);
    }
  }

  return output;
};

// Mirrors lib/jwt.ts payload decoding (no signature check needed client-side).
export const decodeJwtExp = (token: string): number | null => {
  try {
    const claims = JSON.parse(decodeBase64Url(token.split(".")[1]));

    return typeof claims.exp === "number" ? claims.exp : null;
  } catch {
    return null;
  }
};

export const SessionProvider = ({ children }: { children: ReactNode }) => {
  const [isLoaded, setIsLoaded] = useState(false);
  const [user, setUser] = useState<SessionUser | null>(null);

  useEffect(() => {
    let cancelled = false;

    const restore = async () => {
      try {
        const [token, storedUser] = await Promise.all([
          SecureStore.getItemAsync(TOKEN_KEY),
          SecureStore.getItemAsync(USER_KEY),
        ]);

        if (!token || !storedUser) return;

        const exp = decodeJwtExp(token);

        // Only discard a session we can positively prove is expired. Treating
        // an unreadable token as expired is what made every launch sign the
        // user back out; if it really is bad, the next request gets a 401.
        if (exp === null) {
          console.warn("[SESSION_RESTORE]: could not read token expiry");
        } else if (exp * 1000 < Date.now()) {
          await SecureStore.deleteItemAsync(TOKEN_KEY);
          await SecureStore.deleteItemAsync(USER_KEY);
          return;
        }

        if (cancelled) return;

        setAuthToken(token);
        setUser(JSON.parse(storedUser) as SessionUser);
      } catch (error) {
        console.error("[SESSION_RESTORE]: ", error);
      } finally {
        if (!cancelled) setIsLoaded(true);
      }
    };

    void restore();

    return () => {
      cancelled = true;
    };
  }, []);

  const setSession = useCallback(async (result: AuthResult) => {
    setAuthToken(result.token);
    setUser(result.user);

    await SecureStore.setItemAsync(TOKEN_KEY, result.token);
    await SecureStore.setItemAsync(USER_KEY, JSON.stringify(result.user));
  }, []);

  const setUserRole = useCallback(
    (role: string) => {
      setUser((currentUser) => {
        if (!currentUser) return currentUser;

        const updated = { ...currentUser, role };

        void SecureStore.setItemAsync(USER_KEY, JSON.stringify(updated));

        return updated;
      });
    },
    [],
  );

  const signOut = useCallback(async () => {
    // Hand the device back before the token goes away — the release call needs
    // to authenticate as the account that currently holds it. Phones get
    // shared, and without this the previous account keeps receiving ride
    // offers on a phone someone else is now signed in on.
    await releaseCurrentPush();

    // A driver who signs out is off shift: tear down the location foreground
    // service too, or they're left with a "you're online" notification and a
    // GPS drain for a session that has ended.
    await stopBackgroundTracking();

    clearAuthToken();
    setUser(null);

    await SecureStore.deleteItemAsync(TOKEN_KEY);
    await SecureStore.deleteItemAsync(USER_KEY);
  }, []);

  // The other half of the bargain struck in restore(): a token we cannot prove
  // is expired stays in use until the server rejects it, and this is what
  // happens when it does. Without it the app kept the dead token and every
  // screen behind the session had to invent its own meaning for the resulting
  // 401 — driver-home read it as "this user has no driver profile" and showed
  // an onboarding form to a driver who had finished onboarding weeks ago.
  useEffect(() => {
    setUnauthorizedHandler(() => void signOut());

    return () => setUnauthorizedHandler(null);
  }, [signOut]);

  const value = useMemo<SessionContextValue>(
    () => ({
      isLoaded,
      isSignedIn: user !== null,
      userId: user?.id ?? null,
      user,
      setSession,
      setUserRole,
      signOut,
    }),
    [isLoaded, user, setSession, setUserRole, signOut],
  );

  return (
    <SessionContext.Provider value={value}>
      {children}
    </SessionContext.Provider>
  );
};

export const useSession = (): SessionContextValue => {
  const context = useContext(SessionContext);

  if (!context) {
    throw new Error("useSession must be used within a SessionProvider.");
  }

  return context;
};

// Alias kept for parity with the previous auth API.
export const useAuth = useSession;

export const TOKEN_TTL_SECONDS = DEFAULT_TOKEN_TTL_SECONDS;