learning_ai_common_plat/packages/react-auth/src/auth-context.tsx

'use client';

import {
  createContext,
  useContext,
  useState,
  useCallback,
  useEffect,
  useRef,
  type ReactNode,
} from 'react';
import { createApiClient } from '@bytelyst/api-client';
import type { AuthConfig, AuthContextValue, AuthProviderInfo, BaseUser } from './types.js';

/**
 * Create a typed auth provider + hook for a specific user type.
 *
 * Supports the full auth lifecycle: login, register, forgot password,
 * change password, delete account, and automatic token refresh.
 *
 * @example
 * ```tsx
 * const { AuthProvider, useAuth } = createAuthProvider<AdminUser>({
 *   storagePrefix: "admin",
 *   loginEndpoint: "/auth/login",
 *   registerEndpoint: "/auth/register",
 *   forgotPasswordEndpoint: "/auth/forgot-password",
 *   changePasswordEndpoint: "/auth/change-password",
 *   deleteAccountEndpoint: "/auth/delete-account",
 *   refreshEndpoint: "/auth/refresh",
 *   mapLoginResponse: (data) => ({
 *     user: data.user,
 *     accessToken: data.accessToken,
 *     refreshToken: data.refreshToken,
 *   }),
 * });
 * ```
 */
export function createAuthProvider<TUser extends BaseUser = BaseUser>(config: AuthConfig<TUser>) {
  const {
    baseUrl: configBaseUrl = '/api',
    storagePrefix,
    loginEndpoint,
    registerEndpoint,
    forgotPasswordEndpoint,
    changePasswordEndpoint,
    deleteAccountEndpoint,
    refreshEndpoint,
    refreshIntervalMs = 45 * 60 * 1000,
    mapLoginResponse,
    onLoginFallback,
    onInit,
    onLogout,
    oauthEndpoint = '/auth/oauth',
    providersEndpoint = '/auth/providers',
    linkProviderEndpoint = '/auth/providers/link',
    mfaVerifyEndpoint = '/auth/mfa/verify',
    onMfaRequired,
    productId: configProductId,
  } = config;

  const USER_KEY = `${storagePrefix}_auth_user`;
  const TOKEN_KEY = `${storagePrefix}_access_token`;
  const REFRESH_KEY = `${storagePrefix}_refresh_token`;

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

  function getStoredUser(): TUser | null {
    if (typeof window === 'undefined') return null;
    try {
      const stored = localStorage.getItem(USER_KEY);
      return stored ? JSON.parse(stored) : null;
    } catch {
      return null;
    }
  }

  function saveSession(user: TUser, accessToken: string, refreshToken: string) {
    localStorage.setItem(USER_KEY, JSON.stringify(user));
    localStorage.setItem(TOKEN_KEY, accessToken);
    localStorage.setItem(REFRESH_KEY, refreshToken);
  }

  function clearSession() {
    localStorage.removeItem(USER_KEY);
    localStorage.removeItem(TOKEN_KEY);
    localStorage.removeItem(REFRESH_KEY);
  }

  function AuthProvider({ children }: { children: ReactNode }) {
    const [user, setUser] = useState<TUser | null>(() => {
      // Allow onInit to provide an initial session (e.g. from SSO cookies)
      if (onInit) {
        const initResult = onInit();
        if (initResult) {
          saveSession(initResult.user, initResult.accessToken, initResult.refreshToken);
          return initResult.user;
        }
      }
      return getStoredUser();
    });
    const [isLoading, setIsLoading] = useState(false);
    const [error, setError] = useState<string | null>(null);
    const [success, setSuccess] = useState<string | null>(null);
    const [providers, setProviders] = useState<AuthProviderInfo[]>([]);
    const [mfaRequired, setMfaRequired] = useState(false);
    const [mfaMethods, setMfaMethods] = useState<string[]>([]);
    const [mfaChallenge, setMfaChallenge] = useState<string | null>(null);
    const refreshTimerRef = useRef<ReturnType<typeof setInterval> | null>(null);

    const api = createApiClient({
      baseUrl: configBaseUrl,
      getToken: () => (typeof window !== 'undefined' ? localStorage.getItem(TOKEN_KEY) : null),
    });

    const clearMessages = useCallback(() => {
      setError(null);
      setSuccess(null);
    }, []);

    // ── Token refresh ──────────────────────────────

    const refreshAccessToken = useCallback(async () => {
      if (!refreshEndpoint) return;
      const rt = typeof window !== 'undefined' ? localStorage.getItem(REFRESH_KEY) : null;
      if (!rt) return;

      try {
        const data = await api.fetch<{ accessToken: string; refreshToken: string }>(
          refreshEndpoint,
          { method: 'POST', body: JSON.stringify({ refreshToken: rt }) }
        );
        localStorage.setItem(TOKEN_KEY, data.accessToken);
        localStorage.setItem(REFRESH_KEY, data.refreshToken);
      } catch {
        // Token expired — force logout
        setUser(null);
        clearSession();
        onLogout?.();
      }
    }, [api]);

    useEffect(() => {
      if (!user || !refreshEndpoint) return;
      refreshTimerRef.current = setInterval(refreshAccessToken, refreshIntervalMs);
      return () => {
        if (refreshTimerRef.current) clearInterval(refreshTimerRef.current);
      };
    }, [user, refreshAccessToken, refreshIntervalMs]);

    // ── MFA challenge helper ─────────────────────────

    function handleMfaChallenge(data: Record<string, unknown>): boolean {
      if (data && typeof data === 'object' && 'mfaRequired' in data && data.mfaRequired === true) {
        const challenge = data.mfaChallenge as string;
        const methods = data.methods as string[];
        setMfaRequired(true);
        setMfaChallenge(challenge);
        setMfaMethods(methods ?? []);
        onMfaRequired?.(challenge, methods ?? []);
        return true;
      }
      return false;
    }

    // ── Login ──────────────────────────────────────

    const login = useCallback(
      async (email: string, password: string) => {
        setIsLoading(true);
        setError(null);
        setMfaRequired(false);
        setMfaChallenge(null);
        setMfaMethods([]);
        try {
          const { data, error: fetchError } = await api.safeFetch<unknown>(loginEndpoint, {
            method: 'POST',
            body: JSON.stringify(
              configProductId
                ? { email, password, productId: configProductId }
                : { email, password }
            ),
          });

          if (data && !fetchError) {
            if (handleMfaChallenge(data as Record<string, unknown>)) {
              return false;
            }
            const mapped = mapLoginResponse(data);
            setUser(mapped.user);
            saveSession(mapped.user, mapped.accessToken, mapped.refreshToken);
            return true;
          }

          if (fetchError && onLoginFallback) {
            const fallback = await onLoginFallback(email, password, fetchError);
            if (fallback) {
              setUser(fallback.user);
              saveSession(fallback.user, fallback.accessToken, fallback.refreshToken);
              return true;
            }
          }

          setError(fetchError || 'Login failed');
          return false;
        } finally {
          setIsLoading(false);
        }
      },
      [api]
    );

    // ── Register ───────────────────────────────────

    const register = useCallback(
      async (email: string, password: string, displayName: string) => {
        if (!registerEndpoint) {
          setError('Registration not supported');
          return false;
        }
        setIsLoading(true);
        setError(null);
        try {
          const { data, error: fetchError } = await api.safeFetch<unknown>(registerEndpoint, {
            method: 'POST',
            body: JSON.stringify(
              configProductId
                ? { email, password, displayName, productId: configProductId }
                : { email, password, displayName }
            ),
          });

          if (data && !fetchError) {
            const mapped = mapLoginResponse(data);
            setUser(mapped.user);
            saveSession(mapped.user, mapped.accessToken, mapped.refreshToken);
            return true;
          }

          setError(fetchError || 'Registration failed');
          return false;
        } finally {
          setIsLoading(false);
        }
      },
      [api]
    );

    // ── Social login (Phase 1C) ────────────────────

    const loginWithOAuth = useCallback(
      async (provider: string, idToken: string) => {
        setIsLoading(true);
        setError(null);
        setMfaRequired(false);
        setMfaChallenge(null);
        setMfaMethods([]);
        try {
          const oauthBody: Record<string, string> = { idToken };
          if (configProductId) oauthBody.productId = configProductId;
          const { data, error: fetchError } = await api.safeFetch<unknown>(
            `${oauthEndpoint}/${provider}`,
            { method: 'POST', body: JSON.stringify(oauthBody) }
          );

          if (data && !fetchError) {
            if (handleMfaChallenge(data as Record<string, unknown>)) {
              return false;
            }
            const mapped = mapLoginResponse(data);
            setUser(mapped.user);
            saveSession(mapped.user, mapped.accessToken, mapped.refreshToken);
            return true;
          }

          setError(fetchError || `${provider} login failed`);
          return false;
        } finally {
          setIsLoading(false);
        }
      },
      [api]
    );

    const loginWithGoogle = useCallback(
      (idToken: string) => loginWithOAuth('google', idToken),
      [loginWithOAuth]
    );

    const loginWithMicrosoft = useCallback(
      (idToken: string) => loginWithOAuth('microsoft', idToken),
      [loginWithOAuth]
    );

    const loginWithApple = useCallback(
      (idToken: string) => loginWithOAuth('apple', idToken),
      [loginWithOAuth]
    );

    // ── Provider management (Phase 1C) ────────────

    const refreshProviders = useCallback(async () => {
      try {
        const data = await api.fetch<{ providers: AuthProviderInfo[] }>(providersEndpoint, {
          method: 'GET',
        });
        setProviders(data.providers ?? []);
      } catch {
        // non-fatal — providers list is supplementary
      }
    }, [api]);

    const linkProvider = useCallback(
      async (provider: string, idToken: string) => {
        setIsLoading(true);
        setError(null);
        try {
          const { error: fetchError } = await api.safeFetch<void>(linkProviderEndpoint, {
            method: 'POST',
            body: JSON.stringify({ provider, idToken }),
          });
          if (fetchError) {
            setError(fetchError);
            return false;
          }
          await refreshProviders();
          return true;
        } finally {
          setIsLoading(false);
        }
      },
      [api, refreshProviders]
    );

    const unlinkProvider = useCallback(
      async (provider: string) => {
        setIsLoading(true);
        setError(null);
        try {
          const { error: fetchError } = await api.safeFetch<void>(
            `${providersEndpoint}/${provider}`,
            { method: 'DELETE' }
          );
          if (fetchError) {
            setError(fetchError);
            return false;
          }
          await refreshProviders();
          return true;
        } finally {
          setIsLoading(false);
        }
      },
      [api, refreshProviders]
    );

    // ── MFA verify (Phase 2D) ─────────────────────

    const verifyMfa = useCallback(
      async (code: string, method: 'totp' | 'recovery') => {
        if (!mfaChallenge) {
          setError('No MFA challenge in progress');
          return false;
        }
        setIsLoading(true);
        setError(null);
        try {
          const { data, error: fetchError } = await api.safeFetch<unknown>(mfaVerifyEndpoint, {
            method: 'POST',
            body: JSON.stringify({ challengeToken: mfaChallenge, code, method }),
          });

          if (data && !fetchError) {
            const mapped = mapLoginResponse(data);
            setUser(mapped.user);
            saveSession(mapped.user, mapped.accessToken, mapped.refreshToken);
            setMfaRequired(false);
            setMfaChallenge(null);
            setMfaMethods([]);
            return true;
          }

          setError(fetchError || 'MFA verification failed');
          return false;
        } finally {
          setIsLoading(false);
        }
      },
      [api, mfaChallenge]
    );

    // ── Logout ─────────────────────────────────────

    const logout = useCallback(() => {
      setUser(null);
      clearSession();
      setProviders([]);
      setMfaRequired(false);
      setMfaChallenge(null);
      setMfaMethods([]);
      if (refreshTimerRef.current) clearInterval(refreshTimerRef.current);
      onLogout?.();
    }, []);

    // ── Forgot password ────────────────────────────

    const forgotPassword = useCallback(
      async (email: string) => {
        if (!forgotPasswordEndpoint) {
          setError('Forgot password not supported');
          return false;
        }
        setIsLoading(true);
        setError(null);
        setSuccess(null);
        try {
          const { error: fetchError } = await api.safeFetch<{ message: string }>(
            forgotPasswordEndpoint,
            { method: 'POST', body: JSON.stringify({ email }) }
          );
          if (fetchError) {
            setError(fetchError);
            return false;
          }
          setSuccess('If that email exists, a reset link has been sent.');
          return true;
        } finally {
          setIsLoading(false);
        }
      },
      [api]
    );

    // ── Change password ────────────────────────────

    const changePassword = useCallback(
      async (currentPassword: string, newPassword: string) => {
        if (!changePasswordEndpoint) {
          setError('Change password not supported');
          return false;
        }
        setIsLoading(true);
        setError(null);
        setSuccess(null);
        try {
          const { error: fetchError } = await api.safeFetch<{ message: string }>(
            changePasswordEndpoint,
            { method: 'POST', body: JSON.stringify({ currentPassword, newPassword }) }
          );
          if (fetchError) {
            setError(fetchError);
            return false;
          }
          setSuccess('Password changed successfully.');
          return true;
        } finally {
          setIsLoading(false);
        }
      },
      [api]
    );

    // ── Update user (local state + localStorage) ──

    const updateUser = useCallback((updates: Partial<TUser>) => {
      setUser(prev => {
        if (!prev) return null;
        const updated = { ...prev, ...updates };
        localStorage.setItem(USER_KEY, JSON.stringify(updated));
        return updated;
      });
    }, []);

    // ── Delete account ─────────────────────────────

    const deleteAccount = useCallback(
      async (password: string) => {
        if (!deleteAccountEndpoint) {
          setError('Account deletion not supported');
          return false;
        }
        setIsLoading(true);
        setError(null);
        try {
          const { error: fetchError } = await api.safeFetch<{ message: string }>(
            deleteAccountEndpoint,
            { method: 'DELETE', body: JSON.stringify({ password }) }
          );
          if (fetchError) {
            setError(fetchError);
            return false;
          }
          setUser(null);
          clearSession();
          if (refreshTimerRef.current) clearInterval(refreshTimerRef.current);
          onLogout?.();
          return true;
        } finally {
          setIsLoading(false);
        }
      },
      [api]
    );

    return (
      <AuthContext.Provider
        value={{
          user,
          isAuthenticated: !!user,
          isLoading,
          error,
          success,
          login,
          register,
          logout,
          forgotPassword,
          changePassword,
          deleteAccount,
          updateUser,
          clearMessages,
          // SmartAuth: Social login (Phase 1C)
          loginWithGoogle,
          loginWithMicrosoft,
          loginWithApple,
          // SmartAuth: Provider management (Phase 1C)
          providers,
          linkProvider,
          unlinkProvider,
          refreshProviders,
          // SmartAuth: MFA state (Phase 2D)
          mfaRequired,
          mfaMethods,
          mfaChallenge,
          verifyMfa,
        }}
      >
        {children}
      </AuthContext.Provider>
    );
  }

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

  return { AuthProvider, useAuth };
}