igny8/frontend/src/store/authStore.ts

/**
 * Authentication Store (Zustand)
 * Manages authentication state, user, and account information
 */
import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { fetchAPI } from '../services/api';

interface User {
  id: number;
  email: string;
  username: string;
  role: string;
  account?: {
    id: number;
    name: string;
    slug: string;
    credits: number;
    status: string;
  };
}

interface AuthState {
  user: User | null;
  token: string | null;
  refreshToken: string | null;
  isAuthenticated: boolean;
  loading: boolean;
  
  // Actions
  login: (email: string, password: string) => Promise<void>;
  logout: () => void;
  register: (data: any) => Promise<void>;
  setUser: (user: User | null) => void;
  setToken: (token: string | null) => void;
  refreshToken: () => Promise<void>;
  refreshUser: () => Promise<void>;
}

export const useAuthStore = create<AuthState>()(
  persist<AuthState>(
    (set, get) => ({
      user: null,
      token: null,
      isAuthenticated: false,
      loading: false, // Always start with loading false - will be set true only during login/register
      
      login: async (email, password) => {
        set({ loading: true });
        try {
          const API_BASE_URL = import.meta.env.VITE_BACKEND_URL || 'https://api.igny8.com/api';
          const response = await fetch(`${API_BASE_URL}/v1/auth/login/`, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
            },
            body: JSON.stringify({ email, password }),
          });

          const data = await response.json();

          if (!response.ok || !data.success) {
            throw new Error(data.message || 'Login failed');
          }

          // Store user and JWT tokens
          set({ 
            user: data.user, 
            token: data.tokens?.access || null,
            refreshToken: data.tokens?.refresh || null,
            isAuthenticated: true,
            loading: false 
          });
        } catch (error: any) {
          // ALWAYS reset loading on error - critical to prevent stuck state
          set({ loading: false });
          throw new Error(error.message || 'Login failed');
        } finally {
          // Extra safety: ensure loading is ALWAYS false after login attempt completes
          // This handles edge cases like network timeouts, browser crashes, etc.
          const current = get();
          if (current.loading) {
            set({ loading: false });
          }
        }
      },
      
      logout: () => {
        set({ user: null, token: null, refreshToken: null, isAuthenticated: false, loading: false });
      },
      
      register: async (registerData) => {
        set({ loading: true });
        try {
          const API_BASE_URL = import.meta.env.VITE_BACKEND_URL || 'https://api.igny8.com/api';
          const response = await fetch(`${API_BASE_URL}/v1/auth/register/`, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
            },
            body: JSON.stringify({
              ...registerData,
              password_confirm: registerData.password,  // Add password_confirm
            }),
          });

          const data = await response.json();

          if (!response.ok || !data.success) {
            // Handle validation errors
            const errorMessage = data.message || 
              (data.errors && typeof data.errors === 'object' 
                ? JSON.stringify(data.errors) 
                : data.errors) || 
              'Registration failed';
            throw new Error(errorMessage);
          }

          // Store user and JWT tokens
          set({ 
            user: data.user, 
            token: data.tokens?.access || null,
            refreshToken: data.tokens?.refresh || null,
            isAuthenticated: true,
            loading: false 
          });
        } catch (error: any) {
          // ALWAYS reset loading on error - critical to prevent stuck state
          set({ loading: false });
          throw new Error(error.message || 'Registration failed');
        } finally {
          // Extra safety: ensure loading is ALWAYS false after register attempt completes
          // This handles edge cases like network timeouts, browser crashes, etc.
          const current = get();
          if (current.loading) {
            set({ loading: false });
          }
        }
      },
      
      setUser: (user) => {
        set({ user, isAuthenticated: !!user });
      },
      
      setToken: (token) => {
        set({ token });
      },
      
      refreshToken: async () => {
        const state = get();
        if (!state.refreshToken) {
          throw new Error('No refresh token available');
        }
        
        try {
          const API_BASE_URL = import.meta.env.VITE_BACKEND_URL || 'https://api.igny8.com/api';
          const response = await fetch(`${API_BASE_URL}/v1/auth/refresh/`, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json',
            },
            body: JSON.stringify({ refresh: state.refreshToken }),
          });

          const data = await response.json();

          if (!response.ok || !data.success) {
            throw new Error(data.message || 'Token refresh failed');
          }

          // Update access token
          set({ token: data.access });
          
          // Also refresh user data to get latest account/plan information
          // This ensures account/plan changes are reflected immediately
          try {
            await get().refreshUser();
          } catch (userRefreshError) {
            // If user refresh fails, don't fail the token refresh
            console.warn('Failed to refresh user data during token refresh:', userRefreshError);
          }
        } catch (error: any) {
          // If refresh fails, logout user
          set({ user: null, token: null, refreshToken: null, isAuthenticated: false });
          throw new Error(error.message || 'Token refresh failed');
        }
      },
      
      refreshUser: async () => {
        const state = get();
        if (!state.token && !state.isAuthenticated) {
          throw new Error('Not authenticated');
        }
        
        try {
          // Use fetchAPI which handles token automatically and extracts data from unified format
          // fetchAPI is already imported at the top of the file
          const response = await fetchAPI('/v1/auth/me/');
          
          // fetchAPI extracts data field, so response is {user: {...}}
          if (!response || !response.user) {
            throw new Error('Failed to refresh user data');
          }

          // Update user data with latest from server
          // This ensures account/plan changes are reflected immediately
          set({ user: response.user });
        } catch (error: any) {
          // If refresh fails, don't logout - just log the error
          // User might still be authenticated, just couldn't refresh data
          console.warn('Failed to refresh user data:', error);
          // Don't throw - just log the warning to prevent error accumulation
        }
      },
    }),
    {
      name: 'auth-storage',
      // NEVER persist loading - it should always be false on app start
      partialize: (state) => ({ 
        user: state.user, 
        token: state.token,
        refreshToken: state.refreshToken,
        isAuthenticated: state.isAuthenticated,
        // Explicitly exclude loading from persistence
      }),
      // On rehydration, ensure loading is always false
      onRehydrateStorage: () => {
        return (state, error) => {
          // Always reset loading to false on rehydration - critical fix for stuck loading state
          if (state) {
            state.loading = false;
          }
          if (error) {
            console.error('Auth store rehydration error:', error);
            // If rehydration fails, ensure loading is false
            if (state) {
              state.loading = false;
            }
          }
        };
      },
    }
  )
);