bakery-ia/frontend/src/contexts/AuthContext.tsx

// frontend/src/contexts/AuthContext.tsx - UPDATED TO USE NEW REGISTRATION FLOW
import React, { createContext, useContext, useEffect, useState, useCallback } from 'react';
import { authService, UserProfile, RegisterData } from '../api/services/authService';
import { tokenManager } from '../api/auth/tokenManager';

interface AuthContextType {
  user: UserProfile | null;
  isAuthenticated: boolean;
  isLoading: boolean;
  login: (email: string, password: string) => Promise<void>;
  register: (data: RegisterData) => Promise<void>; // SIMPLIFIED - no longer needs auto-login
  logout: () => Promise<void>;
  updateProfile: (updates: Partial<UserProfile>) => Promise<void>;
  refreshUser: () => Promise<void>;
}

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

export const useAuth = () => {
  const context = useContext(AuthContext);
  if (!context) {
    throw new Error('useAuth must be used within an AuthProvider');
  }
  return context;
};

export const AuthProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
  const [user, setUser] = useState<UserProfile | null>(null);
  const [isLoading, setIsLoading] = useState(true);

  // Initialize auth state
  useEffect(() => {
    const initAuth = async () => {
      try {
        await tokenManager.initialize();

        if (authService.isAuthenticated()) {
          // Get user from token first (faster), then validate with API
          const tokenUser = tokenManager.getUserFromToken();
          if (tokenUser) {
            setUser({
              id: tokenUser.user_id,
              email: tokenUser.email,
              full_name: tokenUser.full_name,
              is_active: true,
              is_verified: tokenUser.is_verified,
              created_at: '', // Will be filled by API call
            });
          }

          // Validate with API and get complete profile
          try {
            const profile = await authService.getCurrentUser();
            setUser(profile);
          } catch (error) {
            console.error('Failed to fetch user profile:', error);
            // Keep token-based user data if API fails
          }
        }
      } catch (error) {
        console.error('Auth initialization failed:', error);
        // Clear potentially corrupted tokens
        tokenManager.clearTokens();
      } finally {
        setIsLoading(false);
      }
    };

    initAuth();
  }, []);

  const login = useCallback(async (email: string, password: string) => {
    setIsLoading(true);
    try {
      const profile = await authService.login({ email, password });
      setUser(profile);
    } catch (error) {
      setIsLoading(false);
      throw error; // Re-throw to let components handle the error
    } finally {
      setIsLoading(false);
    }
  }, []);

  const register = useCallback(async (data: RegisterData) => {
    setIsLoading(true);
    try {
      // NEW: Registration now handles tokens internally - no auto-login needed!
      const profile = await authService.register(data);
      setUser(profile);
    } catch (error) {
      setIsLoading(false);
      throw error; // Re-throw to let components handle the error
    } finally {
      setIsLoading(false);
    }
  }, []);

  const logout = useCallback(async () => {
    setIsLoading(true);
    try {
      await authService.logout();
      setUser(null);
    } catch (error) {
      console.error('Logout error:', error);
      // Clear local state even if API call fails
      setUser(null);
      tokenManager.clearTokens();
    } finally {
      setIsLoading(false);
    }
  }, []);

  const updateProfile = useCallback(async (updates: Partial<UserProfile>) => {
    if (!user) return;
    
    try {
      const updated = await authService.updateProfile(updates);
      setUser(updated);
    } catch (error) {
      console.error('Profile update error:', error);
      throw error;
    }
  }, [user]);

  const refreshUser = useCallback(async () => {
    if (!authService.isAuthenticated()) return;
    
    try {
      const profile = await authService.getCurrentUser();
      setUser(profile);
    } catch (error) {
      console.error('User refresh error:', error);
      // If refresh fails with 401, user might need to re-login
      if (error.status === 401) {
        await logout();
      }
    }
  }, [logout]);

  // Set up token refresh interval
  useEffect(() => {
    if (!user) return;

    const interval = setInterval(async () => {
      try {
        await tokenManager.refreshAccessToken();
      } catch (error) {
        console.error('Scheduled token refresh failed:', error);
        // If token refresh fails, user needs to re-login
        await logout();
      }
    }, 60000); // Check every 1 minute

    return () => clearInterval(interval);
  }, [user, logout]);

  // Monitor token expiration
  useEffect(() => {
    if (!user) return;

    const checkTokenValidity = () => {
      if (!authService.isAuthenticated()) {
        console.warn('Token became invalid, logging out user');
        logout();
      }
    };

    // Check token validity every 30 seconds
    const interval = setInterval(checkTokenValidity, 30000);
    return () => clearInterval(interval);
  }, [user, logout]);

  const contextValue = {
    user,
    isAuthenticated: !!user && authService.isAuthenticated(),
    isLoading,
    login,
    register,
    logout,
    updateProfile,
    refreshUser,
  };

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