mimic/frontend/src/hooks/useC2.ts

import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import {
  deleteC2Config,
  executeC2,
  getC2Config,
  getC2Tasks,
  importC2,
  listCallbackHistory,
  listCallbacks,
  putC2Config,
  testC2Config,
} from '@/api/c2';
import type {
  C2ConfigInput,
  C2ExecuteInput,
  C2ImportInput,
  C2TasksResponse,
} from '@/api/types';

function c2ConfigKey(engagementId: number) {
  return ['c2-config', engagementId] as const;
}

function c2CallbacksKey(engagementId: number) {
  return ['c2-callbacks', engagementId] as const;
}

function c2TasksKey(simulationId: number) {
  return ['c2-tasks', simulationId] as const;
}

function c2HistoryKey(engagementId: number, callbackDisplayId: number, page: number, pageSize: number) {
  return ['c2-history', engagementId, callbackDisplayId, page, pageSize] as const;
}

function simulationKey(id: number) {
  return ['simulations', id] as const;
}

export function useC2Config(engagementId: number | undefined, options?: { enabled?: boolean }) {
  const enabled =
    typeof engagementId === 'number' &&
    !Number.isNaN(engagementId) &&
    (options?.enabled !== false);

  return useQuery({
    queryKey: engagementId ? c2ConfigKey(engagementId) : ['c2-config', 'none'],
    queryFn: () => getC2Config(engagementId as number),
    enabled,
  });
}

export function useUpdateC2Config(engagementId: number) {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: (input: C2ConfigInput) => putC2Config(engagementId, input),
    onSuccess: () => qc.invalidateQueries({ queryKey: c2ConfigKey(engagementId) }),
  });
}

export function useDeleteC2Config(engagementId: number) {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: () => deleteC2Config(engagementId),
    onSuccess: () => qc.invalidateQueries({ queryKey: c2ConfigKey(engagementId) }),
  });
}

export function useTestC2Config(engagementId: number) {
  return useMutation({
    mutationFn: () => testC2Config(engagementId),
  });
}

export function useC2Callbacks(engagementId: number | undefined, options?: { enabled?: boolean }) {
  const enabled =
    typeof engagementId === 'number' &&
    !Number.isNaN(engagementId) &&
    (options?.enabled !== false);

  return useQuery({
    queryKey: engagementId ? c2CallbacksKey(engagementId) : ['c2-callbacks', 'none'],
    queryFn: () => listCallbacks(engagementId as number),
    enabled,
  });
}

export function useC2Tasks(simulationId: number | undefined, options?: { enabled?: boolean }) {
  const enabled =
    typeof simulationId === 'number' &&
    !Number.isNaN(simulationId) &&
    (options?.enabled !== false);

  return useQuery({
    queryKey: simulationId ? c2TasksKey(simulationId) : ['c2-tasks', 'none'],
    queryFn: () => getC2Tasks(simulationId as number),
    enabled,
    // Poll every 2500 ms while any task is incomplete; stop when all done.
    refetchInterval: (query: { state: { data?: C2TasksResponse } }) =>
      query.state.data?.tasks?.some((t) => !t.completed) ? 2500 : false,
  });
}

export function useExecuteC2(simulationId: number, engagementId: number) {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: (input: C2ExecuteInput) => executeC2(simulationId, input),
    onSuccess: () => {
      qc.invalidateQueries({ queryKey: simulationKey(simulationId) });
      qc.invalidateQueries({ queryKey: ['engagements', engagementId, 'simulations'] });
      qc.invalidateQueries({ queryKey: c2TasksKey(simulationId) });
    },
  });
}

export function useC2CallbackHistory(
  engagementId: number | undefined,
  callbackDisplayId: number | null,
  options?: { page?: number; pageSize?: number; enabled?: boolean },
) {
  const page = options?.page ?? 1;
  const pageSize = options?.pageSize ?? 25;
  const enabled =
    typeof engagementId === 'number' &&
    !Number.isNaN(engagementId) &&
    callbackDisplayId !== null &&
    (options?.enabled !== false);

  return useQuery({
    queryKey:
      engagementId && callbackDisplayId !== null
        ? c2HistoryKey(engagementId, callbackDisplayId, page, pageSize)
        : ['c2-history', 'none'],
    queryFn: () =>
      listCallbackHistory(engagementId as number, callbackDisplayId as number, {
        page,
        pageSize,
      }),
    enabled,
  });
}

export function useImportC2(simulationId: number) {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: (input: C2ImportInput) => importC2(simulationId, input),
    onSuccess: () => {
      qc.invalidateQueries({ queryKey: c2TasksKey(simulationId) });
      qc.invalidateQueries({ queryKey: simulationKey(simulationId) });
    },
  });
}