frontend/src/hooks/useSegmentation.ts

import { useState, useCallback, useRef, useEffect } from "react";
import { apiClient, ApiError } from "../api/client";
import type { SegmentationResult, JobStatus } from "../types";

// Polling interval in milliseconds
const POLLING_INTERVAL = 2000;

// Cold start retry configuration
const MAX_COLD_START_RETRIES = 5;
const INITIAL_RETRY_DELAY = 2000; // 2 seconds
const MAX_RETRY_DELAY = 30000; // 30 seconds

/**
 * Sleep utility for async delays
 */
const sleep = (ms: number): Promise<void> =>
  new Promise((resolve) => setTimeout(resolve, ms));

/**
 * Hook for running segmentation with async job polling.
 *
 * Instead of waiting for the full inference to complete (which can timeout
 * on HuggingFace Spaces), this hook:
 * 1. Creates a job that returns immediately with a job ID
 * 2. Polls for job status/progress every 2 seconds
 * 3. Returns results when the job completes
 *
 * This avoids the ~60s gateway timeout on HF Spaces while providing
 * real-time progress updates to the user.
 */
export function useSegmentation() {
  // Result state
  const [result, setResult] = useState<SegmentationResult | null>(null);
  const [error, setError] = useState<string | null>(null);

  // Job tracking state
  const [jobId, setJobId] = useState<string | null>(null);
  const [jobStatus, setJobStatus] = useState<JobStatus | null>(null);
  const [progress, setProgress] = useState(0);
  const [progressMessage, setProgressMessage] = useState("");
  const [elapsedSeconds, setElapsedSeconds] = useState<number | undefined>(
    undefined,
  );

  // Loading state - true from job creation until completion/failure
  const [isLoading, setIsLoading] = useState(false);

  // Refs for managing async operations
  const currentJobRef = useRef<string | null>(null);
  const pollingIntervalRef = useRef<number | null>(null);
  const abortControllerRef = useRef<AbortController | null>(null);

  /**
   * Stop polling for job status
   */
  const stopPolling = useCallback(() => {
    if (pollingIntervalRef.current) {
      clearInterval(pollingIntervalRef.current);
      pollingIntervalRef.current = null;
    }
  }, []);

  /**
   * Poll for job status and update state
   */
  const pollJobStatus = useCallback(
    async (id: string, signal: AbortSignal) => {
      // Don't poll if this isn't the current job
      if (id !== currentJobRef.current) {
        stopPolling();
        return;
      }

      try {
        const response = await apiClient.getJobStatus(id, signal);

        // Ignore results if job changed
        if (id !== currentJobRef.current) return;

        // Update progress state
        setJobStatus(response.status);
        setProgress(response.progress);
        setProgressMessage(response.progressMessage);
        setElapsedSeconds(response.elapsedSeconds);

        // Handle completion
        if (response.status === "completed" && response.result) {
          stopPolling();
          setIsLoading(false);
          setResult({
            dwiUrl: response.result.dwiUrl,
            predictionUrl: response.result.predictionUrl,
            metrics: {
              caseId: response.result.caseId,
              diceScore: response.result.diceScore,
              volumeMl: response.result.volumeMl,
              elapsedSeconds: response.result.elapsedSeconds,
              warning: response.result.warning,
            },
          });
        }

        // Handle failure
        if (response.status === "failed") {
          stopPolling();
          setIsLoading(false);
          setError(response.error || "Job failed");
          setResult(null);
        }
      } catch (err) {
        // Ignore abort errors
        if (err instanceof Error && err.name === "AbortError") return;

        // Don't stop polling on transient network errors - retry next interval
        console.warn("Polling error (will retry):", err);
      }
    },
    [stopPolling],
  );

  /**
   * Start segmentation job and begin polling
   *
   * @param caseId - The case ID to process
   * @param fastMode - Whether to use fast inference mode
   */
  const runSegmentation = useCallback(
    async (caseId: string, fastMode = true) => {
      // Cancel any existing job/polling
      stopPolling();
      abortControllerRef.current?.abort();

      const abortController = new AbortController();
      abortControllerRef.current = abortController;

      // Reset state
      setError(null);
      setResult(null);
      setProgress(0);
      setProgressMessage("Creating job...");
      setJobStatus("pending");
      setElapsedSeconds(undefined);
      setIsLoading(true);

      // Retry loop for cold start handling (replaces recursive call)
      let retryCount = 0;
      while (retryCount <= MAX_COLD_START_RETRIES) {
        try {
          // Create the job
          const response = await apiClient.createSegmentJob(
            caseId,
            fastMode,
            abortController.signal,
          );

          // Store job reference
          const newJobId = response.jobId;
          setJobId(newJobId);
          currentJobRef.current = newJobId;
          setJobStatus(response.status);
          setProgressMessage(response.message);

          // Start polling
          pollingIntervalRef.current = window.setInterval(() => {
            pollJobStatus(newJobId, abortController.signal);
          }, POLLING_INTERVAL);

          // Do an initial poll immediately
          await pollJobStatus(newJobId, abortController.signal);

          // Success - exit retry loop
          return;
        } catch (err) {
          // Ignore abort errors
          if (err instanceof Error && err.name === "AbortError") return;

          // Detect cold start (503 Service Unavailable or network failure)
          const is503 = err instanceof ApiError && err.status === 503;
          const isNetworkError =
            err instanceof TypeError &&
            err.message.toLowerCase().includes("fetch");

          // Retry on cold start errors with exponential backoff
          if (
            (is503 || isNetworkError) &&
            retryCount < MAX_COLD_START_RETRIES
          ) {
            retryCount++;
            setJobStatus("waking_up");
            setProgressMessage(
              `Backend is waking up... Please wait (~30-60s). Retry ${retryCount}/${MAX_COLD_START_RETRIES}`,
            );
            setProgress(0);

            // Exponential backoff: 2s, 4s, 8s, 16s, 30s (capped)
            const delay = Math.min(
              INITIAL_RETRY_DELAY * Math.pow(2, retryCount - 1),
              MAX_RETRY_DELAY,
            );
            await sleep(delay);

            // Continue to next iteration of retry loop
            continue;
          }

          // Max retries exceeded or non-retryable error
          const message =
            is503 || isNetworkError
              ? "Backend failed to wake up. Please try again later."
              : err instanceof Error
                ? err.message
                : "Failed to start job";
          setError(message);
          setIsLoading(false);
          setJobStatus("failed");
          return;
        }
      }
    },
    [pollJobStatus, stopPolling],
  );

  /**
   * Cancel the current job (stops polling, clears loading state)
   */
  const cancelJob = useCallback(() => {
    stopPolling();
    abortControllerRef.current?.abort();
    currentJobRef.current = null;
    setIsLoading(false);
    setJobStatus(null);
    setProgress(0);
    setProgressMessage("");
  }, [stopPolling]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      stopPolling();
      abortControllerRef.current?.abort();
    };
  }, [stopPolling]);

  return {
    // Result data
    result,
    error,

    // Job status
    jobId,
    jobStatus,
    progress,
    progressMessage,
    elapsedSeconds,

    // Loading state
    isLoading,

    // Actions
    runSegmentation,
    cancelJob,
  };
}