breakpilot-lehrer/admin-lehrer/components/ocr-pipeline/StepLlmReview.tsx

'use client'

import { useCallback, useEffect, useRef, useState } from 'react'
import type { GridCell, GridResult, WordEntry, ColumnMeta } from '@/app/(admin)/ai/ocr-kombi/types'
import { usePixelWordPositions } from './usePixelWordPositions'
import type { LlmChange, StepLlmReviewProps, ReviewMeta, StreamProgress, RowStatus } from './llm-review-types'
import { COL_TYPE_TO_FIELD, KLAUSUR_API } from './llm-review-types'
import { LoadingScreen, ErrorScreen, AppliedScreen, NoSessionScreen } from './LlmReviewStatusScreens'
import { LlmReviewVocabTable } from './LlmReviewVocabTable'
import { LlmReviewOverlay } from './LlmReviewOverlay'
import { LlmReviewCorrections } from './LlmReviewCorrections'

export function StepLlmReview({ sessionId, onNext }: StepLlmReviewProps) {
  // Core state
  const [status, setStatus] = useState<'idle' | 'loading' | 'ready' | 'running' | 'done' | 'error' | 'applied'>('idle')
  const [meta, setMeta] = useState<ReviewMeta | null>(null)
  const [changes, setChanges] = useState<LlmChange[]>([])
  const [progress, setProgress] = useState<StreamProgress | null>(null)
  const [totalDuration, setTotalDuration] = useState(0)
  const [error, setError] = useState('')
  const [accepted, setAccepted] = useState<Set<number>>(new Set())
  const [applying, setApplying] = useState(false)

  // Full vocab table state
  const [vocabEntries, setVocabEntries] = useState<WordEntry[]>([])
  const [columnsUsed, setColumnsUsed] = useState<ColumnMeta[]>([])
  const [activeRowIndices, setActiveRowIndices] = useState<Set<number>>(new Set())
  const [reviewedRows, setReviewedRows] = useState<Set<number>>(new Set())
  const [skippedRows, setSkippedRows] = useState<Set<number>>(new Set())
  const [correctedMap, setCorrectedMap] = useState<Map<number, LlmChange[]>>(new Map())

  // Image
  const [imageNaturalSize, setImageNaturalSize] = useState<{ w: number; h: number } | null>(null)

  // Overlay view state
  const [viewMode, setViewMode] = useState<'table' | 'overlay'>('table')
  const [fontScale, setFontScale] = useState(0.7)
  const [leftPaddingPct, setLeftPaddingPct] = useState(0)
  const [globalBold, setGlobalBold] = useState(false)
  const [cells, setCells] = useState<GridCell[]>([])

  // Pixel-analysed word positions via shared hook
  const overlayImageUrl = sessionId
    ? `${KLAUSUR_API}/api/v1/ocr-pipeline/sessions/${sessionId}/image/cropped`
    : ''
  const cellWordPositions = usePixelWordPositions(overlayImageUrl, cells, viewMode === 'overlay')

  const tableRef = useRef<HTMLDivElement>(null)
  const activeRowRef = useRef<HTMLTableRowElement>(null)

  // Load session data on mount
  useEffect(() => {
    if (!sessionId) return
    loadSessionData()
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [sessionId])

  const loadSessionData = async () => {
    if (!sessionId) return
    setStatus('loading')
    try {
      const res = await fetch(`${KLAUSUR_API}/api/v1/ocr-pipeline/sessions/${sessionId}`)
      if (!res.ok) throw new Error(`HTTP ${res.status}`)
      const data = await res.json()

      const wordResult: GridResult | undefined = data.word_result
      if (!wordResult) {
        setError('Keine Worterkennungsdaten gefunden. Bitte zuerst Schritt 5 abschliessen.')
        setStatus('error')
        return
      }

      const entries = wordResult.vocab_entries || wordResult.entries || []
      setVocabEntries(entries)
      setColumnsUsed(wordResult.columns_used || [])
      setCells(wordResult.cells || [])

      // Check if LLM review was already run
      const llmReview = wordResult.llm_review
      if (llmReview && llmReview.changes) {
        const existingChanges: LlmChange[] = llmReview.changes as LlmChange[]
        setChanges(existingChanges)
        setTotalDuration(llmReview.duration_ms || 0)

        // Mark all rows as reviewed
        const allReviewed = new Set(entries.map((_: WordEntry, i: number) => i))
        setReviewedRows(allReviewed)

        // Build corrected map
        const cMap = new Map<number, LlmChange[]>()
        for (const c of existingChanges) {
          const existing = cMap.get(c.row_index) || []
          existing.push(c)
          cMap.set(c.row_index, existing)
        }
        setCorrectedMap(cMap)

        // Default: all accepted
        setAccepted(new Set(existingChanges.map((_: LlmChange, i: number) => i)))

        setMeta({
          total_entries: entries.length,
          to_review: llmReview.entries_corrected !== undefined ? entries.length : entries.length,
          skipped: 0,
          model: llmReview.model_used || 'unknown',
        })
        setStatus('done')
      } else {
        setStatus('ready')
      }
    } catch (e: unknown) {
      setError(e instanceof Error ? e.message : String(e))
      setStatus('error')
    }
  }

  const runReview = useCallback(async () => {
    if (!sessionId) return
    setStatus('running')
    setError('')
    setChanges([])
    setProgress(null)
    setMeta(null)
    setTotalDuration(0)
    setActiveRowIndices(new Set())
    setReviewedRows(new Set())
    setSkippedRows(new Set())
    setCorrectedMap(new Map())

    try {
      const res = await fetch(
        `${KLAUSUR_API}/api/v1/ocr-pipeline/sessions/${sessionId}/llm-review?stream=true`,
        { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({}) },
      )

      if (!res.ok) {
        const data = await res.json().catch(() => ({}))
        throw new Error(data.detail || `HTTP ${res.status}`)
      }

      const reader = res.body!.getReader()
      const decoder = new TextDecoder()
      let buffer = ''
      let allChanges: LlmChange[] = []
      let allReviewed = new Set<number>()
      let allSkipped = new Set<number>()
      let cMap = new Map<number, LlmChange[]>()

      while (true) {
        const { done, value } = await reader.read()
        if (done) break
        buffer += decoder.decode(value, { stream: true })

        while (buffer.includes('\n\n')) {
          const idx = buffer.indexOf('\n\n')
          const chunk = buffer.slice(0, idx).trim()
          buffer = buffer.slice(idx + 2)

          if (!chunk.startsWith('data: ')) continue
          const dataStr = chunk.slice(6)

          let event: any
          try { event = JSON.parse(dataStr) } catch { continue }

          if (event.type === 'meta') {
            setMeta({
              total_entries: event.total_entries,
              to_review: event.to_review,
              skipped: event.skipped,
              model: event.model,
              skipped_indices: event.skipped_indices,
            })
            if (event.skipped_indices) {
              allSkipped = new Set(event.skipped_indices)
              setSkippedRows(allSkipped)
            }
          }

          if (event.type === 'batch') {
            const batchChanges: LlmChange[] = event.changes || []
            const batchRows: number[] = event.entries_reviewed || []

            setActiveRowIndices(new Set(batchRows))

            allChanges = [...allChanges, ...batchChanges]
            setChanges(allChanges)
            setProgress(event.progress)

            for (const c of batchChanges) {
              const existing = cMap.get(c.row_index) || []
              existing.push(c)
              cMap.set(c.row_index, [...existing])
            }
            setCorrectedMap(new Map(cMap))

            for (const r of batchRows) {
              allReviewed.add(r)
            }
            setReviewedRows(new Set(allReviewed))

            setTimeout(() => {
              activeRowRef.current?.scrollIntoView({ behavior: 'smooth', block: 'center' })
            }, 50)
          }

          if (event.type === 'complete') {
            setActiveRowIndices(new Set())
            setTotalDuration(event.duration_ms)
            setAccepted(new Set(allChanges.map((_: LlmChange, i: number) => i)))
            const allEntryIndices = vocabEntries.map((_: WordEntry, i: number) => i)
            for (const i of allEntryIndices) {
              if (!allSkipped.has(i)) allReviewed.add(i)
            }
            setReviewedRows(new Set(allReviewed))
            setStatus('done')
          }

          if (event.type === 'error') {
            throw new Error(event.detail || 'Unbekannter Fehler')
          }
        }
      }

      if (allChanges.length === 0) {
        setStatus('done')
      }
    } catch (e: unknown) {
      const msg = e instanceof Error ? e.message : String(e)
      setError(msg)
      setStatus('error')
    }
  }, [sessionId, vocabEntries])

  const toggleChange = (index: number) => {
    setAccepted(prev => {
      const next = new Set(prev)
      if (next.has(index)) next.delete(index)
      else next.add(index)
      return next
    })
  }

  const toggleAll = () => {
    if (accepted.size === changes.length) {
      setAccepted(new Set())
    } else {
      setAccepted(new Set(changes.map((_: LlmChange, i: number) => i)))
    }
  }

  const applyChanges = useCallback(async () => {
    if (!sessionId) return
    setApplying(true)
    try {
      const res = await fetch(`${KLAUSUR_API}/api/v1/ocr-pipeline/sessions/${sessionId}/llm-review/apply`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ accepted_indices: Array.from(accepted) }),
      })
      if (!res.ok) {
        const data = await res.json().catch(() => ({}))
        throw new Error(data.detail || `HTTP ${res.status}`)
      }
      setStatus('applied')
    } catch (e: unknown) {
      setError(e instanceof Error ? e.message : String(e))
    } finally {
      setApplying(false)
    }
  }, [sessionId, accepted])

  const getRowStatus = (rowIndex: number): RowStatus => {
    if (activeRowIndices.has(rowIndex)) return 'active'
    if (skippedRows.has(rowIndex)) return 'skipped'
    if (correctedMap.has(rowIndex)) return 'corrected'
    if (reviewedRows.has(rowIndex)) return 'reviewed'
    return 'pending'
  }

  const dewarpedUrl = sessionId
    ? `${KLAUSUR_API}/api/v1/ocr-pipeline/sessions/${sessionId}/image/cropped`
    : ''

  // --- Early returns for non-main states ---
  if (!sessionId) return <NoSessionScreen />
  if (status === 'loading' || status === 'idle') return <LoadingScreen />
  if (status === 'error') {
    return <ErrorScreen error={error} onRetry={() => { setError(''); loadSessionData() }} onSkip={onNext} />
  }
  if (status === 'applied') {
    return <AppliedScreen acceptedCount={accepted.size} totalChanges={changes.length} onNext={onNext} />
  }

  // Active entry for highlighting on image
  const activeEntry = vocabEntries.find((_: WordEntry, i: number) => activeRowIndices.has(i))
  const pct = progress ? Math.round((progress.current / progress.total) * 100) : 0

  /** Handle inline edit of a cell in the overlay */
  const handleCellEdit = (cellId: string, rowIndex: number, newText: string | null) => {
    if (newText === null) return
    setCells(prev => prev.map(c => c.cell_id === cellId ? { ...c, text: newText } : c))
    const cell = cells.find(c => c.cell_id === cellId)
    if (cell) {
      const field = COL_TYPE_TO_FIELD[cell.col_type]
      if (field) {
        setVocabEntries(prev => prev.map((e, i) =>
          i === rowIndex ? { ...e, [field]: newText } : e
        ))
      }
    }
  }

  // --- Ready / Running / Done: 2-column layout ---
  return (
    <div className="space-y-4">
      {/* Header */}
      <div className="flex items-center justify-between">
        <div>
          <h3 className="text-base font-medium text-gray-700 dark:text-gray-300">
            Schritt 6: Korrektur
          </h3>
          <p className="text-xs text-gray-400 mt-0.5">
            {status === 'ready' && `${vocabEntries.length} Eintraege bereit zur Pruefung`}
            {status === 'running' && meta && `${meta.model} · ${meta.to_review} zu pruefen, ${meta.skipped} uebersprungen`}
            {status === 'done' && (
              <>
                {changes.length} Korrektur{changes.length !== 1 ? 'en' : ''} gefunden
                {meta && <> · {meta.skipped} uebersprungen</>}
                {' '}· {totalDuration}ms · {meta?.model}
              </>
            )}
          </p>
        </div>
        <div className="flex items-center gap-2">
          {status === 'ready' && (
            <button onClick={runReview}
              className="px-5 py-2 bg-teal-600 text-white rounded-lg hover:bg-teal-700 transition-colors text-sm font-medium">
              Korrektur starten
            </button>
          )}
          {status === 'running' && (
            <div className="flex items-center gap-2 text-sm text-teal-600 dark:text-teal-400">
              <div className="animate-spin rounded-full h-4 w-4 border-b-2 border-teal-500" />
              {progress ? `${progress.current}/${progress.total}` : 'Startet...'}
            </div>
          )}
          {status === 'done' && changes.length > 0 && (
            <button onClick={toggleAll}
              className="text-xs px-3 py-1.5 border border-gray-300 dark:border-gray-600 rounded-lg hover:bg-gray-50 dark:hover:bg-gray-700 transition-colors text-gray-600 dark:text-gray-400">
              {accepted.size === changes.length ? 'Keine' : 'Alle'} auswaehlen
            </button>
          )}
        </div>
      </div>

      {/* Progress bar (while running) */}
      {status === 'running' && progress && (
        <div className="space-y-1">
          <div className="flex justify-between text-xs text-gray-400">
            <span>{progress.current} / {progress.total} Eintraege geprueft</span>
            <span>{pct}%</span>
          </div>
          <div className="w-full bg-gray-200 dark:bg-gray-700 rounded-full h-2">
            <div className="bg-teal-500 h-2 rounded-full transition-all duration-500" style={{ width: `${pct}%` }} />
          </div>
        </div>
      )}

      {/* View mode toggle */}
      <div className="flex items-center gap-1">
        <button
          onClick={() => setViewMode('table')}
          className={`px-3 py-1.5 text-xs rounded-l-lg border transition-colors ${
            viewMode === 'table'
              ? 'bg-teal-600 text-white border-teal-600'
              : 'bg-white dark:bg-gray-800 text-gray-600 dark:text-gray-400 border-gray-300 dark:border-gray-600 hover:bg-gray-50 dark:hover:bg-gray-700'
          }`}
        >
          Tabelle
        </button>
        <button
          onClick={() => setViewMode('overlay')}
          className={`px-3 py-1.5 text-xs rounded-r-lg border transition-colors ${
            viewMode === 'overlay'
              ? 'bg-teal-600 text-white border-teal-600'
              : 'bg-white dark:bg-gray-800 text-gray-600 dark:text-gray-400 border-gray-300 dark:border-gray-600 hover:bg-gray-50 dark:hover:bg-gray-700'
          }`}
        >
          Overlay
        </button>
      </div>

      {/* Overlay toolbar */}
      {viewMode === 'overlay' && (
        <div className="flex items-center gap-4 flex-wrap bg-gray-50 dark:bg-gray-800/50 rounded-lg px-3 py-2">
          <label className="flex items-center gap-2 text-xs text-gray-600 dark:text-gray-400">
            Schrift
            <input
              type="range" min={30} max={120} value={Math.round(fontScale * 100)}
              onChange={e => setFontScale(Number(e.target.value) / 100)}
              className="w-24 h-1 accent-teal-600"
            />
            <span className="w-8 text-right font-mono">{Math.round(fontScale * 100)}%</span>
          </label>
          <label className="flex items-center gap-2 text-xs text-gray-600 dark:text-gray-400">
            Einrueckung
            <input
              type="range" min={0} max={20} step={0.5} value={leftPaddingPct}
              onChange={e => setLeftPaddingPct(Number(e.target.value))}
              className="w-24 h-1 accent-teal-600"
            />
            <span className="w-8 text-right font-mono">{leftPaddingPct}%</span>
          </label>
          <button
            onClick={() => setGlobalBold(b => !b)}
            className={`px-2 py-1 text-xs rounded border transition-colors font-bold ${
              globalBold
                ? 'bg-teal-600 text-white border-teal-600'
                : 'bg-white dark:bg-gray-700 text-gray-600 dark:text-gray-400 border-gray-300 dark:border-gray-600'
            }`}
          >
            B
          </button>
        </div>
      )}

      {/* 2-column layout: Image + Table/Overlay */}
      <div className={`grid gap-4 ${viewMode === 'overlay' ? 'grid-cols-2' : 'grid-cols-3'}`}>
        {/* Left: Dewarped Image with highlight overlay */}
        <div className="col-span-1">
          <div className="text-xs font-medium text-gray-500 dark:text-gray-400 mb-1">
            Originalbild
          </div>
          <div className="border rounded-lg overflow-hidden dark:border-gray-700 bg-gray-50 dark:bg-gray-900 relative sticky top-4">
            {/* eslint-disable-next-line @next/next/no-img-element */}
            <img
              src={dewarpedUrl}
              alt="Dewarped"
              className="w-full h-auto"
              onLoad={(e) => {
                const img = e.target as HTMLImageElement
                setImageNaturalSize({ w: img.naturalWidth, h: img.naturalHeight })
              }}
            />
            {/* Highlight overlay for active row */}
            {activeEntry?.bbox && (
              <div
                className="absolute border-2 border-yellow-400 bg-yellow-400/20 pointer-events-none animate-pulse"
                style={{
                  left: `${activeEntry.bbox.x}%`,
                  top: `${activeEntry.bbox.y}%`,
                  width: `${activeEntry.bbox.w}%`,
                  height: `${activeEntry.bbox.h}%`,
                }}
              />
            )}
          </div>
        </div>

        {/* Right: Table or Overlay */}
        <div className={viewMode === 'table' ? 'col-span-2' : 'col-span-1'} ref={tableRef}>
          {viewMode === 'table' ? (
            <LlmReviewVocabTable
              vocabEntries={vocabEntries}
              columnsUsed={columnsUsed}
              getRowStatus={getRowStatus}
              correctedMap={correctedMap}
              activeRowRef={activeRowRef}
            />
          ) : (
            <LlmReviewOverlay
              cells={cells}
              imageNaturalSize={imageNaturalSize}
              fontScale={fontScale}
              leftPaddingPct={leftPaddingPct}
              globalBold={globalBold}
              cellWordPositions={cellWordPositions}
            />
          )}
        </div>
      </div>

      {/* Done state: summary + actions */}
      {status === 'done' && (
        <LlmReviewCorrections
          changes={changes}
          accepted={accepted}
          meta={meta}
          totalDuration={totalDuration}
          applying={applying}
          onToggleChange={toggleChange}
          onToggleAll={toggleAll}
          onApply={applyChanges}
          onNext={onNext}
        />
      )}
    </div>
  )
}