MobilityOps/frontend/src/context/AuthContext.tsx

import { createContext, useCallback, useContext, useEffect, useState, type ReactNode } from "react";
import { api, ApiError, onUnauthorized } from "../api/client";
import type { CurrentUser, Role, SystemStatus } from "../api/types";

interface AuthState {
  user: CurrentUser | null;
  loading: boolean;
  demoMode: boolean;
  oidcEnabled: boolean;
  oidcProviderName: string | null;
  loginAs: (role: Role) => Promise<void>;
  loginWithPassword: (email: string, password: string) => Promise<void>;
  logout: () => Promise<void>;
}

const AuthContext = createContext<AuthState | undefined>(undefined);

const STORAGE_KEY = "mobilityops.demo-user";

function readCachedUser(): CurrentUser | null {
  try {
    const stored = sessionStorage.getItem(STORAGE_KEY);
    return stored ? (JSON.parse(stored) as CurrentUser) : null;
  } catch {
    // A corrupted or blocked sessionStorage must never prevent the app from booting;
    // the session endpoint remains the source of truth.
    return null;
  }
}

function cacheUser(user: CurrentUser | null) {
  try {
    if (user) {
      sessionStorage.setItem(STORAGE_KEY, JSON.stringify(user));
    } else {
      sessionStorage.removeItem(STORAGE_KEY);
    }
  } catch {
    // Session storage is only a paint optimisation. Browser privacy settings or a full
    // storage quota must never turn a successful server-side login/logout into a client
    // failure; the HttpOnly session cookie remains authoritative.
  }
}

export function AuthProvider({ children }: { children: ReactNode }) {
  // The cached value only avoids a login-screen flash while the server round trip below is
  // in flight; it is never trusted on its own. `loading` stays true until the server has
  // confirmed (or rejected) the session cookie.
  const [user, setUser] = useState<CurrentUser | null>(() => readCachedUser());
  const [loading, setLoading] = useState(true);
  const [demoMode, setDemoMode] = useState(true);
  const [oidcEnabled, setOidcEnabled] = useState(false);
  const [oidcProviderName, setOidcProviderName] = useState<string | null>(null);

  useEffect(() => {
    const controller = new AbortController();

    // System capabilities and session identity are independent reads. In particular, a
    // slow public status endpoint must not hold an already valid session (or the login
    // controls) behind its timeout.
    void api.get<SystemStatus>("/api/v1/system/status", { signal: controller.signal })
      .then((system) => {
        if (controller.signal.aborted) return;
        setDemoMode(system.demo_mode);
        setOidcEnabled(system.oidc_enabled);
        setOidcProviderName(system.oidc_provider_name);
      })
      .catch(() => {
        if (!controller.signal.aborted) setDemoMode(true);
      });

    void api
      .get<CurrentUser>("/api/v1/auth/session", { signal: controller.signal })
      .then((confirmed) => {
        if (controller.signal.aborted) return;
        setUser(confirmed);
        cacheUser(confirmed);
      })
      .catch(() => {
        if (controller.signal.aborted) return;
        setUser(null);
        cacheUser(null);
      })
      .finally(() => {
        if (!controller.signal.aborted) setLoading(false);
      });

    return () => controller.abort();
  }, []);

  useEffect(
    () =>
      onUnauthorized(() => {
        setUser(null);
        cacheUser(null);
      }),
    [],
  );

  const loginAs = useCallback(async (role: Role) => {
    setLoading(true);
    try {
      const loggedIn = await api.post<CurrentUser>("/api/v1/demo/login", { role });
      setUser(loggedIn);
      cacheUser(loggedIn);
    } finally {
      setLoading(false);
    }
  }, []);

  const loginWithPassword = useCallback(async (email: string, password: string) => {
    setLoading(true);
    try {
      const loggedIn = await api.post<CurrentUser>("/api/v1/auth/login", { email, password });
      setUser(loggedIn);
      cacheUser(loggedIn);
    } finally {
      setLoading(false);
    }
  }, []);

  const logout = useCallback(async () => {
    try {
      // A logout response expires an HttpOnly cookie. Chromium can expose a very small
      // race between resolving fetch() and applying that Set-Cookie header to an
      // immediate top-level navigation. Confirm the server now rejects the session
      // before allowing the router to continue; retrying logout is idempotent.
      for (let attempt = 0; attempt < 3; attempt += 1) {
        await api.post("/api/v1/auth/logout");
        try {
          await api.get<CurrentUser>("/api/v1/auth/session");
        } catch (err) {
          if (err instanceof ApiError && err.status === 401) return;
          throw err;
        }
      }
    } catch {
      // Best effort: the cookie is cleared server-side when it works, and the client has
      // still drops its own state when the request itself is unavailable.
    } finally {
      // Keep RequireAuth from redirecting to /login while the server-side invalidation
      // is still in flight. Otherwise a very fast navigation can race ahead of logout.
      setUser(null);
      cacheUser(null);
    }
  }, []);

  return (
    <AuthContext.Provider value={{ user, loading, demoMode, oidcEnabled, oidcProviderName, loginAs, loginWithPassword, logout }}>{children}</AuthContext.Provider>
  );
}

export function useAuth(): AuthState {
  const ctx = useContext(AuthContext);
  if (!ctx) throw new Error("useAuth must be used within AuthProvider");
  return ctx;
}