chrysopedia/frontend/src/context/AuthContext.tsx

import {
  createContext,
  useContext,
  useState,
  useEffect,
  useCallback,
  type ReactNode,
} from "react";
import {
  AUTH_TOKEN_KEY,
  authLogin,
  authGetMe,
  authRegister,
  impersonateUser,
  stopImpersonation as apiStopImpersonation,
  ApiError,
  type UserResponse,
  type RegisterRequest,
} from "../api";

const ADMIN_TOKEN_KEY = "chrysopedia_admin_token";

interface AuthContextValue {
  user: UserResponse | null;
  token: string | null;
  isAuthenticated: boolean;
  isImpersonating: boolean;
  loading: boolean;
  login: (email: string, password: string) => Promise<void>;
  register: (data: RegisterRequest) => Promise<UserResponse>;
  logout: () => void;
  startImpersonation: (userId: string) => Promise<void>;
  exitImpersonation: () => Promise<void>;
}

const AuthContext = createContext<AuthContextValue | null>(null);

export function AuthProvider({ children }: { children: ReactNode }) {
  const [user, setUser] = useState<UserResponse | null>(null);
  const [token, setToken] = useState<string | null>(() => {
    try {
      return localStorage.getItem(AUTH_TOKEN_KEY);
    } catch {
      return null;
    }
  });
  const [loading, setLoading] = useState(!!token);

  // Rehydrate session from stored token on mount
  useEffect(() => {
    if (!token) return;
    let cancelled = false;
    authGetMe(token)
      .then((u) => {
        if (!cancelled) setUser(u);
      })
      .catch(() => {
        // Token expired or invalid — clear it
        if (!cancelled) {
          localStorage.removeItem(AUTH_TOKEN_KEY);
          setToken(null);
          setUser(null);
        }
      })
      .finally(() => {
        if (!cancelled) setLoading(false);
      });
    return () => {
      cancelled = true;
    };
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  const login = useCallback(async (email: string, password: string) => {
    const resp = await authLogin(email, password);
    localStorage.setItem(AUTH_TOKEN_KEY, resp.access_token);
    setToken(resp.access_token);
    const me = await authGetMe(resp.access_token);
    setUser(me);
  }, []);

  const register = useCallback(async (data: RegisterRequest) => {
    return authRegister(data);
  }, []);

  const logout = useCallback(() => {
    localStorage.removeItem(AUTH_TOKEN_KEY);
    sessionStorage.removeItem(ADMIN_TOKEN_KEY);
    setToken(null);
    setUser(null);
  }, []);

  const startImpersonation = useCallback(async (userId: string) => {
    if (!token) return;
    // Save admin token so we can restore it later
    sessionStorage.setItem(ADMIN_TOKEN_KEY, token);
    const resp = await impersonateUser(token, userId);
    localStorage.setItem(AUTH_TOKEN_KEY, resp.access_token);
    setToken(resp.access_token);
    const me = await authGetMe(resp.access_token);
    setUser(me);
  }, [token]);

  const exitImpersonation = useCallback(async () => {
    // Try to call stop endpoint for audit log
    if (token) {
      try {
        await apiStopImpersonation(token);
      } catch {
        // Best effort — still restore admin session
      }
    }
    // Restore admin token
    const adminToken = sessionStorage.getItem(ADMIN_TOKEN_KEY);
    sessionStorage.removeItem(ADMIN_TOKEN_KEY);
    if (adminToken) {
      localStorage.setItem(AUTH_TOKEN_KEY, adminToken);
      setToken(adminToken);
      const me = await authGetMe(adminToken);
      setUser(me);
    } else {
      // Fallback: just logout
      logout();
    }
  }, [token, logout]);

  return (
    <AuthContext.Provider
      value={{
        user,
        token,
        isAuthenticated: !!user,
        isImpersonating: !!user?.impersonating,
        loading,
        login,
        register,
        logout,
        startImpersonation,
        exitImpersonation,
      }}
    >
      {children}
    </AuthContext.Provider>
  );
}

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

export { ApiError };