MAIA/interfaces/nativeapp/src/contexts/AuthContext.tsx

// src/contexts/AuthContext.tsx
import React, { createContext, useState, useEffect, useContext, useMemo, useCallback } from 'react';
import { Platform, ActivityIndicator, View, StyleSheet } from 'react-native'; // Import Platform
import apiClient from '../api/client';
import { useTheme } from 'react-native-paper';
import { getAccessToken, getRefreshToken, storeTokens, clearTokens } from '../api/client';

// Define UserRole enum matching the backend
enum UserRole {
  USER = 'user',
  ADMIN = 'admin',
}

interface UserData {
  username: string;
  name: string;
  role: UserRole;
  // Add other user fields if needed
}

interface AuthContextData {
  isAuthenticated: boolean;
  isLoading: boolean;
  user: UserData | null; // Add user data to context
  login: (username: string, password: string) => Promise<void>;
  logout: () => Promise<void>;
  register: (username: string, password: string, name: string) => Promise<void>; // Add register function
}

const AuthContext = createContext<AuthContextData>({
  isAuthenticated: false,
  isLoading: true,
  user: null, // Initialize user as null
  login: async () => { throw new Error('AuthContext not initialized'); },
  logout: async () => { throw new Error('AuthContext not initialized'); },
  register: async () => { throw new Error('AuthContext not initialized'); }, // Add register initializer
});

interface AuthProviderProps {
  children: React.ReactNode;
}

export const AuthProvider: React.FC<AuthProviderProps> = ({ children }) => {
  const [isLoading, setIsLoading] = useState<boolean>(true);
  const [isAuthenticatedState, setIsAuthenticatedState] = useState<boolean>(false);
  const [userState, setUserState] = useState<UserData | null>(null); // State for user data

  // Function to fetch user data
  const fetchUserData = useCallback(async () => {
    try {
      console.log("[AuthContext] fetchUserData: Fetching /user/me");
      const response = await apiClient.get<UserData>('/user/me');
      console.log("[AuthContext] fetchUserData: User data received:", response.data);
      setUserState(response.data);
      return response.data;
    } catch (error: any) {
      console.error("[AuthContext] fetchUserData: Error fetching user data:", error.response?.data || error.message);
      // If fetching user fails (e.g., token expired mid-session), log out
      await clearTokens();
      setIsAuthenticatedState(false);
      setUserState(null);
      return null;
    }
  }, []);

  const checkAuthStatus = useCallback(async () => {
    const token = await getAccessToken();
    const hasToken = !!token;
    if (hasToken) {
      console.log("[AuthContext] checkAuthStatus: Token found, fetching user data.");
      const userData = await fetchUserData();
      if (userData) {
        setIsAuthenticatedState(true);
      } else {
        // Fetch failed, already handled logout in fetchUserData
        setIsAuthenticatedState(false);
      }
    } else {
      console.log("[AuthContext] checkAuthStatus: No token found.");
      setIsAuthenticatedState(false);
      setUserState(null);
    }
    return isAuthenticatedState; // Return the updated state
  }, [fetchUserData, isAuthenticatedState]); // Added isAuthenticatedState dependency

  const loadInitialAuth = useCallback(async () => {
    setIsLoading(true);
    try {
      console.log("[AuthContext] loadInitialAuth: Checking initial auth status");
      await checkAuthStatus();
      console.log("[AuthContext] loadInitialAuth: Initial check complete.");
    } catch (error) {
      console.error("[AuthContext] loadInitialAuth: Error during initial auth check:", error);
      await clearTokens();
      setIsAuthenticatedState(false);
      setUserState(null);
    } finally {
      setIsLoading(false);
    }
  }, [checkAuthStatus]);

  useEffect(() => {
    loadInitialAuth();
  }, [loadInitialAuth]);

  const login = useCallback(async (username: string, password: string) => {
    console.log("[AuthContext] login: Function called with:", username);
    try {
      // ... (existing login API call) ...
      const response = await apiClient.post(
        '/auth/login',
        'grant_type=password&username=' + encodeURIComponent(username) + '&password=' + encodeURIComponent(password) + '&scope=&client_id=&client_secret=',
        {
          headers: {
            'accept': 'application/json',
            'Content-Type': 'application/x-www-form-urlencoded',
          },
        }
      );

      const { access_token, refresh_token } = response.data;
      // ... (existing token validation) ...
      if (!access_token || typeof access_token !== 'string' || !refresh_token) {
         console.error("[AuthContext] login: Invalid token structure received:", response.data);
         throw new Error('Invalid token received from server.');
      }

      console.log('[AuthContext] login: Login successful, storing tokens.');
      await storeTokens(access_token, refresh_token);

      // Fetch user data immediately after successful login
      const userData = await fetchUserData();
      if (userData) {
        setIsAuthenticatedState(true);
      } else {
        // Should not happen if login succeeded, but handle defensively
        throw new Error('Failed to fetch user data after login.');
      }

    } catch (error: any) {
      // ... (existing error handling) ...
      console.error("[AuthContext] login: Caught Error Object:", error);
      await clearTokens();
      setIsAuthenticatedState(false);
      setUserState(null);
      throw error;
    }
  }, [fetchUserData]); // Added fetchUserData dependency

  const register = useCallback(async (username: string, password: string, name: string) => {
    console.log("[AuthContext] register: Function called with:", username, name);
    try {
      // Call the backend register endpoint
      const response = await apiClient.post('/auth/register', {
        username,
        password,
        name,
      }, {
        headers: {
          'accept': 'application/json',
          'Content-Type': 'application/json',
        },
      });

      console.log('[AuthContext] register: Registration successful:', response.data);
      // Optionally, you could automatically log the user in here
      // For now, we'll just let the user log in manually after registering
      // Or display a success message and navigate back to login

    } catch (error: any) {
      console.error("[AuthContext] register: Caught Error Object:", error);
      // Rethrow the error so the UI can handle it (e.g., display specific messages)
      throw error;
    }
  }, []); // No dependencies needed for register itself

  const logout = useCallback(async () => {
    console.log('[AuthContext] logout: Logging out.');
    const refreshToken = await getRefreshToken();
    // ... (existing backend logout call) ...
    try {
        if (refreshToken) {
            console.log('[AuthContext] logout: Calling backend /auth/logout');
            await apiClient.post("/auth/logout", { refresh_token: refreshToken });
            console.log('[AuthContext] logout: Backend logout call successful (or ignored error).');
        }
    } catch (error: any) {
        console.error('[AuthContext] logout: Error calling backend logout:', error.response?.data || error.message);
    } finally {
        await clearTokens();
        setIsAuthenticatedState(false);
        setUserState(null); // Clear user data on logout
        console.log('[AuthContext] logout: Local tokens cleared and state updated.');
    }
  }, []);

  const contextValue = useMemo(() => ({
    isAuthenticated: isAuthenticatedState,
    isLoading,
    user: userState, // Provide user state
    login,
    logout,
    register, // Add register to context value
  }), [isAuthenticatedState, isLoading, userState, login, logout, register]); // Added register dependency

  // ... (rest of the component: Provider, useAuth, AuthLoadingScreen) ...
  return (
    <AuthContext.Provider value={contextValue}>
      {children}
    </AuthContext.Provider>
  );
};

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

export const AuthLoadingScreen: React.FC = () => {
    const theme = useTheme();
    const styles = StyleSheet.create({
        container: { flex: 1, justifyContent: 'center', alignItems: 'center', backgroundColor: theme.colors.background }
    });
    return (<View style={styles.container}><ActivityIndicator size="large" color={theme.colors.primary} /></View>);
}

// Export UserRole if needed elsewhere
export { UserRole };