CIAgent/apps/web/hooks/use-auth.ts

"use client";

import { useRouter } from "next/navigation";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { ApiError, api, clearTokens, getRefreshToken, setTokens } from "@/lib/api-client";
import type {
  BanIpPayload,
  ChangePasswordPayload,
  ConfirmPasswordResetPayload,
  DeleteAccountPayload,
  LoginPayload,
  RegisterPayload,
  RequestPasswordResetPayload,
  ResendVerificationPayload,
  SetSystemSecretPayload,
  SetUserApiKeyPayload,
  SystemSecretStatus,
  UnbanRequestPayload,
  UserApiKeyStatus,
  VerifyEmailPayload,
} from "@/lib/types";

export function useSystemStatus() {
  return useQuery({
    queryKey: ["system-status"],
    queryFn: api.systemStatus,
    staleTime: 60_000,
  });
}

export function useSystemSecrets() {
  return useQuery({
    queryKey: ["system-secrets"],
    queryFn: api.listSystemSecrets,
  });
}

export function useSetSystemSecret() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: ({ key, payload }: { key: string; payload: SetSystemSecretPayload }) =>
      api.setSystemSecret(key, payload),
    onSuccess: (updated) => {
      queryClient.setQueryData<SystemSecretStatus[]>(["system-secrets"], (prev) =>
        prev?.map((s) => (s.key === updated.key ? updated : s)),
      );
      queryClient.invalidateQueries({ queryKey: ["system-status"] });
    },
  });
}

export function useCreateDbViewerSession() {
  return useMutation({
    mutationFn: () => api.createDbViewerSession(),
  });
}

export function useSystemLogs() {
  return useQuery({
    queryKey: ["system-logs"],
    queryFn: api.systemLogs,
    refetchInterval: 5_000,
  });
}

export function useCurrentUser() {
  return useQuery({
    queryKey: ["me"],
    queryFn: api.me,
    retry: false,
  });
}

export function useLogin() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (payload: LoginPayload) => api.login(payload),
    onSuccess: (tokens) => {
      setTokens(tokens.access_token, tokens.refresh_token);
      return queryClient.invalidateQueries({ queryKey: ["me"] });
    },
  });
}

export function useRegister() {
  return useMutation({
    mutationFn: (payload: RegisterPayload) => api.register(payload),
  });
}

export function useLogout() {
  const router = useRouter();
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: async () => {
      const refreshToken = getRefreshToken();
      if (refreshToken) {
        await api.logout(refreshToken).catch(() => undefined);
      }
    },
    onSuccess: () => {
      clearTokens();
      // We already know definitively the user is logged out here - clear
      // the whole cache and navigate immediately, rather than invalidating
      // and waiting for a background refetch of ["me"] to fail with 401.
      // That indirection left a window where the dashboard kept rendering
      // stale cached data until the reactive layout redirect eventually
      // caught up, flickering between dashboard and login on every logout.
      queryClient.clear();
      router.replace("/login");
    },
  });
}

export function useDeleteAccount() {
  const router = useRouter();
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (payload: DeleteAccountPayload) => api.deleteAccount(payload),
    onSuccess: () => {
      // Same reasoning as useLogout: we know for certain the account is
      // gone, so clear and navigate immediately rather than waiting on a
      // background refetch of ["me"] to fail.
      clearTokens();
      queryClient.clear();
      router.replace("/");
    },
  });
}

export function useChangePassword() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (payload: ChangePasswordPayload) => api.changePassword(payload),
    onSuccess: (tokens) => {
      setTokens(tokens.access_token, tokens.refresh_token);
      return queryClient.invalidateQueries({ queryKey: ["me"] });
    },
  });
}

export function useVerifyEmail() {
  return useMutation({
    mutationFn: (payload: VerifyEmailPayload) => api.verifyEmail(payload),
  });
}

export function useResendVerification() {
  return useMutation({
    mutationFn: (payload: ResendVerificationPayload) => api.resendVerification(payload),
  });
}

export function useRequestPasswordReset() {
  return useMutation({
    mutationFn: (payload: RequestPasswordResetPayload) => api.requestPasswordReset(payload),
  });
}

export function useConfirmPasswordReset() {
  return useMutation({
    mutationFn: (payload: ConfirmPasswordResetPayload) => api.confirmPasswordReset(payload),
  });
}

export function useSecurityEvents() {
  return useQuery({
    queryKey: ["security-events"],
    queryFn: api.securityEvents,
  });
}

export function useSubmitUnbanRequest() {
  return useMutation({
    mutationFn: (payload: UnbanRequestPayload) => api.submitUnbanRequest(payload),
  });
}

export function useIpBans() {
  return useQuery({
    queryKey: ["ip-bans"],
    queryFn: api.listIpBans,
  });
}

export function useUnbanRequests() {
  return useQuery({
    queryKey: ["unban-requests"],
    queryFn: api.listUnbanRequests,
  });
}

export function useDeleteIpBan() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (ipAddress: string) => api.deleteIpBan(ipAddress),
    onSuccess: () => queryClient.invalidateQueries({ queryKey: ["ip-bans"] }),
  });
}

export function useCreateIpBan() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (payload: BanIpPayload) => api.createIpBan(payload),
    onSuccess: () => queryClient.invalidateQueries({ queryKey: ["ip-bans"] }),
  });
}

export function useAcceptUnbanRequest() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (requestId: string) => api.acceptUnbanRequest(requestId),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["ip-bans"] });
      queryClient.invalidateQueries({ queryKey: ["unban-requests"] });
    },
  });
}

export function useRejectUnbanRequest() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (requestId: string) => api.rejectUnbanRequest(requestId),
    onSuccess: () => queryClient.invalidateQueries({ queryKey: ["unban-requests"] }),
  });
}

export function useUserApiKeys() {
  return useQuery({
    queryKey: ["user-api-keys"],
    queryFn: api.listUserApiKeys,
  });
}

export function useSetUserApiKey() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: ({ provider, payload }: { provider: string; payload: SetUserApiKeyPayload }) =>
      api.setUserApiKey(provider, payload),
    onSuccess: (updated) => {
      queryClient.setQueryData<UserApiKeyStatus[]>(["user-api-keys"], (prev) =>
        prev?.map((k) => (k.provider === updated.provider ? updated : k)),
      );
    },
  });
}

/** Seconds until a throttled action (resend/reset) can be retried, or null
 * if the most recent error wasn't a throttle response - drives a live
 * countdown on the resend button. */
export function retryAfterSeconds(error: unknown): number | null {
  if (error instanceof ApiError && typeof error.body?.retry_after_seconds === "number") {
    return error.body.retry_after_seconds;
  }
  return null;
}

export type LoginFailureKind =
  | "verify_email"
  | "account_locked"
  | "throttled"
  | "banned"
  | "generic";

export interface LoginFailure {
  kind: LoginFailureKind;
  message: string;
  retryAfterSeconds: number | null;
}

/** The backend keeps login error messages generic on purpose (never
 * confirms account existence), so the UI distinguishes cases by status code
 * and message shape rather than a structured error code. */
export function classifyLoginError(error: unknown): LoginFailure {
  const message = authErrorMessage(error);
  if (error instanceof ApiError && error.status === 429) {
    const retry = retryAfterSeconds(error);
    return { kind: retry !== null ? "throttled" : "banned", message, retryAfterSeconds: retry };
  }
  if (/verify your email/i.test(message)) {
    return { kind: "verify_email", message, retryAfterSeconds: null };
  }
  if (/account locked/i.test(message)) {
    return { kind: "account_locked", message, retryAfterSeconds: null };
  }
  return { kind: "generic", message, retryAfterSeconds: null };
}

export function authErrorMessage(error: unknown): string {
  if (error instanceof ApiError) {
    if (error.body?.detail) return error.body.detail;
    if (error.body?.errors?.length) {
      return error.body.errors
        .map((e) => e.msg)
        .filter(Boolean)
        .join(", ");
    }
  }
  return "Something went wrong. Please try again.";
}