geointel/frontend/src/hooks/useDetectionWorkflow.ts

import { useState } from 'react'
import { detectionApi } from '../services/api'
import type {
  DatasetCreateResponse,
  DetectionModelCapability,
  DetectionQaResult,
  DetectionRead,
  DetectionRunRead,
  DetectionRunResponse,
  QualityCheckRead,
} from '../types'
import { formatError } from '../lib/formatError'

interface DetectionWorkflowOptions {
  selectedProjectId: string | null
  rasterDatasets: DatasetCreateResponse[]
  qaIouThreshold: number
  loadProjectData: (projectId: string) => Promise<unknown>
  loadQualityChecks: (projectId?: string | null) => Promise<QualityCheckRead[] | void>
}

export function useDetectionWorkflow({
  selectedProjectId,
  rasterDatasets,
  qaIouThreshold,
  loadProjectData,
  loadQualityChecks,
}: DetectionWorkflowOptions) {
  const [detectionModels, setDetectionModels] = useState<DetectionModelCapability[]>([])
  const [loadingDetectionModels, setLoadingDetectionModels] = useState(false)
  const [detectionModelError, setDetectionModelError] = useState<string | null>(null)
  const [selectedDetectionDatasetId, setSelectedDetectionDatasetId] = useState('')
  const [selectedDetectionModelId, setSelectedDetectionModelId] = useState('yolo-placeholder')
  const [detectionTileManifestPath, setDetectionTileManifestPath] = useState('')
  const [detectionConfidenceThreshold, setDetectionConfidenceThreshold] = useState(0.5)
  const [runningDetection, setRunningDetection] = useState(false)
  const [detectionRunResult, setDetectionRunResult] = useState<DetectionRunResponse | null>(null)
  const [detectionRunError, setDetectionRunError] = useState<string | null>(null)
  const [detectionRuns, setDetectionRuns] = useState<DetectionRunRead[]>([])
  const [selectedDetectionRunId, setSelectedDetectionRunId] = useState('')
  const [detectionItems, setDetectionItems] = useState<DetectionRead[]>([])
  const [detectionGeoJson, setDetectionGeoJson] = useState<GeoJSON.FeatureCollection | null>(null)
  const [detectionClassFilter, setDetectionClassFilter] = useState('')
  const [detectionMinConfidenceFilter, setDetectionMinConfidenceFilter] = useState(0)
  const [loadingDetectionResults, setLoadingDetectionResults] = useState(false)
  const [detectionReferenceDatasetId, setDetectionReferenceDatasetId] = useState('')
  const [detectionQaResult, setDetectionQaResult] = useState<DetectionQaResult | null>(null)
  const [detectionQaError, setDetectionQaError] = useState<string | null>(null)
  const [runningDetectionQa, setRunningDetectionQa] = useState(false)

  const loadDetectionModels = async () => {
    setLoadingDetectionModels(true)
    setDetectionModelError(null)
    try {
      const response = await detectionApi.listModels()
      setDetectionModels(response.models)
      if (!response.models.some((model) => model.model_id === selectedDetectionModelId) && response.models.length > 0) {
        setSelectedDetectionModelId(response.models[0].model_id)
      }
    } catch (error) {
      setDetectionModelError(formatError(error, 'Failed to load detection models'))
    } finally {
      setLoadingDetectionModels(false)
    }
  }

  const loadDetectionRuns = async (projectId = selectedProjectId) => {
    if (!projectId) {
      setDetectionRuns([])
      return
    }
    try {
      const response = await detectionApi.listRuns({ project_id: projectId })
      setDetectionRuns(response.items)
      if (!selectedDetectionRunId && response.items.length > 0) {
        setSelectedDetectionRunId(response.items[0].id)
      }
    } catch (error) {
      setDetectionRunError(formatError(error, 'Failed to load detection runs'))
    }
  }

  const loadDetectionResults = async (analysisRunId = selectedDetectionRunId) => {
    if (!analysisRunId) {
      setDetectionItems([])
      setDetectionGeoJson(null)
      return
    }
    setLoadingDetectionResults(true)
    setDetectionRunError(null)
    try {
      const params = {
        class_name: detectionClassFilter || null,
        min_confidence: detectionMinConfidenceFilter > 0 ? detectionMinConfidenceFilter : null,
      }
      const [detectionsResponse, geoJsonResponse] = await Promise.all([
        detectionApi.listDetections(analysisRunId, params),
        detectionApi.getRunGeoJson(analysisRunId, params),
      ])
      setDetectionItems(detectionsResponse.items)
      setDetectionGeoJson(geoJsonResponse)
    } catch (error) {
      setDetectionRunError(formatError(error, 'Failed to load detection results'))
    } finally {
      setLoadingDetectionResults(false)
    }
  }

  const runDetection = async () => {
    if (!selectedProjectId) {
      setDetectionRunError('Select a project first')
      return
    }
    const datasetId = selectedDetectionDatasetId || rasterDatasets[0]?.id
    if (!datasetId) {
      setDetectionRunError('Select a raster dataset')
      return
    }
    setDetectionRunError(null)
    setDetectionRunResult(null)
    setRunningDetection(true)
    try {
      const result = await detectionApi.run({
        project_id: selectedProjectId,
        dataset_id: datasetId,
        model_id: selectedDetectionModelId,
        confidence_threshold: detectionConfidenceThreshold,
        tile_manifest_path: detectionTileManifestPath.trim() || null,
        parameters_json: {},
      })
      setDetectionRunResult(result)
      setSelectedDetectionRunId(result.analysis_run_id)
      await loadDetectionRuns(selectedProjectId)
      await loadDetectionResults(result.analysis_run_id)
      await loadProjectData(selectedProjectId)
    } catch (error) {
      setDetectionRunError(formatError(error, 'Detection run failed'))
    } finally {
      setRunningDetection(false)
    }
  }

  const runDetectionQa = async () => {
    if (!selectedDetectionRunId) {
      setDetectionQaError('Select a detection run')
      return
    }
    if (!detectionReferenceDatasetId) {
      setDetectionQaError('Select a reference dataset')
      return
    }
    setDetectionQaError(null)
    setDetectionQaResult(null)
    setRunningDetectionQa(true)
    try {
      const result = await detectionApi.compareWithReference(selectedDetectionRunId, {
        reference_dataset_id: detectionReferenceDatasetId,
        iou_threshold: qaIouThreshold,
        class_name: detectionClassFilter || null,
        min_confidence: detectionMinConfidenceFilter > 0 ? detectionMinConfidenceFilter : null,
      })
      setDetectionQaResult(result)
      await loadQualityChecks(selectedProjectId)
    } catch (error) {
      setDetectionQaError(formatError(error, 'Detection QA failed'))
    } finally {
      setRunningDetectionQa(false)
    }
  }

  const resetDetectionForProject = () => {
    setSelectedDetectionDatasetId('')
    setDetectionRuns([])
    setSelectedDetectionRunId('')
    setDetectionItems([])
    setDetectionGeoJson(null)
    setDetectionRunResult(null)
  }

  return {
    detectionModels,
    loadingDetectionModels,
    detectionModelError,
    selectedDetectionDatasetId,
    selectedDetectionModelId,
    detectionTileManifestPath,
    detectionConfidenceThreshold,
    runningDetection,
    detectionRunResult,
    detectionRunError,
    detectionRuns,
    selectedDetectionRunId,
    detectionItems,
    detectionGeoJson,
    detectionClassFilter,
    detectionMinConfidenceFilter,
    loadingDetectionResults,
    detectionReferenceDatasetId,
    detectionQaResult,
    detectionQaError,
    runningDetectionQa,
    loadDetectionModels,
    loadDetectionRuns,
    loadDetectionResults,
    runDetection,
    runDetectionQa,
    resetDetectionForProject,
    setSelectedDetectionDatasetId,
    setSelectedDetectionModelId,
    setDetectionTileManifestPath,
    setDetectionConfidenceThreshold,
    setSelectedDetectionRunId,
    setDetectionClassFilter,
    setDetectionMinConfidenceFilter,
    setDetectionReferenceDatasetId,
  }
}