ui/src/features/annotations/CanvasEditor.tsx

import { useRef, useEffect, useState, useCallback, forwardRef, useImperativeHandle } from 'react'
import { MediaType } from '../../types'
import type { Media, AnnotationListItem, Detection, Affiliation, CombatReadiness } from '../../types'
import { getClassColor, getPhotoModeSuffix, getClassNameFallback } from './classColors'

interface Props {
  media: Media
  annotation: AnnotationListItem | null
  detections: Detection[]
  onDetectionsChange: (dets: Detection[]) => void
  selectedClassNum: number
  currentTime: number
  annotations: AnnotationListItem[]
}

export interface CanvasEditorHandle {
  deleteSelected: () => void
  deleteAll: () => void
  hasSelection: () => boolean
}

interface DragState {
  type: 'draw' | 'move' | 'resize'
  startX: number
  startY: number
  detectionIndex?: number
  handle?: string
}

const HANDLE_SIZE = 6
const MIN_BOX_SIZE = 12

const AFFILIATION_COLORS: Record<number, string> = {
  0: '#FFD700',
  1: '#228be6',
  2: '#fa5252',
}

const CanvasEditor = forwardRef<CanvasEditorHandle, Props>(function CanvasEditor(
  { media, annotation, detections, onDetectionsChange, selectedClassNum, currentTime, annotations },
  ref,
) {
  const canvasRef = useRef<HTMLCanvasElement>(null)
  const containerRef = useRef<HTMLDivElement>(null)
  const imgRef = useRef<HTMLImageElement | null>(null)
  const [zoom, setZoom] = useState(1)
  const [pan, setPan] = useState({ x: 0, y: 0 })
  const [selected, setSelected] = useState<Set<number>>(new Set())
  const [dragState, setDragState] = useState<DragState | null>(null)
  const [drawRect, setDrawRect] = useState<{ x: number; y: number; w: number; h: number } | null>(null)
  const [imgSize, setImgSize] = useState({ w: 0, h: 0 })

  useImperativeHandle(ref, () => ({
    deleteSelected() {
      if (selected.size === 0) return
      onDetectionsChange(detections.filter((_, i) => !selected.has(i)))
      setSelected(new Set())
    },
    deleteAll() {
      onDetectionsChange([])
      setSelected(new Set())
    },
    hasSelection() {
      return selected.size > 0
    },
  }), [selected, detections, onDetectionsChange])

  const isVideo = media.mediaType === MediaType.Video

  const loadImage = useCallback(() => {
    if (isVideo) {
      // Use natural size based on container; no image load
      imgRef.current = null
      return
    }
    const img = new Image()
    img.crossOrigin = 'anonymous'
    if (annotation && !media.path.startsWith('blob:')) {
      img.src = `/api/annotations/annotations/${annotation.id}/image`
    } else if (media.path.startsWith('blob:')) {
      img.src = media.path
    } else {
      img.src = `/api/annotations/media/${media.id}/file`
    }
    img.onload = () => {
      imgRef.current = img
      setImgSize({ w: img.naturalWidth, h: img.naturalHeight })
    }
  }, [media, annotation, isVideo])

  useEffect(() => { loadImage() }, [loadImage])

  useEffect(() => {
    if (!isVideo || !containerRef.current) return
    const update = () => {
      const c = containerRef.current
      if (c) setImgSize({ w: c.clientWidth, h: c.clientHeight })
    }
    update()
    const ro = new ResizeObserver(update)
    ro.observe(containerRef.current)
    return () => ro.disconnect()
  }, [isVideo])

  const toCanvas = useCallback((nx: number, ny: number) => ({
    x: nx * imgSize.w * zoom + pan.x,
    y: ny * imgSize.h * zoom + pan.y,
  }), [imgSize, zoom, pan])

  const fromCanvas = useCallback((cx: number, cy: number) => ({
    x: Math.max(0, Math.min(1, (cx - pan.x) / (imgSize.w * zoom))),
    y: Math.max(0, Math.min(1, (cy - pan.y) / (imgSize.h * zoom))),
  }), [imgSize, zoom, pan])

  const draw = useCallback(() => {
    const canvas = canvasRef.current
    const ctx = canvas?.getContext('2d')
    if (!canvas || !ctx) return
    if (!isVideo && !imgRef.current) return

    const container = containerRef.current
    if (container) {
      canvas.width = container.clientWidth
      canvas.height = container.clientHeight
    }

    ctx.clearRect(0, 0, canvas.width, canvas.height)
    ctx.save()
    if (!isVideo && imgRef.current) {
      ctx.drawImage(imgRef.current, pan.x, pan.y, imgSize.w * zoom, imgSize.h * zoom)
    }

    const timeWindowDets = getTimeWindowDetections()
    const allDets = [...detections, ...timeWindowDets]

    allDets.forEach((det, i) => {
      const isSelected = selected.has(i) && i < detections.length
      const cx = (det.centerX - det.width / 2) * imgSize.w * zoom + pan.x
      const cy = (det.centerY - det.height / 2) * imgSize.h * zoom + pan.y
      const w = det.width * imgSize.w * zoom
      const h = det.height * imgSize.h * zoom

      const color = getClassColor(det.classNum)
      ctx.strokeStyle = color
      ctx.lineWidth = isSelected ? 2 : 1
      ctx.strokeRect(cx, cy, w, h)

      ctx.fillStyle = color
      ctx.globalAlpha = 0.1
      ctx.fillRect(cx, cy, w, h)
      ctx.globalAlpha = 1

      const name = det.label || getClassNameFallback(det.classNum)
      const modeSuffix = getPhotoModeSuffix(det.classNum)
      const confSuffix = det.confidence < 0.995 ? ` ${(det.confidence * 100).toFixed(0)}%` : ''
      const label = `${name}${modeSuffix}${confSuffix}`

      ctx.font = '11px sans-serif'
      const metrics = ctx.measureText(label)
      const padX = 3
      const labelH = 14
      const labelW = metrics.width + padX * 2
      ctx.fillStyle = color
      ctx.fillRect(cx, cy - labelH, labelW, labelH)
      ctx.fillStyle = '#000'
      ctx.fillText(label, cx + padX, cy - 3)

      if (det.combatReadiness === 1) {
        ctx.fillStyle = '#40c057'
        ctx.beginPath()
        ctx.arc(cx + w - 6, cy + 6, 3, 0, Math.PI * 2)
        ctx.fill()
      }

      if (isSelected) {
        const handles = getHandles(cx, cy, w, h)
        handles.forEach(hp => {
          ctx.fillStyle = '#fff'
          ctx.fillRect(hp.x - HANDLE_SIZE / 2, hp.y - HANDLE_SIZE / 2, HANDLE_SIZE, HANDLE_SIZE)
          ctx.strokeStyle = color
          ctx.strokeRect(hp.x - HANDLE_SIZE / 2, hp.y - HANDLE_SIZE / 2, HANDLE_SIZE, HANDLE_SIZE)
        })
      }
    })

    if (drawRect) {
      ctx.strokeStyle = '#fd7e14'
      ctx.lineWidth = 1
      ctx.setLineDash([4, 4])
      ctx.strokeRect(drawRect.x, drawRect.y, drawRect.w, drawRect.h)
      ctx.setLineDash([])
    }

    ctx.restore()
  }, [detections, selected, zoom, pan, imgSize, drawRect, currentTime, annotations])

  useEffect(() => {
    const id = requestAnimationFrame(draw)
    return () => cancelAnimationFrame(id)
  }, [draw])

  useEffect(() => {
    const container = containerRef.current
    if (!container) return
    const obs = new ResizeObserver(() => draw())
    obs.observe(container)
    return () => obs.disconnect()
  }, [draw])

  const getTimeWindowDetections = (): Detection[] => {
    if (media.mediaType !== MediaType.Video) return []
    if (annotation) return []
    const timeTicks = currentTime * 10_000_000
    return annotations
      .filter(a => {
        if (!a.time) return false
        const parts = a.time.split(':').map(Number)
        const annTime = (parts[0] * 3600 + parts[1] * 60 + parts[2]) * 10_000_000
        return Math.abs(annTime - timeTicks) < 2_000_000
      })
      .flatMap(a => a.detections)
  }

  const getHandles = (x: number, y: number, w: number, h: number) => [
    { x, y, cursor: 'nw-resize', name: 'tl' },
    { x: x + w / 2, y, cursor: 'n-resize', name: 'tc' },
    { x: x + w, y, cursor: 'ne-resize', name: 'tr' },
    { x: x + w, y: y + h / 2, cursor: 'e-resize', name: 'mr' },
    { x: x + w, y: y + h, cursor: 'se-resize', name: 'br' },
    { x: x + w / 2, y: y + h, cursor: 's-resize', name: 'bc' },
    { x, y: y + h, cursor: 'sw-resize', name: 'bl' },
    { x, y: y + h / 2, cursor: 'w-resize', name: 'ml' },
  ]

  const hitTest = (cx: number, cy: number) => {
    for (let i = detections.length - 1; i >= 0; i--) {
      const d = detections[i]
      const bx = (d.centerX - d.width / 2) * imgSize.w * zoom + pan.x
      const by = (d.centerY - d.height / 2) * imgSize.h * zoom + pan.y
      const bw = d.width * imgSize.w * zoom
      const bh = d.height * imgSize.h * zoom

      if (selected.has(i)) {
        const handles = getHandles(bx, by, bw, bh)
        for (const h of handles) {
          if (Math.abs(cx - h.x) < HANDLE_SIZE && Math.abs(cy - h.y) < HANDLE_SIZE) {
            return { type: 'handle' as const, index: i, handle: h.name }
          }
        }
      }

      if (cx >= bx && cx <= bx + bw && cy >= by && cy <= by + bh) {
        return { type: 'box' as const, index: i }
      }
    }
    return null
  }

  const handleMouseDown = (e: React.MouseEvent) => {
    const rect = canvasRef.current?.getBoundingClientRect()
    if (!rect) return
    const mx = e.clientX - rect.left
    const my = e.clientY - rect.top

    if (e.ctrlKey && e.button === 0) {
      setDragState({ type: 'draw', startX: mx, startY: my })
      return
    }

    const hit = hitTest(mx, my)
    if (hit?.type === 'handle') {
      setDragState({ type: 'resize', startX: mx, startY: my, detectionIndex: hit.index, handle: hit.handle })
    } else if (hit?.type === 'box') {
      if (e.ctrlKey) {
        setSelected(prev => { const n = new Set(prev); n.has(hit.index) ? n.delete(hit.index) : n.add(hit.index); return n })
      } else {
        setSelected(new Set([hit.index]))
      }
      setDragState({ type: 'move', startX: mx, startY: my, detectionIndex: hit.index })
    } else {
      setSelected(new Set())
      setDragState({ type: 'draw', startX: mx, startY: my })
    }
  }

  const handleMouseMove = (e: React.MouseEvent) => {
    if (!dragState) return
    const rect = canvasRef.current?.getBoundingClientRect()
    if (!rect) return
    const mx = e.clientX - rect.left
    const my = e.clientY - rect.top

    if (dragState.type === 'draw') {
      setDrawRect({
        x: Math.min(dragState.startX, mx),
        y: Math.min(dragState.startY, my),
        w: Math.abs(mx - dragState.startX),
        h: Math.abs(my - dragState.startY),
      })
    } else if (dragState.type === 'move' && dragState.detectionIndex !== undefined) {
      const dx = (mx - dragState.startX) / (imgSize.w * zoom)
      const dy = (my - dragState.startY) / (imgSize.h * zoom)
      const newDets = [...detections]
      const indices = selected.size > 0 ? Array.from(selected) : [dragState.detectionIndex]
      indices.forEach(i => {
        if (newDets[i]) {
          newDets[i] = {
            ...newDets[i],
            centerX: Math.max(newDets[i].width / 2, Math.min(1 - newDets[i].width / 2, newDets[i].centerX + dx)),
            centerY: Math.max(newDets[i].height / 2, Math.min(1 - newDets[i].height / 2, newDets[i].centerY + dy)),
          }
        }
      })
      onDetectionsChange(newDets)
      setDragState({ ...dragState, startX: mx, startY: my })
    } else if (dragState.type === 'resize' && dragState.detectionIndex !== undefined && dragState.handle) {
      const idx = dragState.detectionIndex
      const d = detections[idx]
      const norm = fromCanvas(mx, my)
      const newDets = [...detections]
      let x1 = d.centerX - d.width / 2, y1 = d.centerY - d.height / 2
      let x2 = d.centerX + d.width / 2, y2 = d.centerY + d.height / 2

      if (dragState.handle.includes('l')) x1 = norm.x
      if (dragState.handle.includes('r')) x2 = norm.x
      if (dragState.handle.includes('t')) y1 = norm.y
      if (dragState.handle.includes('b')) y2 = norm.y

      const w = Math.abs(x2 - x1), h = Math.abs(y2 - y1)
      if (w * imgSize.w * zoom >= MIN_BOX_SIZE && h * imgSize.h * zoom >= MIN_BOX_SIZE) {
        newDets[idx] = {
          ...d,
          centerX: Math.min(x1, x2) + w / 2,
          centerY: Math.min(y1, y2) + h / 2,
          width: w,
          height: h,
        }
        onDetectionsChange(newDets)
      }
    }
  }

  const handleMouseUp = () => {
    if (dragState?.type === 'draw' && drawRect) {
      const w = drawRect.w / (imgSize.w * zoom)
      const h = drawRect.h / (imgSize.h * zoom)
      if (w * imgSize.w >= MIN_BOX_SIZE && h * imgSize.h >= MIN_BOX_SIZE) {
        const center = fromCanvas(drawRect.x + drawRect.w / 2, drawRect.y + drawRect.h / 2)
        const newDet: Detection = {
          id: crypto.randomUUID(),
          classNum: selectedClassNum,
          label: '',
          confidence: 1,
          affiliation: 0 as Affiliation,
          combatReadiness: 0 as CombatReadiness,
          centerX: center.x,
          centerY: center.y,
          width: w,
          height: h,
        }
        onDetectionsChange([...detections, newDet])
        setSelected(new Set([detections.length]))
      }
      setDrawRect(null)
    }
    setDragState(null)
  }

  const handleWheel = (e: React.WheelEvent) => {
    if (!e.ctrlKey) return
    e.preventDefault()
    const delta = e.deltaY > 0 ? 0.9 : 1.1
    setZoom(z => Math.max(0.1, Math.min(10, z * delta)))
  }

  useEffect(() => {
    const handler = (e: KeyboardEvent) => {
      if (e.target instanceof HTMLInputElement) return
      if (e.key === 'Delete' && selected.size > 0) {
        onDetectionsChange(detections.filter((_, i) => !selected.has(i)))
        setSelected(new Set())
      }
      if (e.key === 'x' || e.key === 'X') {
        if (e.target instanceof HTMLInputElement) return
        onDetectionsChange([])
        setSelected(new Set())
      }
    }
    window.addEventListener('keydown', handler)
    return () => window.removeEventListener('keydown', handler)
  }, [detections, selected, onDetectionsChange])

  return (
    <div ref={containerRef} className="w-full h-full flex-1 relative overflow-hidden cursor-crosshair">
      <canvas
        ref={canvasRef}
        className="absolute inset-0"
        onMouseDown={handleMouseDown}
        onMouseMove={handleMouseMove}
        onMouseUp={handleMouseUp}
        onMouseLeave={handleMouseUp}
        onWheel={handleWheel}
      />
    </div>
  )
})

export default CanvasEditor