UMBRA/frontend/src/hooks/useAuth.ts

import { useState } from 'react';
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import api from '@/lib/api';
import type { AuthStatus, LoginResponse } from '@/types';

export function useAuth() {
  const queryClient = useQueryClient();
  const [mfaToken, setMfaToken] = useState<string | null>(null);
  const [mfaSetupRequired, setMfaSetupRequired] = useState(false);

  const authQuery = useQuery({
    queryKey: ['auth'],
    queryFn: async () => {
      const { data } = await api.get<AuthStatus>('/auth/status');
      return data;
    },
    retry: false,
  });

  const loginMutation = useMutation({
    mutationFn: async ({ username, password }: { username: string; password: string }) => {
      const { data } = await api.post<LoginResponse>('/auth/login', { username, password });
      return data;
    },
    onSuccess: (data) => {
      if ('mfa_setup_required' in data && data.mfa_setup_required) {
        // MFA enforcement — user must set up TOTP before accessing app
        setMfaSetupRequired(true);
        setMfaToken(data.mfa_token);
      } else if ('mfa_token' in data && 'totp_required' in data && data.totp_required) {
        // Regular TOTP challenge
        setMfaToken(data.mfa_token);
        setMfaSetupRequired(false);
      } else {
        setMfaToken(null);
        setMfaSetupRequired(false);
        // Optimistically mark authenticated to prevent form flash during refetch
        if ('authenticated' in data && data.authenticated && !('must_change_password' in data && data.must_change_password)) {
          queryClient.setQueryData(['auth'], (old: AuthStatus | undefined) => {
            if (!old) return old; // let invalidateQueries handle it
            return { ...old, authenticated: true };
          });
        }
        queryClient.invalidateQueries({ queryKey: ['auth'] });
      }
    },
  });

  const registerMutation = useMutation({
    mutationFn: async ({ username, password, email, date_of_birth, preferred_name }: {
      username: string; password: string; email: string; date_of_birth: string;
      preferred_name?: string;
    }) => {
      const payload: Record<string, string> = { username, password, email, date_of_birth };
      if (preferred_name) payload.preferred_name = preferred_name;
      const { data } = await api.post<LoginResponse & { message?: string }>('/auth/register', payload);
      return data;
    },
    onSuccess: (data) => {
      if ('mfa_setup_required' in data && data.mfa_setup_required) {
        setMfaSetupRequired(true);
        setMfaToken(data.mfa_token);
      } else {
        setMfaToken(null);
        setMfaSetupRequired(false);
        queryClient.invalidateQueries({ queryKey: ['auth'] });
      }
    },
  });

  const totpVerifyMutation = useMutation({
    mutationFn: async ({ code, isBackup }: { code: string; isBackup: boolean }) => {
      const payload: Record<string, string> = { mfa_token: mfaToken! };
      if (isBackup) {
        payload.backup_code = code;
      } else {
        payload.code = code;
      }
      const { data } = await api.post('/auth/totp-verify', payload);
      return data;
    },
    onSuccess: () => {
      setMfaToken(null);
      setMfaSetupRequired(false);
      queryClient.invalidateQueries({ queryKey: ['auth'] });
    },
  });

  const setupMutation = useMutation({
    mutationFn: async ({ username, password }: { username: string; password: string }) => {
      const { data } = await api.post('/auth/setup', { username, password });
      return data;
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ['auth'] });
    },
  });

  const logoutMutation = useMutation({
    mutationFn: async () => {
      const { data } = await api.post('/auth/logout');
      return data;
    },
    onSuccess: () => {
      setMfaToken(null);
      setMfaSetupRequired(false);
      queryClient.invalidateQueries({ queryKey: ['auth'] });
    },
  });

  return {
    authStatus: authQuery.data,
    isLoading: authQuery.isLoading,
    role: authQuery.data?.role ?? null,
    isAdmin: authQuery.data?.role === 'admin',
    mfaRequired: mfaToken !== null && !mfaSetupRequired,
    mfaSetupRequired,
    mfaToken,
    login: loginMutation.mutateAsync,
    register: registerMutation.mutateAsync,
    verifyTotp: totpVerifyMutation.mutateAsync,
    setup: setupMutation.mutateAsync,
    logout: logoutMutation.mutateAsync,
    isLoginPending: loginMutation.isPending,
    isRegisterPending: registerMutation.isPending,
    isTotpPending: totpVerifyMutation.isPending,
    isSetupPending: setupMutation.isPending,
  };
}