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
   * @param retryCount - Internal retry counter for cold start handling (do not set manually)
   */
  const runSegmentation = useCallback(
    async (caseId: string, fastMode = true, retryCount = 0) => {
      // Only reset state on first attempt (not retries)
      if (retryCount === 0) {
        // 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)
      }

      const abortController = abortControllerRef.current
      if (!abortController) return

      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)
      } 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) {
          setJobStatus('waking_up')
          setProgressMessage(
            `Backend is waking up... Please wait (~30-60s). Retry ${retryCount + 1}/${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),
            MAX_RETRY_DELAY
          )
          await sleep(delay)

          // Recursive retry
          return runSegmentation(caseId, fastMode, retryCount + 1)
        }

        // 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')
      }
    },
    [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,
  }
}