import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
import {
  listTasks,
  getTask,
  startTask,
  stopTask,
  getTaskOutput,
  sendTaskMessage,
  type TaskSummary,
  type TaskWithSubtasks,
  type TaskOutputResponse,
} from '../lib/api';

// Query keys for consistent cache management
export const taskKeys = {
  all: ['tasks'] as const,
  lists: () => [...taskKeys.all, 'list'] as const,
  list: () => [...taskKeys.lists()] as const,
  details: () => [...taskKeys.all, 'detail'] as const,
  detail: (id: string) => [...taskKeys.details(), id] as const,
  output: (id: string) => [...taskKeys.all, 'output', id] as const,
};

/**
 * Hook to fetch the list of all tasks
 * Automatically refetches every 5 seconds for live updates
 */
export function useTasks() {
  return useQuery({
    queryKey: taskKeys.list(),
    queryFn: async () => {
      const response = await listTasks();
      return response.tasks;
    },
    refetchInterval: 5000, // Poll every 5 seconds for updates
    staleTime: 2000, // Consider data stale after 2 seconds
  });
}

/**
 * Hook to fetch a specific task with its subtasks
 */
export function useTask(taskId: string | null) {
  return useQuery({
    queryKey: taskKeys.detail(taskId ?? ''),
    queryFn: () => getTask(taskId!),
    enabled: !!taskId,
    refetchInterval: 5000,
  });
}

/**
 * Hook to fetch task output history
 */
export function useTaskOutput(taskId: string | null) {
  return useQuery({
    queryKey: taskKeys.output(taskId ?? ''),
    queryFn: () => getTaskOutput(taskId!),
    enabled: !!taskId,
    refetchInterval: 3000, // More frequent updates for output
  });
}

/**
 * Hook to start a task
 */
export function useStartTask() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: startTask,
    onSuccess: (updatedTask) => {
      // Invalidate task list to refetch
      queryClient.invalidateQueries({ queryKey: taskKeys.lists() });
      // Update the specific task in cache
      queryClient.setQueryData(
        taskKeys.detail(updatedTask.id),
        (old: TaskWithSubtasks | undefined) => {
          if (old) {
            return { ...old, ...updatedTask };
          }
          return old;
        }
      );
    },
  });
}

/**
 * Hook to stop a task
 */
export function useStopTask() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: stopTask,
    onSuccess: (updatedTask) => {
      // Invalidate task list to refetch
      queryClient.invalidateQueries({ queryKey: taskKeys.lists() });
      // Update the specific task in cache
      queryClient.setQueryData(
        taskKeys.detail(updatedTask.id),
        (old: TaskWithSubtasks | undefined) => {
          if (old) {
            return { ...old, ...updatedTask };
          }
          return old;
        }
      );
    },
  });
}

/**
 * Hook to send a message to a task
 */
export function useSendTaskMessage() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: ({ taskId, message }: { taskId: string; message: string }) =>
      sendTaskMessage(taskId, message),
    onSuccess: (_, { taskId }) => {
      // Invalidate task output to refetch
      queryClient.invalidateQueries({ queryKey: taskKeys.output(taskId) });
    },
  });
}

/**
 * Helper to group tasks by status for display
 */
export function groupTasksByStatus(tasks: TaskSummary[]) {
  const groups = {
    running: [] as TaskSummary[],
    pending: [] as TaskSummary[],
    blocked: [] as TaskSummary[],
    completed: [] as TaskSummary[],
  };

  for (const task of tasks) {
    switch (task.status) {
      case 'running':
      case 'initializing':
      case 'starting':
        groups.running.push(task);
        break;
      case 'pending':
        groups.pending.push(task);
        break;
      case 'blocked':
      case 'paused':
        groups.blocked.push(task);
        break;
      case 'done':
      case 'failed':
      case 'merged':
        groups.completed.push(task);
        break;
    }
  }

  return groups;
}

/**
 * Get counts for dashboard display
 */
export function getTaskCounts(tasks: TaskSummary[]) {
  const counts = {
    total: tasks.length,
    running: 0,
    pending: 0,
    blocked: 0,
    completed: 0,
    failed: 0,
  };

  for (const task of tasks) {
    switch (task.status) {
      case 'running':
      case 'initializing':
      case 'starting':
        counts.running++;
        break;
      case 'pending':
        counts.pending++;
        break;
      case 'blocked':
      case 'paused':
        counts.blocked++;
        break;
      case 'done':
      case 'merged':
        counts.completed++;
        break;
      case 'failed':
        counts.failed++;
        break;
    }
  }

  return counts;
}