nats-console/apps/web/stores/auth.ts

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { api } from '@/lib/api';

interface User {
  id: string;
  email: string;
  firstName: string | null;
  lastName: string | null;
  avatarUrl: string | null;
  mfaEnabled?: boolean;
  role?: string;
}

interface MfaPending {
  mfaToken: string;
  userId: string;
}

interface LoginResult {
  success: boolean;
  mfaRequired?: boolean;
  mfaPending?: MfaPending;
}

interface AuthState {
  user: User | null;
  accessToken: string | null;
  refreshToken: string | null;
  orgId: string | null;
  isLoading: boolean;
  isAuthenticated: boolean;
  _hasHydrated: boolean;
  mfaPending: MfaPending | null;

  login: (email: string, password: string) => Promise<LoginResult>;
  loginWithMfa: (code: string) => Promise<void>;
  clearMfaPending: () => void;
  register: (data: { email: string; password: string; firstName: string; lastName: string }) => Promise<void>;
  logout: () => Promise<void>;
  refreshTokens: () => Promise<void>;
  setUser: (user: User) => void;
  setHasHydrated: (state: boolean) => void;
}

export const useAuthStore = create<AuthState>()(
  persist(
    (set, get) => ({
      user: null,
      accessToken: null,
      refreshToken: null,
      orgId: null,
      isLoading: false,
      isAuthenticated: false,
      _hasHydrated: false,
      mfaPending: null,

      setHasHydrated: (state: boolean) => {
        set({ _hasHydrated: state });
      },

      login: async (email, password) => {
        set({ isLoading: true });
        try {
          const result = await api.auth.login(email, password);

          // Check if MFA is required
          if (result.mfaRequired) {
            set({
              isLoading: false,
              mfaPending: {
                mfaToken: result.mfaToken,
                userId: result.userId,
              },
            });
            return { success: false, mfaRequired: true, mfaPending: { mfaToken: result.mfaToken, userId: result.userId } };
          }

          // No MFA required, complete login
          const { user, tokens, orgId } = result;

          // Store token in localStorage for API client
          localStorage.setItem('accessToken', tokens.accessToken);

          set({
            user,
            accessToken: tokens.accessToken,
            refreshToken: tokens.refreshToken,
            orgId,
            isAuthenticated: true,
            isLoading: false,
            mfaPending: null,
          });

          return { success: true };
        } catch (error) {
          set({ isLoading: false });
          throw error;
        }
      },

      loginWithMfa: async (code: string) => {
        const { mfaPending } = get();
        if (!mfaPending) {
          throw new Error('No MFA pending');
        }

        set({ isLoading: true });
        try {
          const { user, tokens, orgId } = await api.auth.loginWithMfa(mfaPending.mfaToken, code);

          // Store token in localStorage for API client
          localStorage.setItem('accessToken', tokens.accessToken);

          set({
            user,
            accessToken: tokens.accessToken,
            refreshToken: tokens.refreshToken,
            orgId,
            isAuthenticated: true,
            isLoading: false,
            mfaPending: null,
          });
        } catch (error) {
          set({ isLoading: false });
          throw error;
        }
      },

      clearMfaPending: () => {
        set({ mfaPending: null });
      },

      register: async (data) => {
        set({ isLoading: true });
        try {
          const { user, tokens, orgId } = await api.auth.register(data);

          localStorage.setItem('accessToken', tokens.accessToken);

          set({
            user,
            accessToken: tokens.accessToken,
            refreshToken: tokens.refreshToken,
            orgId,
            isAuthenticated: true,
            isLoading: false,
          });
        } catch (error) {
          set({ isLoading: false });
          throw error;
        }
      },

      logout: async () => {
        try {
          await api.auth.logout();
        } catch {
          // Ignore errors during logout
        }

        localStorage.removeItem('accessToken');

        set({
          user: null,
          accessToken: null,
          refreshToken: null,
          orgId: null,
          isAuthenticated: false,
        });
      },

      refreshTokens: async () => {
        const { refreshToken } = get();
        if (!refreshToken) {
          throw new Error('No refresh token');
        }

        try {
          const { tokens } = await api.auth.refresh(refreshToken);

          localStorage.setItem('accessToken', tokens.accessToken);

          set({
            accessToken: tokens.accessToken,
            refreshToken: tokens.refreshToken,
          });
        } catch (error) {
          // Refresh failed, logout
          get().logout();
          throw error;
        }
      },

      setUser: (user) => set({ user }),
    }),
    {
      name: 'auth-storage',
      partialize: (state) => ({
        user: state.user,
        accessToken: state.accessToken,
        refreshToken: state.refreshToken,
        orgId: state.orgId,
        isAuthenticated: state.isAuthenticated,
      }),
      onRehydrateStorage: () => (state) => {
        // Called after hydration is complete
        state?.setHasHydrated(true);
        // Sync localStorage with hydrated token
        if (state?.accessToken) {
          localStorage.setItem('accessToken', state.accessToken);
        }
      },
    }
  )
);