src/hooks/useSegmentation.ts

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

// Polling interval in milliseconds
const POLLING_INTERVAL = 2000

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

        // 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
   */
  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)

      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

        const message = 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,
  }
}