frontend/src/hooks/useTaskPolling.ts

import { useState, useCallback, useRef } from "react";
import { api } from "@/lib/api";

interface TaskStatus {
  status:
    | "pending"
    | "running"
    | "completed"
    | "failed"
    | "COMPLETED"
    | "FAILED";
  progress?: number;
  message?: string;
  error?: string;
}

interface UseTaskPollingOptions {
  onSuccess?: (taskId: string) => void;
  onError?: (error: string, taskId: string) => void;
  onProgress?: (progress: number, message?: string) => void;
  maxAttempts?: number;
  interval?: number;
  enableExponentialBackoff?: boolean;
}

interface UseTaskPollingReturn {
  pollTaskStatus: (taskId: string) => void;
  stopPolling: () => void;
  isPolling: boolean;
  currentAttempts: number;
}

export function useTaskPolling(
  options: UseTaskPollingOptions = {}
): UseTaskPollingReturn {
  const {
    onSuccess,
    onError,
    onProgress,
    maxAttempts = 180, // Maximum 15 minutes with exponential backoff
    interval = 5000,
    enableExponentialBackoff = true,
  } = options;

  const [isPolling, setIsPolling] = useState(false);
  const [currentAttempts, setCurrentAttempts] = useState(0);
  const pollTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const attemptsRef = useRef(0);
  const consecutiveErrorsRef = useRef(0);

  const stopPolling = useCallback(() => {
    if (pollTimeoutRef.current) {
      clearTimeout(pollTimeoutRef.current);
      pollTimeoutRef.current = null;
    }
    setIsPolling(false);
    setCurrentAttempts(0);
    attemptsRef.current = 0;
    consecutiveErrorsRef.current = 0;
  }, []);

  const getPollingInterval = useCallback(
    (attempt: number): number => {
      if (!enableExponentialBackoff) {
        return interval;
      }

      // Exponential backoff: start with 5s, max out at 30s
      // First 12 attempts: 5s
      // Next 18 attempts: 10s
      // Next 30 attempts: 15s
      // Remaining attempts: 30s
      if (attempt <= 12) {
        return 5000;
      } else if (attempt <= 30) {
        return 10000;
      } else if (attempt <= 60) {
        return 15000;
      } else {
        return 30000;
      }
    },
    [interval, enableExponentialBackoff]
  );

  const pollTaskStatus = useCallback(
    async (taskId: string) => {
      setIsPolling(true);
      attemptsRef.current = 0;
      consecutiveErrorsRef.current = 0;
      setCurrentAttempts(0);

      const poll = async () => {
        try {
          const taskStatus: TaskStatus = await api.tasks.get(taskId);

          // Reset consecutive errors on successful request
          consecutiveErrorsRef.current = 0;

          // Update progress if available
          if (taskStatus.progress && onProgress) {
            onProgress(taskStatus.progress, taskStatus.message);
          }

          // Check for completion
          if (
            taskStatus.status === "completed" ||
            taskStatus.status === "COMPLETED"
          ) {
            stopPolling();
            if (onSuccess) {
              onSuccess(taskId);
            }
            return;
          }

          // Check for failure
          if (
            taskStatus.status === "failed" ||
            taskStatus.status === "FAILED"
          ) {
            stopPolling();
            if (onError) {
              onError(taskStatus.error || "Task failed", taskId);
            }
            return;
          }

          // Continue polling if still running
          attemptsRef.current++;
          setCurrentAttempts(attemptsRef.current);

          if (attemptsRef.current < maxAttempts) {
            const nextInterval = getPollingInterval(attemptsRef.current);
            pollTimeoutRef.current = setTimeout(poll, nextInterval);
          } else {
            // Timeout reached - provide more helpful message based on last known progress
            stopPolling();
            if (onError) {
              let timeoutMessage = "Task polling timeout after 15 minutes";

              if (taskStatus.progress && taskStatus.progress > 0) {
                timeoutMessage = `Task was ${Math.round(
                  taskStatus.progress
                )}% complete when polling timed out. The process may still be running in the background. You can refresh the page or try again in a few minutes.`;
              } else {
                timeoutMessage =
                  "Task polling timed out. The process may still be running in the background. Please refresh the page or check your task manually.";
              }

              onError(timeoutMessage, taskId);
            }
          }
        } catch (error) {
          console.error("Error polling task status:", error);
          consecutiveErrorsRef.current++;

          // If we have too many consecutive errors, stop polling
          if (consecutiveErrorsRef.current >= 5) {
            stopPolling();
            if (onError) {
              onError(
                "Multiple polling errors occurred. Please check your connection and try again.",
                taskId
              );
            }
            return;
          }

          // On error, continue polling for a few more attempts with exponential backoff
          attemptsRef.current++;
          setCurrentAttempts(attemptsRef.current);

          if (attemptsRef.current < maxAttempts) {
            // Use longer interval after errors
            const errorInterval = Math.min(
              getPollingInterval(attemptsRef.current) * 2,
              60000
            );
            pollTimeoutRef.current = setTimeout(poll, errorInterval);
          } else {
            stopPolling();
            if (onError) {
              onError(
                error instanceof Error
                  ? `Polling failed: ${error.message}`
                  : "Unknown polling error",
                taskId
              );
            }
          }
        }
      };

      // Start polling with a 2-second delay to give the task time to start
      pollTimeoutRef.current = setTimeout(poll, 2000);
    },
    [
      maxAttempts,
      getPollingInterval,
      onSuccess,
      onError,
      onProgress,
      stopPolling,
    ]
  );

  return {
    pollTaskStatus,
    stopPolling,
    isPolling,
    currentAttempts,
  };
}