utp_service/frontend/src/api/tasks.ts

import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
import { api } from './client'
import { useSSE } from '../hooks/useSSE'

export interface TaskEvent {
  id: number
  message: string
  level: string
  created_at: string
}

export interface TaskResult {
  id: number
  device_id: number
  device_ip: string | null
  device_hostname: string | null
  object_id: number | null
  object_name: string | null
  rack_id: number | null
  rack_name: string | null
  status: string
  stdout: string | null
  stderr: string | null
  exit_code: number | null
  data: Record<string, unknown> | null
  duration_ms: number | null
  executed_at: string | null
}

export interface Task {
  id: string
  type: string
  status: string
  cancel_requested: boolean
  target_scope: Record<string, unknown>
  target_label: string | null
  params: Record<string, unknown>
  result: { total: number; ok: number; failed: number } | null
  error_message: string | null
  created_by: number | null
  created_at: string
  started_at: string | null
  finished_at: string | null
}

export interface TaskDetail extends Task {
  events: TaskEvent[]
  device_results: TaskResult[]
}

export interface ActionInfo {
  name: string
  display_name: string
  compatible_categories: string[]
  params_schema: { name: string; type: string; label: string; required?: boolean }[]
}

export interface TaskPreviewResponse {
  total: number
  compatible: number
  skipped: number
  by_category: Record<string, number>
  compatible_by_category: Record<string, number>
  skipped_categories: string[]
}

export const useTasks = (page = 1, pageSize = 50) =>
  useQuery({
    queryKey: ['tasks', page, pageSize],
    queryFn: () =>
      api
        .get<Task[]>('/api/v1/tasks', { params: { skip: (page - 1) * pageSize, limit: pageSize } })
        .then((r) => r.data),
    refetchInterval: 5000,
  })

export const useTask = (id: string | undefined, options?: { refetchInterval?: number | false }) =>
  useQuery({
    queryKey: ['tasks', id],
    queryFn: () => api.get<TaskDetail>(`/api/v1/tasks/${id}`).then((r) => r.data),
    enabled: !!id,
    refetchInterval: options?.refetchInterval,
  })

// Subscribe to SSE for a running task — merges results in cache, invalidates on completion.
export const useTaskSSE = (id: string | undefined) => {
  const qc = useQueryClient()
  useSSE(id, {
    onTaskResult: (data) => {
      qc.setQueryData<TaskDetail>(['tasks', id], (prev) => {
        if (!prev) return prev
        const existing = prev.device_results.find((r) => r.device_id === data.device_id)
        const updated: TaskResult = existing
          ? { ...existing, status: data.status, stdout: data.stdout, stderr: data.stderr, duration_ms: data.duration_ms }
          : {
              id: 0,
              device_id: data.device_id,
              device_ip: data.device_ip,
              device_hostname: data.device_hostname,
              object_id: null,
              object_name: null,
              rack_id: null,
              rack_name: null,
              status: data.status,
              stdout: data.stdout,
              stderr: data.stderr,
              exit_code: null,
              data: null,
              duration_ms: data.duration_ms,
              executed_at: null,
            }
        return {
          ...prev,
          device_results: existing
            ? prev.device_results.map((r) => (r.device_id === data.device_id ? updated : r))
            : [...prev.device_results, updated],
        }
      })
    },
    onTaskStatus: () => {
      // Task finished — do a full refetch to get final status and complete results
      qc.invalidateQueries({ queryKey: ['tasks', id] })
      qc.invalidateQueries({ queryKey: ['tasks'] })
    },
  })
}

export const useActions = () =>
  useQuery({
    queryKey: ['actions'],
    queryFn: () => api.get<ActionInfo[]>('/api/v1/tasks/actions').then((r) => r.data),
  })

export const useCreateTask = () => {
  const qc = useQueryClient()
  return useMutation({
    mutationFn: (body: { type: string; target_scope: object; params: object }) =>
      api.post<Task>('/api/v1/tasks', body).then((r) => r.data),
    onSuccess: () => qc.invalidateQueries({ queryKey: ['tasks'] }),
  })
}

export const useCancelTask = () => {
  const qc = useQueryClient()
  return useMutation({
    mutationFn: (taskId: string) =>
      api.post(`/api/v1/tasks/${taskId}/cancel`).then((r) => r.data),
    onSuccess: (_data, taskId) => {
      qc.invalidateQueries({ queryKey: ['tasks', taskId] })
      qc.invalidateQueries({ queryKey: ['tasks'] })
    },
  })
}

export const usePreviewTask = (
  type: string | undefined,
  target_scope: object | undefined,
  enabled: boolean,
) =>
  useQuery({
    queryKey: ['task-preview', type, target_scope],
    queryFn: () =>
      api
        .post<TaskPreviewResponse>('/api/v1/tasks/preview', { type, target_scope, params: {} })
        .then((r) => r.data),
    enabled: enabled && !!type && !!target_scope,
    staleTime: 10_000,
  })

export const useDeviceTasks = (deviceId: number) =>
  useQuery({
    queryKey: ['tasks', 'device', deviceId],
    queryFn: () =>
      api
        .get<Task[]>('/api/v1/tasks', { params: { device_id: deviceId, limit: 100 } })
        .then((r) => r.data),
    refetchInterval: 10_000,
  })