breakpilot-lehrer/klausur-service/backend/ocr_pipeline_api.py

"""
OCR Pipeline API - Schrittweise Seitenrekonstruktion.

Zerlegt den OCR-Prozess in 8 einzelne Schritte:
1. Deskewing - Scan begradigen
2. Dewarping - Buchwoelbung entzerren
3. Spaltenerkennung - Unsichtbare Spalten finden
4. Zeilenerkennung - Horizontale Zeilen + Kopf-/Fusszeilen
5. Worterkennung - OCR mit Bounding Boxes
6. LLM-Korrektur - OCR-Fehler per LLM korrigieren
7. Seitenrekonstruktion - Seite nachbauen
8. Ground Truth Validierung - Gesamtpruefung

Lizenz: Apache 2.0
DATENSCHUTZ: Alle Verarbeitung erfolgt lokal.
"""

import json
import logging
import os
import time
import uuid
from dataclasses import asdict
from datetime import datetime
from typing import Any, Dict, List, Optional

import cv2
import numpy as np
from fastapi import APIRouter, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel

from cv_vocab_pipeline import (
    OLLAMA_REVIEW_MODEL,
    DocumentTypeResult,
    PageRegion,
    RowGeometry,
    _cells_to_vocab_entries,
    _detect_header_footer_gaps,
    _detect_sub_columns,
    _fix_character_confusion,
    _fix_phonetic_brackets,
    analyze_layout,
    analyze_layout_by_words,
    build_cell_grid,
    build_cell_grid_streaming,
    build_cell_grid_v2,
    build_cell_grid_v2_streaming,
    build_word_grid,
    classify_column_types,
    create_layout_image,
    create_ocr_image,
    deskew_image,
    deskew_image_by_word_alignment,
    detect_column_geometry,
    detect_document_type,
    detect_row_geometry,
    expand_narrow_columns,
    _apply_shear,
    dewarp_image,
    dewarp_image_manual,
    llm_review_entries,
    llm_review_entries_streaming,
    render_image_high_res,
    render_pdf_high_res,
)
from ocr_pipeline_session_store import (
    create_session_db,
    delete_session_db,
    get_session_db,
    get_session_image,
    init_ocr_pipeline_tables,
    list_sessions_db,
    update_session_db,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/v1/ocr-pipeline", tags=["ocr-pipeline"])

# ---------------------------------------------------------------------------
# In-memory cache for active sessions (BGR numpy arrays for processing)
# DB is source of truth, cache holds BGR arrays during active processing.
# ---------------------------------------------------------------------------

_cache: Dict[str, Dict[str, Any]] = {}


async def _load_session_to_cache(session_id: str) -> Dict[str, Any]:
    """Load session from DB into cache, decoding PNGs to BGR arrays."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    if session_id in _cache:
        return _cache[session_id]

    cache_entry: Dict[str, Any] = {
        "id": session_id,
        **session,
        "original_bgr": None,
        "deskewed_bgr": None,
        "dewarped_bgr": None,
    }

    # Decode images from DB into BGR numpy arrays
    for img_type, bgr_key in [
        ("original", "original_bgr"),
        ("deskewed", "deskewed_bgr"),
        ("dewarped", "dewarped_bgr"),
    ]:
        png_data = await get_session_image(session_id, img_type)
        if png_data:
            arr = np.frombuffer(png_data, dtype=np.uint8)
            bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
            cache_entry[bgr_key] = bgr

    _cache[session_id] = cache_entry
    return cache_entry


def _get_cached(session_id: str) -> Dict[str, Any]:
    """Get from cache or raise 404."""
    entry = _cache.get(session_id)
    if not entry:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not in cache — reload first")
    return entry


# ---------------------------------------------------------------------------
# Pydantic Models
# ---------------------------------------------------------------------------

class ManualDeskewRequest(BaseModel):
    angle: float


class DeskewGroundTruthRequest(BaseModel):
    is_correct: bool
    corrected_angle: Optional[float] = None
    notes: Optional[str] = None


class ManualDewarpRequest(BaseModel):
    shear_degrees: float


class DewarpGroundTruthRequest(BaseModel):
    is_correct: bool
    corrected_shear: Optional[float] = None
    notes: Optional[str] = None


class RenameSessionRequest(BaseModel):
    name: str


class ManualColumnsRequest(BaseModel):
    columns: List[Dict[str, Any]]


class ColumnGroundTruthRequest(BaseModel):
    is_correct: bool
    corrected_columns: Optional[List[Dict[str, Any]]] = None
    notes: Optional[str] = None


class ManualRowsRequest(BaseModel):
    rows: List[Dict[str, Any]]


class RowGroundTruthRequest(BaseModel):
    is_correct: bool
    corrected_rows: Optional[List[Dict[str, Any]]] = None
    notes: Optional[str] = None


class RemoveHandwritingRequest(BaseModel):
    method: str = "auto"       # "auto" | "telea" | "ns"
    target_ink: str = "all"    # "all" | "colored" | "pencil"
    dilation: int = 2          # mask dilation iterations (0-5)
    use_source: str = "auto"   # "original" | "deskewed" | "auto"


# ---------------------------------------------------------------------------
# Session Management Endpoints
# ---------------------------------------------------------------------------

@router.get("/sessions")
async def list_sessions():
    """List all OCR pipeline sessions."""
    sessions = await list_sessions_db()
    return {"sessions": sessions}


@router.post("/sessions")
async def create_session(
    file: UploadFile = File(...),
    name: Optional[str] = Form(None),
):
    """Upload a PDF or image file and create a pipeline session."""
    file_data = await file.read()
    filename = file.filename or "upload"
    content_type = file.content_type or ""

    session_id = str(uuid.uuid4())
    is_pdf = content_type == "application/pdf" or filename.lower().endswith(".pdf")

    try:
        if is_pdf:
            img_bgr = render_pdf_high_res(file_data, page_number=0, zoom=3.0)
        else:
            img_bgr = render_image_high_res(file_data)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Could not process file: {e}")

    # Encode original as PNG bytes
    success, png_buf = cv2.imencode(".png", img_bgr)
    if not success:
        raise HTTPException(status_code=500, detail="Failed to encode image")

    original_png = png_buf.tobytes()
    session_name = name or filename

    # Persist to DB
    await create_session_db(
        session_id=session_id,
        name=session_name,
        filename=filename,
        original_png=original_png,
    )

    # Cache BGR array for immediate processing
    _cache[session_id] = {
        "id": session_id,
        "filename": filename,
        "name": session_name,
        "original_bgr": img_bgr,
        "deskewed_bgr": None,
        "dewarped_bgr": None,
        "deskew_result": None,
        "dewarp_result": None,
        "ground_truth": {},
        "current_step": 1,
    }

    logger.info(f"OCR Pipeline: created session {session_id} from {filename} "
                f"({img_bgr.shape[1]}x{img_bgr.shape[0]})")

    return {
        "session_id": session_id,
        "filename": filename,
        "name": session_name,
        "image_width": img_bgr.shape[1],
        "image_height": img_bgr.shape[0],
        "original_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/original",
    }


@router.get("/sessions/{session_id}")
async def get_session_info(session_id: str):
    """Get session info including deskew/dewarp/column results for step navigation."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    # Get image dimensions from original PNG
    original_png = await get_session_image(session_id, "original")
    if original_png:
        arr = np.frombuffer(original_png, dtype=np.uint8)
        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
        img_w, img_h = img.shape[1], img.shape[0] if img is not None else (0, 0)
    else:
        img_w, img_h = 0, 0

    result = {
        "session_id": session["id"],
        "filename": session.get("filename", ""),
        "name": session.get("name", ""),
        "image_width": img_w,
        "image_height": img_h,
        "original_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/original",
        "current_step": session.get("current_step", 1),
    }

    if session.get("deskew_result"):
        result["deskew_result"] = session["deskew_result"]
    if session.get("dewarp_result"):
        result["dewarp_result"] = session["dewarp_result"]
    if session.get("column_result"):
        result["column_result"] = session["column_result"]
    if session.get("row_result"):
        result["row_result"] = session["row_result"]
    if session.get("word_result"):
        result["word_result"] = session["word_result"]

    return result


@router.put("/sessions/{session_id}")
async def rename_session(session_id: str, req: RenameSessionRequest):
    """Rename a session."""
    updated = await update_session_db(session_id, name=req.name)
    if not updated:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    return {"session_id": session_id, "name": req.name}


@router.delete("/sessions/{session_id}")
async def delete_session(session_id: str):
    """Delete a session."""
    _cache.pop(session_id, None)
    deleted = await delete_session_db(session_id)
    if not deleted:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    return {"session_id": session_id, "deleted": True}


# ---------------------------------------------------------------------------
# Image Endpoints
# ---------------------------------------------------------------------------

@router.get("/sessions/{session_id}/image/{image_type}")
async def get_image(session_id: str, image_type: str):
    """Serve session images: original, deskewed, dewarped, binarized, columns-overlay, or rows-overlay."""
    valid_types = {"original", "deskewed", "dewarped", "binarized", "columns-overlay", "rows-overlay", "words-overlay", "clean"}
    if image_type not in valid_types:
        raise HTTPException(status_code=400, detail=f"Unknown image type: {image_type}")

    if image_type == "columns-overlay":
        return await _get_columns_overlay(session_id)

    if image_type == "rows-overlay":
        return await _get_rows_overlay(session_id)

    if image_type == "words-overlay":
        return await _get_words_overlay(session_id)

    # Try cache first for fast serving
    cached = _cache.get(session_id)
    if cached:
        png_key = f"{image_type}_png" if image_type != "original" else None
        bgr_key = f"{image_type}_bgr" if image_type != "binarized" else None

        # For binarized, check if we have it cached as PNG
        if image_type == "binarized" and cached.get("binarized_png"):
            return Response(content=cached["binarized_png"], media_type="image/png")

    # Load from DB
    data = await get_session_image(session_id, image_type)
    if not data:
        raise HTTPException(status_code=404, detail=f"Image '{image_type}' not available yet")

    return Response(content=data, media_type="image/png")


# ---------------------------------------------------------------------------
# Deskew Endpoints
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/deskew")
async def auto_deskew(session_id: str):
    """Run both deskew methods and pick the best one."""
    # Ensure session is in cache
    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    img_bgr = cached.get("original_bgr")
    if img_bgr is None:
        raise HTTPException(status_code=400, detail="Original image not available")

    t0 = time.time()

    # Method 1: Hough Lines
    try:
        deskewed_hough, angle_hough = deskew_image(img_bgr.copy())
    except Exception as e:
        logger.warning(f"Hough deskew failed: {e}")
        deskewed_hough, angle_hough = img_bgr, 0.0

    # Method 2: Word Alignment (needs image bytes)
    success_enc, png_orig = cv2.imencode(".png", img_bgr)
    orig_bytes = png_orig.tobytes() if success_enc else b""

    try:
        deskewed_wa_bytes, angle_wa = deskew_image_by_word_alignment(orig_bytes)
    except Exception as e:
        logger.warning(f"Word alignment deskew failed: {e}")
        deskewed_wa_bytes, angle_wa = orig_bytes, 0.0

    duration = time.time() - t0

    # Pick best method
    if abs(angle_wa) >= abs(angle_hough) or abs(angle_hough) < 0.1:
        method_used = "word_alignment"
        angle_applied = angle_wa
        wa_array = np.frombuffer(deskewed_wa_bytes, dtype=np.uint8)
        deskewed_bgr = cv2.imdecode(wa_array, cv2.IMREAD_COLOR)
        if deskewed_bgr is None:
            deskewed_bgr = deskewed_hough
            method_used = "hough"
            angle_applied = angle_hough
    else:
        method_used = "hough"
        angle_applied = angle_hough
        deskewed_bgr = deskewed_hough

    # Encode as PNG
    success, deskewed_png_buf = cv2.imencode(".png", deskewed_bgr)
    deskewed_png = deskewed_png_buf.tobytes() if success else b""

    # Create binarized version
    binarized_png = None
    try:
        binarized = create_ocr_image(deskewed_bgr)
        success_bin, bin_buf = cv2.imencode(".png", binarized)
        binarized_png = bin_buf.tobytes() if success_bin else None
    except Exception as e:
        logger.warning(f"Binarization failed: {e}")

    confidence = max(0.5, 1.0 - abs(angle_applied) / 5.0)

    deskew_result = {
        "angle_hough": round(angle_hough, 3),
        "angle_word_alignment": round(angle_wa, 3),
        "angle_applied": round(angle_applied, 3),
        "method_used": method_used,
        "confidence": round(confidence, 2),
        "duration_seconds": round(duration, 2),
    }

    # Update cache
    cached["deskewed_bgr"] = deskewed_bgr
    cached["binarized_png"] = binarized_png
    cached["deskew_result"] = deskew_result

    # Persist to DB
    db_update = {
        "deskewed_png": deskewed_png,
        "deskew_result": deskew_result,
        "current_step": 2,
    }
    if binarized_png:
        db_update["binarized_png"] = binarized_png
    await update_session_db(session_id, **db_update)

    logger.info(f"OCR Pipeline: deskew session {session_id}: "
                f"hough={angle_hough:.2f} wa={angle_wa:.2f} -> {method_used} {angle_applied:.2f}")

    return {
        "session_id": session_id,
        **deskew_result,
        "deskewed_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/deskewed",
        "binarized_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/binarized",
    }


@router.post("/sessions/{session_id}/deskew/manual")
async def manual_deskew(session_id: str, req: ManualDeskewRequest):
    """Apply a manual rotation angle to the original image."""
    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    img_bgr = cached.get("original_bgr")
    if img_bgr is None:
        raise HTTPException(status_code=400, detail="Original image not available")

    angle = max(-5.0, min(5.0, req.angle))

    h, w = img_bgr.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(img_bgr, M, (w, h),
                              flags=cv2.INTER_LINEAR,
                              borderMode=cv2.BORDER_REPLICATE)

    success, png_buf = cv2.imencode(".png", rotated)
    deskewed_png = png_buf.tobytes() if success else b""

    # Binarize
    binarized_png = None
    try:
        binarized = create_ocr_image(rotated)
        success_bin, bin_buf = cv2.imencode(".png", binarized)
        binarized_png = bin_buf.tobytes() if success_bin else None
    except Exception:
        pass

    deskew_result = {
        **(cached.get("deskew_result") or {}),
        "angle_applied": round(angle, 3),
        "method_used": "manual",
    }

    # Update cache
    cached["deskewed_bgr"] = rotated
    cached["binarized_png"] = binarized_png
    cached["deskew_result"] = deskew_result

    # Persist to DB
    db_update = {
        "deskewed_png": deskewed_png,
        "deskew_result": deskew_result,
    }
    if binarized_png:
        db_update["binarized_png"] = binarized_png
    await update_session_db(session_id, **db_update)

    logger.info(f"OCR Pipeline: manual deskew session {session_id}: {angle:.2f}")

    return {
        "session_id": session_id,
        "angle_applied": round(angle, 3),
        "method_used": "manual",
        "deskewed_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/deskewed",
    }


@router.post("/sessions/{session_id}/ground-truth/deskew")
async def save_deskew_ground_truth(session_id: str, req: DeskewGroundTruthRequest):
    """Save ground truth feedback for the deskew step."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    gt = {
        "is_correct": req.is_correct,
        "corrected_angle": req.corrected_angle,
        "notes": req.notes,
        "saved_at": datetime.utcnow().isoformat(),
        "deskew_result": session.get("deskew_result"),
    }
    ground_truth["deskew"] = gt

    await update_session_db(session_id, ground_truth=ground_truth)

    # Update cache
    if session_id in _cache:
        _cache[session_id]["ground_truth"] = ground_truth

    logger.info(f"OCR Pipeline: ground truth deskew session {session_id}: "
                f"correct={req.is_correct}, corrected_angle={req.corrected_angle}")

    return {"session_id": session_id, "ground_truth": gt}


# ---------------------------------------------------------------------------
# Dewarp Endpoints
# ---------------------------------------------------------------------------

async def _detect_shear_with_vlm(image_bytes: bytes) -> Dict[str, Any]:
    """Ask qwen2.5vl:32b to estimate the vertical shear angle of a scanned page.

    The VLM is shown the image and asked: are the column/table borders tilted?
    If yes, by how many degrees? Returns a dict with shear_degrees and confidence.
    Confidence is 0.0 if Ollama is unavailable or parsing fails.
    """
    import httpx
    import base64
    import re

    ollama_base = os.getenv("OLLAMA_BASE_URL", "http://host.docker.internal:11434")
    model = os.getenv("OLLAMA_HTR_MODEL", "qwen2.5vl:32b")

    prompt = (
        "This is a scanned vocabulary worksheet. Look at the vertical borders of the table columns. "
        "Are they perfectly vertical, or do they tilt slightly? "
        "If they tilt, estimate the tilt angle in degrees (positive = top tilts right, negative = top tilts left). "
        "Reply with ONLY a JSON object like: {\"shear_degrees\": 1.2, \"confidence\": 0.8} "
        "Use confidence 0.0-1.0 based on how clearly you can see the tilt. "
        "If the columns look straight, return {\"shear_degrees\": 0.0, \"confidence\": 0.9}"
    )

    img_b64 = base64.b64encode(image_bytes).decode("utf-8")
    payload = {
        "model": model,
        "prompt": prompt,
        "images": [img_b64],
        "stream": False,
    }

    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(f"{ollama_base}/api/generate", json=payload)
            resp.raise_for_status()
            text = resp.json().get("response", "")

        # Parse JSON from response (may have surrounding text)
        match = re.search(r'\{[^}]+\}', text)
        if match:
            import json
            data = json.loads(match.group(0))
            shear = float(data.get("shear_degrees", 0.0))
            conf = float(data.get("confidence", 0.0))
            # Clamp to reasonable range
            shear = max(-3.0, min(3.0, shear))
            conf = max(0.0, min(1.0, conf))
            return {"method": "vlm_qwen2.5vl", "shear_degrees": round(shear, 3), "confidence": round(conf, 2)}
    except Exception as e:
        logger.warning(f"VLM dewarp failed: {e}")

    return {"method": "vlm_qwen2.5vl", "shear_degrees": 0.0, "confidence": 0.0}


@router.post("/sessions/{session_id}/dewarp")
async def auto_dewarp(
    session_id: str,
    method: str = Query("ensemble", description="Detection method: ensemble | vlm | cv"),
):
    """Detect and correct vertical shear on the deskewed image.

    Methods:
    - **ensemble** (default): 3-method CV ensemble (vertical edges + projection + Hough)
    - **cv**: CV ensemble only (same as ensemble)
    - **vlm**: Ask qwen2.5vl:32b to estimate the shear angle visually
    """
    if method not in ("ensemble", "cv", "vlm"):
        raise HTTPException(status_code=400, detail="method must be one of: ensemble, cv, vlm")

    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    deskewed_bgr = cached.get("deskewed_bgr")
    if deskewed_bgr is None:
        raise HTTPException(status_code=400, detail="Deskew must be completed before dewarp")

    t0 = time.time()

    if method == "vlm":
        # Encode deskewed image to PNG for VLM
        success, png_buf = cv2.imencode(".png", deskewed_bgr)
        img_bytes = png_buf.tobytes() if success else b""
        vlm_det = await _detect_shear_with_vlm(img_bytes)
        shear_deg = vlm_det["shear_degrees"]
        if abs(shear_deg) >= 0.05 and vlm_det["confidence"] >= 0.3:
            dewarped_bgr = _apply_shear(deskewed_bgr, -shear_deg)
        else:
            dewarped_bgr = deskewed_bgr
        dewarp_info = {
            "method": vlm_det["method"],
            "shear_degrees": shear_deg,
            "confidence": vlm_det["confidence"],
            "detections": [vlm_det],
        }
    else:
        dewarped_bgr, dewarp_info = dewarp_image(deskewed_bgr)

    duration = time.time() - t0

    # Encode as PNG
    success, png_buf = cv2.imencode(".png", dewarped_bgr)
    dewarped_png = png_buf.tobytes() if success else b""

    dewarp_result = {
        "method_used": dewarp_info["method"],
        "shear_degrees": dewarp_info["shear_degrees"],
        "confidence": dewarp_info["confidence"],
        "duration_seconds": round(duration, 2),
        "detections": dewarp_info.get("detections", []),
    }

    # Update cache
    cached["dewarped_bgr"] = dewarped_bgr
    cached["dewarp_result"] = dewarp_result

    # Persist to DB
    await update_session_db(
        session_id,
        dewarped_png=dewarped_png,
        dewarp_result=dewarp_result,
        auto_shear_degrees=dewarp_info.get("shear_degrees", 0.0),
        current_step=3,
    )

    logger.info(f"OCR Pipeline: dewarp session {session_id}: "
                f"method={dewarp_info['method']} shear={dewarp_info['shear_degrees']:.3f} "
                f"conf={dewarp_info['confidence']:.2f} ({duration:.2f}s)")

    return {
        "session_id": session_id,
        **dewarp_result,
        "dewarped_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/dewarped",
    }


@router.post("/sessions/{session_id}/dewarp/manual")
async def manual_dewarp(session_id: str, req: ManualDewarpRequest):
    """Apply shear correction with a manual angle."""
    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    deskewed_bgr = cached.get("deskewed_bgr")
    if deskewed_bgr is None:
        raise HTTPException(status_code=400, detail="Deskew must be completed before dewarp")

    shear_deg = max(-2.0, min(2.0, req.shear_degrees))

    if abs(shear_deg) < 0.001:
        dewarped_bgr = deskewed_bgr
    else:
        dewarped_bgr = dewarp_image_manual(deskewed_bgr, shear_deg)

    success, png_buf = cv2.imencode(".png", dewarped_bgr)
    dewarped_png = png_buf.tobytes() if success else b""

    dewarp_result = {
        **(cached.get("dewarp_result") or {}),
        "method_used": "manual",
        "shear_degrees": round(shear_deg, 3),
    }

    # Update cache
    cached["dewarped_bgr"] = dewarped_bgr
    cached["dewarp_result"] = dewarp_result

    # Persist to DB
    await update_session_db(
        session_id,
        dewarped_png=dewarped_png,
        dewarp_result=dewarp_result,
    )

    logger.info(f"OCR Pipeline: manual dewarp session {session_id}: shear={shear_deg:.3f}")

    return {
        "session_id": session_id,
        "shear_degrees": round(shear_deg, 3),
        "method_used": "manual",
        "dewarped_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/dewarped",
    }


@router.post("/sessions/{session_id}/ground-truth/dewarp")
async def save_dewarp_ground_truth(session_id: str, req: DewarpGroundTruthRequest):
    """Save ground truth feedback for the dewarp step."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    gt = {
        "is_correct": req.is_correct,
        "corrected_shear": req.corrected_shear,
        "notes": req.notes,
        "saved_at": datetime.utcnow().isoformat(),
        "dewarp_result": session.get("dewarp_result"),
    }
    ground_truth["dewarp"] = gt

    await update_session_db(session_id, ground_truth=ground_truth)

    if session_id in _cache:
        _cache[session_id]["ground_truth"] = ground_truth

    logger.info(f"OCR Pipeline: ground truth dewarp session {session_id}: "
                f"correct={req.is_correct}, corrected_shear={req.corrected_shear}")

    return {"session_id": session_id, "ground_truth": gt}


# ---------------------------------------------------------------------------
# Document Type Detection (between Dewarp and Columns)
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/detect-type")
async def detect_type(session_id: str):
    """Detect document type (vocab_table, full_text, generic_table).

    Should be called after dewarp (clean image available).
    Stores result in session for frontend to decide pipeline flow.
    """
    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    dewarped_bgr = cached.get("dewarped_bgr")
    if dewarped_bgr is None:
        raise HTTPException(status_code=400, detail="Dewarp must be completed first")

    t0 = time.time()
    ocr_img = create_ocr_image(dewarped_bgr)
    result = detect_document_type(ocr_img, dewarped_bgr)
    duration = time.time() - t0

    result_dict = {
        "doc_type": result.doc_type,
        "confidence": result.confidence,
        "pipeline": result.pipeline,
        "skip_steps": result.skip_steps,
        "features": result.features,
        "duration_seconds": round(duration, 2),
    }

    # Persist to DB
    await update_session_db(
        session_id,
        doc_type=result.doc_type,
        doc_type_result=result_dict,
    )

    cached["doc_type_result"] = result_dict

    logger.info(f"OCR Pipeline: detect-type session {session_id}: "
                f"{result.doc_type} (confidence={result.confidence}, {duration:.2f}s)")

    return {"session_id": session_id, **result_dict}


# ---------------------------------------------------------------------------
# Column Detection Endpoints (Step 3)
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/columns")
async def detect_columns(session_id: str):
    """Run column detection on the dewarped image."""
    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    dewarped_bgr = cached.get("dewarped_bgr")
    if dewarped_bgr is None:
        raise HTTPException(status_code=400, detail="Dewarp must be completed before column detection")

    t0 = time.time()

    # Binarized image for layout analysis
    ocr_img = create_ocr_image(dewarped_bgr)
    h, w = ocr_img.shape[:2]

    # Phase A: Geometry detection (returns word_dicts + inv for reuse)
    geo_result = detect_column_geometry(ocr_img, dewarped_bgr)

    if geo_result is None:
        # Fallback to projection-based layout
        layout_img = create_layout_image(dewarped_bgr)
        regions = analyze_layout(layout_img, ocr_img)
    else:
        geometries, left_x, right_x, top_y, bottom_y, word_dicts, inv = geo_result
        content_w = right_x - left_x

        # Cache intermediates for row detection (avoids second Tesseract run)
        cached["_word_dicts"] = word_dicts
        cached["_inv"] = inv
        cached["_content_bounds"] = (left_x, right_x, top_y, bottom_y)

        # Detect header/footer early so sub-column clustering ignores them
        header_y, footer_y = _detect_header_footer_gaps(inv, w, h) if inv is not None else (None, None)

        # Split sub-columns (e.g. page references) before classification
        geometries = _detect_sub_columns(geometries, content_w, left_x=left_x,
                                          top_y=top_y, header_y=header_y, footer_y=footer_y)

        # Expand narrow columns (sub-columns are often very narrow)
        geometries = expand_narrow_columns(geometries, content_w, left_x, word_dicts)

        # Phase B: Content-based classification
        regions = classify_column_types(geometries, content_w, top_y, w, h, bottom_y,
                                        left_x=left_x, right_x=right_x, inv=inv)

    duration = time.time() - t0

    columns = [asdict(r) for r in regions]

    # Determine classification methods used
    methods = list(set(
        c.get("classification_method", "") for c in columns
        if c.get("classification_method")
    ))

    column_result = {
        "columns": columns,
        "classification_methods": methods,
        "duration_seconds": round(duration, 2),
    }

    # Persist to DB — also invalidate downstream results (rows, words)
    await update_session_db(
        session_id,
        column_result=column_result,
        row_result=None,
        word_result=None,
        current_step=3,
    )

    # Update cache
    cached["column_result"] = column_result
    cached.pop("row_result", None)
    cached.pop("word_result", None)

    col_count = len([c for c in columns if c["type"].startswith("column")])
    logger.info(f"OCR Pipeline: columns session {session_id}: "
                f"{col_count} columns detected ({duration:.2f}s)")

    return {
        "session_id": session_id,
        **column_result,
    }


@router.post("/sessions/{session_id}/columns/manual")
async def set_manual_columns(session_id: str, req: ManualColumnsRequest):
    """Override detected columns with manual definitions."""
    column_result = {
        "columns": req.columns,
        "duration_seconds": 0,
        "method": "manual",
    }

    await update_session_db(session_id, column_result=column_result,
                           row_result=None, word_result=None)

    if session_id in _cache:
        _cache[session_id]["column_result"] = column_result
        _cache[session_id].pop("row_result", None)
        _cache[session_id].pop("word_result", None)

    logger.info(f"OCR Pipeline: manual columns session {session_id}: "
                f"{len(req.columns)} columns set")

    return {"session_id": session_id, **column_result}


@router.post("/sessions/{session_id}/ground-truth/columns")
async def save_column_ground_truth(session_id: str, req: ColumnGroundTruthRequest):
    """Save ground truth feedback for the column detection step."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    gt = {
        "is_correct": req.is_correct,
        "corrected_columns": req.corrected_columns,
        "notes": req.notes,
        "saved_at": datetime.utcnow().isoformat(),
        "column_result": session.get("column_result"),
    }
    ground_truth["columns"] = gt

    await update_session_db(session_id, ground_truth=ground_truth)

    if session_id in _cache:
        _cache[session_id]["ground_truth"] = ground_truth

    return {"session_id": session_id, "ground_truth": gt}


@router.get("/sessions/{session_id}/ground-truth/columns")
async def get_column_ground_truth(session_id: str):
    """Retrieve saved ground truth for column detection, including auto vs GT diff."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    columns_gt = ground_truth.get("columns")
    if not columns_gt:
        raise HTTPException(status_code=404, detail="No column ground truth saved")

    return {
        "session_id": session_id,
        "columns_gt": columns_gt,
        "columns_auto": session.get("column_result"),
    }


async def _get_columns_overlay(session_id: str) -> Response:
    """Generate dewarped image with column borders drawn on it."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    column_result = session.get("column_result")
    if not column_result or not column_result.get("columns"):
        raise HTTPException(status_code=404, detail="No column data available")

    # Load dewarped image
    dewarped_png = await get_session_image(session_id, "dewarped")
    if not dewarped_png:
        raise HTTPException(status_code=404, detail="Dewarped image not available")

    arr = np.frombuffer(dewarped_png, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(status_code=500, detail="Failed to decode image")

    # Color map for region types (BGR)
    colors = {
        "column_en": (255, 180, 0),      # Blue
        "column_de": (0, 200, 0),         # Green
        "column_example": (0, 140, 255),  # Orange
        "column_text": (200, 200, 0),     # Cyan/Turquoise
        "page_ref": (200, 0, 200),        # Purple
        "column_marker": (0, 0, 220),     # Red
        "column_ignore": (180, 180, 180), # Light Gray
        "header": (128, 128, 128),        # Gray
        "footer": (128, 128, 128),        # Gray
        "margin_top": (100, 100, 100),    # Dark Gray
        "margin_bottom": (100, 100, 100), # Dark Gray
    }

    overlay = img.copy()
    for col in column_result["columns"]:
        x, y = col["x"], col["y"]
        w, h = col["width"], col["height"]
        color = colors.get(col.get("type", ""), (200, 200, 200))

        # Semi-transparent fill
        cv2.rectangle(overlay, (x, y), (x + w, y + h), color, -1)

        # Solid border
        cv2.rectangle(img, (x, y), (x + w, y + h), color, 3)

        # Label with confidence
        label = col.get("type", "unknown").replace("column_", "").upper()
        conf = col.get("classification_confidence")
        if conf is not None and conf < 1.0:
            label = f"{label} {int(conf * 100)}%"
        cv2.putText(img, label, (x + 10, y + 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)

    # Blend overlay at 20% opacity
    cv2.addWeighted(overlay, 0.2, img, 0.8, 0, img)

    success, result_png = cv2.imencode(".png", img)
    if not success:
        raise HTTPException(status_code=500, detail="Failed to encode overlay image")

    return Response(content=result_png.tobytes(), media_type="image/png")


# ---------------------------------------------------------------------------
# Row Detection Endpoints
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/rows")
async def detect_rows(session_id: str):
    """Run row detection on the dewarped image using horizontal gap analysis."""
    if session_id not in _cache:
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    dewarped_bgr = cached.get("dewarped_bgr")
    if dewarped_bgr is None:
        raise HTTPException(status_code=400, detail="Dewarp must be completed before row detection")

    t0 = time.time()

    # Try to reuse cached word_dicts and inv from column detection
    word_dicts = cached.get("_word_dicts")
    inv = cached.get("_inv")
    content_bounds = cached.get("_content_bounds")

    if word_dicts is None or inv is None or content_bounds is None:
        # Not cached — run column geometry to get intermediates
        ocr_img = create_ocr_image(dewarped_bgr)
        geo_result = detect_column_geometry(ocr_img, dewarped_bgr)
        if geo_result is None:
            raise HTTPException(status_code=400, detail="Column geometry detection failed — cannot detect rows")
        _geoms, left_x, right_x, top_y, bottom_y, word_dicts, inv = geo_result
        cached["_word_dicts"] = word_dicts
        cached["_inv"] = inv
        cached["_content_bounds"] = (left_x, right_x, top_y, bottom_y)
    else:
        left_x, right_x, top_y, bottom_y = content_bounds

    # Run row detection
    rows = detect_row_geometry(inv, word_dicts, left_x, right_x, top_y, bottom_y)
    duration = time.time() - t0

    # Build serializable result (exclude words to keep payload small)
    rows_data = []
    for r in rows:
        rows_data.append({
            "index": r.index,
            "x": r.x,
            "y": r.y,
            "width": r.width,
            "height": r.height,
            "word_count": r.word_count,
            "row_type": r.row_type,
            "gap_before": r.gap_before,
        })

    type_counts = {}
    for r in rows:
        type_counts[r.row_type] = type_counts.get(r.row_type, 0) + 1

    row_result = {
        "rows": rows_data,
        "summary": type_counts,
        "total_rows": len(rows),
        "duration_seconds": round(duration, 2),
    }

    # Persist to DB — also invalidate word_result since rows changed
    await update_session_db(
        session_id,
        row_result=row_result,
        word_result=None,
        current_step=4,
    )

    cached["row_result"] = row_result
    cached.pop("word_result", None)

    logger.info(f"OCR Pipeline: rows session {session_id}: "
                f"{len(rows)} rows detected ({duration:.2f}s): {type_counts}")

    return {
        "session_id": session_id,
        **row_result,
    }


@router.post("/sessions/{session_id}/rows/manual")
async def set_manual_rows(session_id: str, req: ManualRowsRequest):
    """Override detected rows with manual definitions."""
    row_result = {
        "rows": req.rows,
        "total_rows": len(req.rows),
        "duration_seconds": 0,
        "method": "manual",
    }

    await update_session_db(session_id, row_result=row_result, word_result=None)

    if session_id in _cache:
        _cache[session_id]["row_result"] = row_result
        _cache[session_id].pop("word_result", None)

    logger.info(f"OCR Pipeline: manual rows session {session_id}: "
                f"{len(req.rows)} rows set")

    return {"session_id": session_id, **row_result}


@router.post("/sessions/{session_id}/ground-truth/rows")
async def save_row_ground_truth(session_id: str, req: RowGroundTruthRequest):
    """Save ground truth feedback for the row detection step."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    gt = {
        "is_correct": req.is_correct,
        "corrected_rows": req.corrected_rows,
        "notes": req.notes,
        "saved_at": datetime.utcnow().isoformat(),
        "row_result": session.get("row_result"),
    }
    ground_truth["rows"] = gt

    await update_session_db(session_id, ground_truth=ground_truth)

    if session_id in _cache:
        _cache[session_id]["ground_truth"] = ground_truth

    return {"session_id": session_id, "ground_truth": gt}


@router.get("/sessions/{session_id}/ground-truth/rows")
async def get_row_ground_truth(session_id: str):
    """Retrieve saved ground truth for row detection."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    rows_gt = ground_truth.get("rows")
    if not rows_gt:
        raise HTTPException(status_code=404, detail="No row ground truth saved")

    return {
        "session_id": session_id,
        "rows_gt": rows_gt,
        "rows_auto": session.get("row_result"),
    }


# ---------------------------------------------------------------------------
# Word Recognition Endpoints (Step 5)
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/words")
async def detect_words(
    session_id: str,
    request: Request,
    engine: str = "auto",
    pronunciation: str = "british",
    stream: bool = False,
):
    """Build word grid from columns × rows, OCR each cell.

    Query params:
        engine: 'auto' (default), 'tesseract', or 'rapid'
        pronunciation: 'british' (default) or 'american' — for IPA dictionary lookup
        stream: false (default) for JSON response, true for SSE streaming
    """
    if session_id not in _cache:
        logger.info("detect_words: session %s not in cache, loading from DB", session_id)
        await _load_session_to_cache(session_id)
    cached = _get_cached(session_id)

    dewarped_bgr = cached.get("dewarped_bgr")
    if dewarped_bgr is None:
        logger.warning("detect_words: dewarped_bgr is None for session %s (cache keys: %s)",
                       session_id, [k for k in cached.keys() if k.endswith('_bgr')])
        raise HTTPException(status_code=400, detail="Dewarp must be completed before word detection")

    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    column_result = session.get("column_result")
    row_result = session.get("row_result")
    if not column_result or not column_result.get("columns"):
        raise HTTPException(status_code=400, detail="Column detection must be completed first")
    if not row_result or not row_result.get("rows"):
        raise HTTPException(status_code=400, detail="Row detection must be completed first")

    # Convert column dicts back to PageRegion objects
    col_regions = [
        PageRegion(
            type=c["type"],
            x=c["x"], y=c["y"],
            width=c["width"], height=c["height"],
            classification_confidence=c.get("classification_confidence", 1.0),
            classification_method=c.get("classification_method", ""),
        )
        for c in column_result["columns"]
    ]

    # Convert row dicts back to RowGeometry objects
    row_geoms = [
        RowGeometry(
            index=r["index"],
            x=r["x"], y=r["y"],
            width=r["width"], height=r["height"],
            word_count=r.get("word_count", 0),
            words=[],
            row_type=r.get("row_type", "content"),
            gap_before=r.get("gap_before", 0),
        )
        for r in row_result["rows"]
    ]

    # Cell-First OCR (v2): no full-page word re-population needed.
    # Each cell is cropped and OCR'd in isolation → no neighbour bleeding.
    # We still need word_count > 0 for row filtering in build_cell_grid_v2,
    # so populate from cached words if available (just for counting).
    word_dicts = cached.get("_word_dicts")
    if word_dicts is None:
        ocr_img_tmp = create_ocr_image(dewarped_bgr)
        geo_result = detect_column_geometry(ocr_img_tmp, dewarped_bgr)
        if geo_result is not None:
            _geoms, left_x, right_x, top_y, bottom_y, word_dicts, inv = geo_result
            cached["_word_dicts"] = word_dicts
            cached["_inv"] = inv
            cached["_content_bounds"] = (left_x, right_x, top_y, bottom_y)

    if word_dicts:
        content_bounds = cached.get("_content_bounds")
        if content_bounds:
            _lx, _rx, top_y, _by = content_bounds
        else:
            top_y = min(r.y for r in row_geoms) if row_geoms else 0

        for row in row_geoms:
            row_y_rel = row.y - top_y
            row_bottom_rel = row_y_rel + row.height
            row.words = [
                w for w in word_dicts
                if row_y_rel <= w['top'] + w['height'] / 2 < row_bottom_rel
            ]
            row.word_count = len(row.words)

    if stream:
        # Cell-First OCR v2: use batch-then-stream approach instead of
        # per-cell streaming. The parallel ThreadPoolExecutor in
        # build_cell_grid_v2 is much faster than sequential streaming.
        return StreamingResponse(
            _word_batch_stream_generator(
                session_id, cached, col_regions, row_geoms,
                dewarped_bgr, engine, pronunciation, request,
            ),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    # --- Non-streaming path ---
    t0 = time.time()

    # Create binarized OCR image (for Tesseract)
    ocr_img = create_ocr_image(dewarped_bgr)
    img_h, img_w = dewarped_bgr.shape[:2]

    # Build cell grid using Cell-First OCR (v2) — each cell cropped in isolation
    cells, columns_meta = build_cell_grid_v2(
        ocr_img, col_regions, row_geoms, img_w, img_h,
        ocr_engine=engine, img_bgr=dewarped_bgr,
    )
    duration = time.time() - t0

    # Layout detection
    col_types = {c['type'] for c in columns_meta}
    is_vocab = bool(col_types & {'column_en', 'column_de'})

    # Count content rows and columns for grid_shape
    n_content_rows = len([r for r in row_geoms if r.row_type == 'content'])
    n_cols = len(columns_meta)

    # Determine which engine was actually used
    used_engine = cells[0].get("ocr_engine", "tesseract") if cells else engine

    # Grid result (always generic)
    word_result = {
        "cells": cells,
        "grid_shape": {
            "rows": n_content_rows,
            "cols": n_cols,
            "total_cells": len(cells),
        },
        "columns_used": columns_meta,
        "layout": "vocab" if is_vocab else "generic",
        "image_width": img_w,
        "image_height": img_h,
        "duration_seconds": round(duration, 2),
        "ocr_engine": used_engine,
        "summary": {
            "total_cells": len(cells),
            "non_empty_cells": sum(1 for c in cells if c.get("text")),
            "low_confidence": sum(1 for c in cells if 0 < c.get("confidence", 0) < 50),
        },
    }

    # For vocab layout: map cells 1:1 to vocab entries (row→entry).
    # No content shuffling — each cell stays at its detected position.
    if is_vocab:
        entries = _cells_to_vocab_entries(cells, columns_meta)
        entries = _fix_character_confusion(entries)
        entries = _fix_phonetic_brackets(entries, pronunciation=pronunciation)
        word_result["vocab_entries"] = entries
        word_result["entries"] = entries
        word_result["entry_count"] = len(entries)
        word_result["summary"]["total_entries"] = len(entries)
        word_result["summary"]["with_english"] = sum(1 for e in entries if e.get("english"))
        word_result["summary"]["with_german"] = sum(1 for e in entries if e.get("german"))

    # Persist to DB
    await update_session_db(
        session_id,
        word_result=word_result,
        current_step=5,
    )

    cached["word_result"] = word_result

    logger.info(f"OCR Pipeline: words session {session_id}: "
                f"layout={word_result['layout']}, "
                f"{len(cells)} cells ({duration:.2f}s), summary: {word_result['summary']}")

    return {
        "session_id": session_id,
        **word_result,
    }


async def _word_batch_stream_generator(
    session_id: str,
    cached: Dict[str, Any],
    col_regions: List[PageRegion],
    row_geoms: List[RowGeometry],
    dewarped_bgr: np.ndarray,
    engine: str,
    pronunciation: str,
    request: Request,
):
    """SSE generator that runs batch OCR (parallel) then streams results.

    Unlike the old per-cell streaming, this uses build_cell_grid_v2 with
    ThreadPoolExecutor for parallel OCR, then emits all cells as SSE events.
    The 'preparing' event keeps the connection alive during OCR processing.
    """
    import asyncio

    t0 = time.time()
    ocr_img = create_ocr_image(dewarped_bgr)
    img_h, img_w = dewarped_bgr.shape[:2]

    _skip_types = {'column_ignore', 'header', 'footer', 'margin_top', 'margin_bottom', 'margin_left', 'margin_right'}
    n_content_rows = len([r for r in row_geoms if r.row_type == 'content'])
    n_cols = len([c for c in col_regions if c.type not in _skip_types])
    col_types = {c.type for c in col_regions if c.type not in _skip_types}
    is_vocab = bool(col_types & {'column_en', 'column_de'})
    total_cells = n_content_rows * n_cols

    # 1. Send meta event immediately
    meta_event = {
        "type": "meta",
        "grid_shape": {"rows": n_content_rows, "cols": n_cols, "total_cells": total_cells},
        "layout": "vocab" if is_vocab else "generic",
    }
    yield f"data: {json.dumps(meta_event)}\n\n"

    # 2. Send preparing event (keepalive for proxy)
    yield f"data: {json.dumps({'type': 'preparing', 'message': 'Cell-First OCR laeuft parallel...'})}\n\n"

    # 3. Run batch OCR in thread pool with periodic keepalive events.
    # The OCR takes 30-60s and proxy servers (Nginx) may drop idle SSE
    # connections after 30-60s. Send keepalive every 5s to prevent this.
    loop = asyncio.get_event_loop()
    ocr_future = loop.run_in_executor(
        None,
        lambda: build_cell_grid_v2(
            ocr_img, col_regions, row_geoms, img_w, img_h,
            ocr_engine=engine, img_bgr=dewarped_bgr,
        ),
    )

    # Send keepalive events every 5 seconds while OCR runs
    keepalive_count = 0
    while not ocr_future.done():
        try:
            cells, columns_meta = await asyncio.wait_for(
                asyncio.shield(ocr_future), timeout=5.0,
            )
            break  # OCR finished
        except asyncio.TimeoutError:
            keepalive_count += 1
            elapsed = int(time.time() - t0)
            yield f"data: {json.dumps({'type': 'keepalive', 'elapsed': elapsed, 'message': f'OCR laeuft... ({elapsed}s)'})}\n\n"
            if await request.is_disconnected():
                logger.info(f"SSE batch: client disconnected during OCR for {session_id}")
                ocr_future.cancel()
                return
    else:
        cells, columns_meta = ocr_future.result()

    if await request.is_disconnected():
        logger.info(f"SSE batch: client disconnected after OCR for {session_id}")
        return

    # 4. Send columns meta
    if columns_meta:
        yield f"data: {json.dumps({'type': 'columns', 'columns_used': columns_meta})}\n\n"

    # 5. Stream all cells
    for idx, cell in enumerate(cells):
        cell_event = {
            "type": "cell",
            "cell": cell,
            "progress": {"current": idx + 1, "total": len(cells)},
        }
        yield f"data: {json.dumps(cell_event)}\n\n"

    # 6. Build final result and persist
    duration = time.time() - t0
    used_engine = cells[0].get("ocr_engine", "tesseract") if cells else engine

    word_result = {
        "cells": cells,
        "grid_shape": {"rows": n_content_rows, "cols": n_cols, "total_cells": len(cells)},
        "columns_used": columns_meta,
        "layout": "vocab" if is_vocab else "generic",
        "image_width": img_w,
        "image_height": img_h,
        "duration_seconds": round(duration, 2),
        "ocr_engine": used_engine,
        "summary": {
            "total_cells": len(cells),
            "non_empty_cells": sum(1 for c in cells if c.get("text")),
            "low_confidence": sum(1 for c in cells if 0 < c.get("confidence", 0) < 50),
        },
    }

    vocab_entries = None
    if is_vocab:
        entries = _cells_to_vocab_entries(cells, columns_meta)
        entries = _fix_character_confusion(entries)
        entries = _fix_phonetic_brackets(entries, pronunciation=pronunciation)
        word_result["vocab_entries"] = entries
        word_result["entries"] = entries
        word_result["entry_count"] = len(entries)
        word_result["summary"]["total_entries"] = len(entries)
        word_result["summary"]["with_english"] = sum(1 for e in entries if e.get("english"))
        word_result["summary"]["with_german"] = sum(1 for e in entries if e.get("german"))
        vocab_entries = entries

    await update_session_db(session_id, word_result=word_result, current_step=5)
    cached["word_result"] = word_result

    logger.info(f"OCR Pipeline SSE batch: words session {session_id}: "
                f"layout={word_result['layout']}, {len(cells)} cells ({duration:.2f}s)")

    # 7. Send complete event
    complete_event = {
        "type": "complete",
        "summary": word_result["summary"],
        "duration_seconds": round(duration, 2),
        "ocr_engine": used_engine,
    }
    if vocab_entries is not None:
        complete_event["vocab_entries"] = vocab_entries
    yield f"data: {json.dumps(complete_event)}\n\n"


async def _word_stream_generator(
    session_id: str,
    cached: Dict[str, Any],
    col_regions: List[PageRegion],
    row_geoms: List[RowGeometry],
    dewarped_bgr: np.ndarray,
    engine: str,
    pronunciation: str,
    request: Request,
):
    """SSE generator that yields cell-by-cell OCR progress."""
    t0 = time.time()

    ocr_img = create_ocr_image(dewarped_bgr)
    img_h, img_w = dewarped_bgr.shape[:2]

    # Compute grid shape upfront for the meta event
    n_content_rows = len([r for r in row_geoms if r.row_type == 'content'])
    _skip_types = {'column_ignore', 'header', 'footer', 'margin_top', 'margin_bottom', 'margin_left', 'margin_right'}
    n_cols = len([c for c in col_regions if c.type not in _skip_types])

    # Determine layout
    col_types = {c.type for c in col_regions if c.type not in _skip_types}
    is_vocab = bool(col_types & {'column_en', 'column_de'})

    # Start streaming — first event: meta
    columns_meta = None  # will be set from first yield
    total_cells = n_content_rows * n_cols

    meta_event = {
        "type": "meta",
        "grid_shape": {"rows": n_content_rows, "cols": n_cols, "total_cells": total_cells},
        "layout": "vocab" if is_vocab else "generic",
    }
    yield f"data: {json.dumps(meta_event)}\n\n"

    # Keepalive: send preparing event so proxy doesn't timeout during OCR init
    yield f"data: {json.dumps({'type': 'preparing', 'message': 'Cell-First OCR wird initialisiert...'})}\n\n"

    # Stream cells one by one
    all_cells: List[Dict[str, Any]] = []
    cell_idx = 0
    last_keepalive = time.time()

    for cell, cols_meta, total in build_cell_grid_v2_streaming(
        ocr_img, col_regions, row_geoms, img_w, img_h,
        ocr_engine=engine, img_bgr=dewarped_bgr,
    ):
        if await request.is_disconnected():
            logger.info(f"SSE: client disconnected during streaming for {session_id}")
            return

        if columns_meta is None:
            columns_meta = cols_meta
            # Send columns_used as part of first cell or update meta
            meta_update = {
                "type": "columns",
                "columns_used": cols_meta,
            }
            yield f"data: {json.dumps(meta_update)}\n\n"

        all_cells.append(cell)
        cell_idx += 1

        cell_event = {
            "type": "cell",
            "cell": cell,
            "progress": {"current": cell_idx, "total": total},
        }
        yield f"data: {json.dumps(cell_event)}\n\n"

    # All cells done — build final result
    duration = time.time() - t0
    if columns_meta is None:
        columns_meta = []

    # Post-OCR: remove rows where ALL cells are empty (inter-row gaps
    # that had stray Tesseract artifacts giving word_count > 0).
    rows_with_text: set = set()
    for c in all_cells:
        if c.get("text", "").strip():
            rows_with_text.add(c["row_index"])
    before_filter = len(all_cells)
    all_cells = [c for c in all_cells if c["row_index"] in rows_with_text]
    empty_rows_removed = (before_filter - len(all_cells)) // max(n_cols, 1)
    if empty_rows_removed > 0:
        logger.info(f"SSE: removed {empty_rows_removed} all-empty rows after OCR")

    used_engine = all_cells[0].get("ocr_engine", "tesseract") if all_cells else engine

    word_result = {
        "cells": all_cells,
        "grid_shape": {
            "rows": n_content_rows,
            "cols": n_cols,
            "total_cells": len(all_cells),
        },
        "columns_used": columns_meta,
        "layout": "vocab" if is_vocab else "generic",
        "image_width": img_w,
        "image_height": img_h,
        "duration_seconds": round(duration, 2),
        "ocr_engine": used_engine,
        "summary": {
            "total_cells": len(all_cells),
            "non_empty_cells": sum(1 for c in all_cells if c.get("text")),
            "low_confidence": sum(1 for c in all_cells if 0 < c.get("confidence", 0) < 50),
        },
    }

    # For vocab layout: map cells 1:1 to vocab entries (row→entry).
    # No content shuffling — each cell stays at its detected position.
    vocab_entries = None
    if is_vocab:
        entries = _cells_to_vocab_entries(all_cells, columns_meta)
        entries = _fix_character_confusion(entries)
        entries = _fix_phonetic_brackets(entries, pronunciation=pronunciation)
        word_result["vocab_entries"] = entries
        word_result["entries"] = entries
        word_result["entry_count"] = len(entries)
        word_result["summary"]["total_entries"] = len(entries)
        word_result["summary"]["with_english"] = sum(1 for e in entries if e.get("english"))
        word_result["summary"]["with_german"] = sum(1 for e in entries if e.get("german"))
        vocab_entries = entries

    # Persist to DB
    await update_session_db(
        session_id,
        word_result=word_result,
        current_step=5,
    )
    cached["word_result"] = word_result

    logger.info(f"OCR Pipeline SSE: words session {session_id}: "
                f"layout={word_result['layout']}, "
                f"{len(all_cells)} cells ({duration:.2f}s)")

    # Final complete event
    complete_event = {
        "type": "complete",
        "summary": word_result["summary"],
        "duration_seconds": round(duration, 2),
        "ocr_engine": used_engine,
    }
    if vocab_entries is not None:
        complete_event["vocab_entries"] = vocab_entries
    yield f"data: {json.dumps(complete_event)}\n\n"


class WordGroundTruthRequest(BaseModel):
    is_correct: bool
    corrected_entries: Optional[List[Dict[str, Any]]] = None
    notes: Optional[str] = None


@router.post("/sessions/{session_id}/ground-truth/words")
async def save_word_ground_truth(session_id: str, req: WordGroundTruthRequest):
    """Save ground truth feedback for the word recognition step."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    gt = {
        "is_correct": req.is_correct,
        "corrected_entries": req.corrected_entries,
        "notes": req.notes,
        "saved_at": datetime.utcnow().isoformat(),
        "word_result": session.get("word_result"),
    }
    ground_truth["words"] = gt

    await update_session_db(session_id, ground_truth=ground_truth)

    if session_id in _cache:
        _cache[session_id]["ground_truth"] = ground_truth

    return {"session_id": session_id, "ground_truth": gt}


@router.get("/sessions/{session_id}/ground-truth/words")
async def get_word_ground_truth(session_id: str):
    """Retrieve saved ground truth for word recognition."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    ground_truth = session.get("ground_truth") or {}
    words_gt = ground_truth.get("words")
    if not words_gt:
        raise HTTPException(status_code=404, detail="No word ground truth saved")

    return {
        "session_id": session_id,
        "words_gt": words_gt,
        "words_auto": session.get("word_result"),
    }


# ---------------------------------------------------------------------------
# LLM Review Endpoints (Step 6)
# ---------------------------------------------------------------------------


@router.post("/sessions/{session_id}/llm-review")
async def run_llm_review(session_id: str, request: Request, stream: bool = False):
    """Run LLM-based correction on vocab entries from Step 5.

    Query params:
        stream: false (default) for JSON response, true for SSE streaming
    """
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=400, detail="No word result found — run Step 5 first")

    entries = word_result.get("vocab_entries") or word_result.get("entries") or []
    if not entries:
        raise HTTPException(status_code=400, detail="No vocab entries found — run Step 5 first")

    # Optional model override from request body
    body = {}
    try:
        body = await request.json()
    except Exception:
        pass
    model = body.get("model") or OLLAMA_REVIEW_MODEL

    if stream:
        return StreamingResponse(
            _llm_review_stream_generator(session_id, entries, word_result, model, request),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
        )

    # Non-streaming path
    try:
        result = await llm_review_entries(entries, model=model)
    except Exception as e:
        import traceback
        logger.error(f"LLM review failed for session {session_id}: {type(e).__name__}: {e}\n{traceback.format_exc()}")
        raise HTTPException(status_code=502, detail=f"LLM review failed ({type(e).__name__}): {e}")

    # Store result inside word_result as a sub-key
    word_result["llm_review"] = {
        "changes": result["changes"],
        "model_used": result["model_used"],
        "duration_ms": result["duration_ms"],
        "entries_corrected": result["entries_corrected"],
    }
    await update_session_db(session_id, word_result=word_result, current_step=6)

    if session_id in _cache:
        _cache[session_id]["word_result"] = word_result

    logger.info(f"LLM review session {session_id}: {len(result['changes'])} changes, "
                f"{result['duration_ms']}ms, model={result['model_used']}")

    return {
        "session_id": session_id,
        "changes": result["changes"],
        "model_used": result["model_used"],
        "duration_ms": result["duration_ms"],
        "total_entries": len(entries),
        "corrections_found": len(result["changes"]),
    }


async def _llm_review_stream_generator(
    session_id: str,
    entries: List[Dict],
    word_result: Dict,
    model: str,
    request: Request,
):
    """SSE generator that yields batch-by-batch LLM review progress."""
    try:
        async for event in llm_review_entries_streaming(entries, model=model):
            if await request.is_disconnected():
                logger.info(f"SSE: client disconnected during LLM review for {session_id}")
                return

            yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"

            # On complete: persist to DB
            if event.get("type") == "complete":
                word_result["llm_review"] = {
                    "changes": event["changes"],
                    "model_used": event["model_used"],
                    "duration_ms": event["duration_ms"],
                    "entries_corrected": event["entries_corrected"],
                }
                await update_session_db(session_id, word_result=word_result, current_step=6)
                if session_id in _cache:
                    _cache[session_id]["word_result"] = word_result

                logger.info(f"LLM review SSE session {session_id}: {event['corrections_found']} changes, "
                            f"{event['duration_ms']}ms, skipped={event['skipped']}, model={event['model_used']}")

    except Exception as e:
        import traceback
        logger.error(f"LLM review SSE failed for {session_id}: {type(e).__name__}: {e}\n{traceback.format_exc()}")
        error_event = {"type": "error", "detail": f"{type(e).__name__}: {e}"}
        yield f"data: {json.dumps(error_event)}\n\n"


@router.post("/sessions/{session_id}/llm-review/apply")
async def apply_llm_corrections(session_id: str, request: Request):
    """Apply selected LLM corrections to vocab entries."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=400, detail="No word result found")

    llm_review = word_result.get("llm_review")
    if not llm_review:
        raise HTTPException(status_code=400, detail="No LLM review found — run /llm-review first")

    body = await request.json()
    accepted_indices = set(body.get("accepted_indices", []))  # indices into changes[]

    changes = llm_review.get("changes", [])
    entries = word_result.get("vocab_entries") or word_result.get("entries") or []

    # Build a lookup: (row_index, field) -> new_value for accepted changes
    corrections = {}
    applied_count = 0
    for idx, change in enumerate(changes):
        if idx in accepted_indices:
            key = (change["row_index"], change["field"])
            corrections[key] = change["new"]
            applied_count += 1

    # Apply corrections to entries
    for entry in entries:
        row_idx = entry.get("row_index", -1)
        for field_name in ("english", "german", "example"):
            key = (row_idx, field_name)
            if key in corrections:
                entry[field_name] = corrections[key]
                entry["llm_corrected"] = True

    # Update word_result
    word_result["vocab_entries"] = entries
    word_result["entries"] = entries
    word_result["llm_review"]["applied_count"] = applied_count
    word_result["llm_review"]["applied_at"] = datetime.utcnow().isoformat()

    await update_session_db(session_id, word_result=word_result)

    if session_id in _cache:
        _cache[session_id]["word_result"] = word_result

    logger.info(f"Applied {applied_count}/{len(changes)} LLM corrections for session {session_id}")

    return {
        "session_id": session_id,
        "applied_count": applied_count,
        "total_changes": len(changes),
    }


@router.post("/sessions/{session_id}/reconstruction")
async def save_reconstruction(session_id: str, request: Request):
    """Save edited cell texts from reconstruction step."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=400, detail="No word result found")

    body = await request.json()
    cell_updates = body.get("cells", [])

    if not cell_updates:
        await update_session_db(session_id, current_step=7)
        return {"session_id": session_id, "updated": 0}

    # Build update map: cell_id -> new text
    update_map = {c["cell_id"]: c["text"] for c in cell_updates}

    # Update cells
    cells = word_result.get("cells", [])
    updated_count = 0
    for cell in cells:
        if cell["cell_id"] in update_map:
            cell["text"] = update_map[cell["cell_id"]]
            cell["status"] = "edited"
            updated_count += 1

    word_result["cells"] = cells

    # Also update vocab_entries if present
    entries = word_result.get("vocab_entries") or word_result.get("entries") or []
    if entries:
        # Map cell_id pattern "R{row}_C{col}" to entry fields
        for entry in entries:
            row_idx = entry.get("row_index", -1)
            # Check each field's cell
            for col_idx, field_name in enumerate(["english", "german", "example"]):
                cell_id = f"R{row_idx:02d}_C{col_idx}"
                # Also try without zero-padding
                cell_id_alt = f"R{row_idx}_C{col_idx}"
                new_text = update_map.get(cell_id) or update_map.get(cell_id_alt)
                if new_text is not None:
                    entry[field_name] = new_text

        word_result["vocab_entries"] = entries
        if "entries" in word_result:
            word_result["entries"] = entries

    await update_session_db(session_id, word_result=word_result, current_step=7)

    if session_id in _cache:
        _cache[session_id]["word_result"] = word_result

    logger.info(f"Reconstruction saved for session {session_id}: {updated_count} cells updated")

    return {
        "session_id": session_id,
        "updated": updated_count,
    }


@router.get("/sessions/{session_id}/reconstruction/fabric-json")
async def get_fabric_json(session_id: str):
    """Return cell grid as Fabric.js-compatible JSON for the canvas editor."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=400, detail="No word result found")

    cells = word_result.get("cells", [])
    img_w = word_result.get("image_width", 800)
    img_h = word_result.get("image_height", 600)

    from services.layout_reconstruction_service import cells_to_fabric_json
    fabric_json = cells_to_fabric_json(cells, img_w, img_h)

    return fabric_json


@router.get("/sessions/{session_id}/reconstruction/export/pdf")
async def export_reconstruction_pdf(session_id: str):
    """Export the reconstructed cell grid as a PDF table."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=400, detail="No word result found")

    cells = word_result.get("cells", [])
    columns_used = word_result.get("columns_used", [])
    grid_shape = word_result.get("grid_shape", {})
    n_rows = grid_shape.get("rows", 0)
    n_cols = grid_shape.get("cols", 0)

    # Build table data: rows × columns
    table_data: list[list[str]] = []
    header = [c.get("label", c.get("type", f"Col {i}")) for i, c in enumerate(columns_used)]
    if not header:
        header = [f"Col {i}" for i in range(n_cols)]
    table_data.append(header)

    for r in range(n_rows):
        row_texts = []
        for ci in range(n_cols):
            cell_id = f"R{r:02d}_C{ci}"
            cell = next((c for c in cells if c.get("cell_id") == cell_id), None)
            row_texts.append(cell.get("text", "") if cell else "")
        table_data.append(row_texts)

    # Generate PDF with reportlab
    try:
        from reportlab.lib.pagesizes import A4
        from reportlab.lib import colors
        from reportlab.platypus import SimpleDocTemplate, Table, TableStyle
        import io as _io

        buf = _io.BytesIO()
        doc = SimpleDocTemplate(buf, pagesize=A4)
        if not table_data or not table_data[0]:
            raise HTTPException(status_code=400, detail="No data to export")

        t = Table(table_data)
        t.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#0d9488')),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.white),
            ('FONTSIZE', (0, 0), (-1, -1), 9),
            ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
            ('VALIGN', (0, 0), (-1, -1), 'TOP'),
            ('WORDWRAP', (0, 0), (-1, -1), True),
        ]))
        doc.build([t])
        buf.seek(0)

        from fastapi.responses import StreamingResponse
        return StreamingResponse(
            buf,
            media_type="application/pdf",
            headers={"Content-Disposition": f'attachment; filename="reconstruction_{session_id}.pdf"'},
        )
    except ImportError:
        raise HTTPException(status_code=501, detail="reportlab not installed")


@router.get("/sessions/{session_id}/reconstruction/export/docx")
async def export_reconstruction_docx(session_id: str):
    """Export the reconstructed cell grid as a DOCX table."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=400, detail="No word result found")

    cells = word_result.get("cells", [])
    columns_used = word_result.get("columns_used", [])
    grid_shape = word_result.get("grid_shape", {})
    n_rows = grid_shape.get("rows", 0)
    n_cols = grid_shape.get("cols", 0)

    try:
        from docx import Document
        from docx.shared import Pt
        import io as _io

        doc = Document()
        doc.add_heading(f'Rekonstruktion – Session {session_id[:8]}', level=1)

        # Build header
        header = [c.get("label", c.get("type", f"Col {i}")) for i, c in enumerate(columns_used)]
        if not header:
            header = [f"Col {i}" for i in range(n_cols)]

        table = doc.add_table(rows=1 + n_rows, cols=max(n_cols, 1))
        table.style = 'Table Grid'

        # Header row
        for ci, h in enumerate(header):
            table.rows[0].cells[ci].text = h

        # Data rows
        for r in range(n_rows):
            for ci in range(n_cols):
                cell_id = f"R{r:02d}_C{ci}"
                cell = next((c for c in cells if c.get("cell_id") == cell_id), None)
                table.rows[r + 1].cells[ci].text = cell.get("text", "") if cell else ""

        buf = _io.BytesIO()
        doc.save(buf)
        buf.seek(0)

        from fastapi.responses import StreamingResponse
        return StreamingResponse(
            buf,
            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            headers={"Content-Disposition": f'attachment; filename="reconstruction_{session_id}.docx"'},
        )
    except ImportError:
        raise HTTPException(status_code=501, detail="python-docx not installed")


@router.post("/sessions/{session_id}/reprocess")
async def reprocess_session(session_id: str, request: Request):
    """Re-run pipeline from a specific step, clearing downstream data.

    Body: {"from_step": 5}  (1-indexed step number)

    Clears downstream results:
    - from_step <= 1: deskew_result, dewarp_result, column_result, row_result, word_result
    - from_step <= 2: dewarp_result, column_result, row_result, word_result
    - from_step <= 3: column_result, row_result, word_result
    - from_step <= 4: row_result, word_result
    - from_step <= 5: word_result (cells, vocab_entries)
    - from_step <= 6: word_result.llm_review only
    """
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    body = await request.json()
    from_step = body.get("from_step", 1)
    if not isinstance(from_step, int) or from_step < 1 or from_step > 7:
        raise HTTPException(status_code=400, detail="from_step must be between 1 and 7")

    update_kwargs: Dict[str, Any] = {"current_step": from_step}

    # Clear downstream data based on from_step
    if from_step <= 5:
        update_kwargs["word_result"] = None
    elif from_step == 6:
        # Only clear LLM review from word_result
        word_result = session.get("word_result")
        if word_result:
            word_result.pop("llm_review", None)
            word_result.pop("llm_corrections", None)
            update_kwargs["word_result"] = word_result

    if from_step <= 4:
        update_kwargs["row_result"] = None
    if from_step <= 3:
        update_kwargs["column_result"] = None
    if from_step <= 2:
        update_kwargs["dewarp_result"] = None
    if from_step <= 1:
        update_kwargs["deskew_result"] = None

    await update_session_db(session_id, **update_kwargs)

    # Also clear cache
    if session_id in _cache:
        for key in list(update_kwargs.keys()):
            if key != "current_step":
                _cache[session_id][key] = update_kwargs[key]
        _cache[session_id]["current_step"] = from_step

    logger.info(f"Session {session_id} reprocessing from step {from_step}")

    return {
        "session_id": session_id,
        "from_step": from_step,
        "cleared": [k for k in update_kwargs if k != "current_step"],
    }


async def _get_rows_overlay(session_id: str) -> Response:
    """Generate dewarped image with row bands drawn on it."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    row_result = session.get("row_result")
    if not row_result or not row_result.get("rows"):
        raise HTTPException(status_code=404, detail="No row data available")

    # Load dewarped image
    dewarped_png = await get_session_image(session_id, "dewarped")
    if not dewarped_png:
        raise HTTPException(status_code=404, detail="Dewarped image not available")

    arr = np.frombuffer(dewarped_png, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(status_code=500, detail="Failed to decode image")

    # Color map for row types (BGR)
    row_colors = {
        "content": (255, 180, 0),     # Blue
        "header": (128, 128, 128),    # Gray
        "footer": (128, 128, 128),    # Gray
        "margin_top": (100, 100, 100),    # Dark Gray
        "margin_bottom": (100, 100, 100), # Dark Gray
    }

    overlay = img.copy()
    for row in row_result["rows"]:
        x, y = row["x"], row["y"]
        w, h = row["width"], row["height"]
        row_type = row.get("row_type", "content")
        color = row_colors.get(row_type, (200, 200, 200))

        # Semi-transparent fill
        cv2.rectangle(overlay, (x, y), (x + w, y + h), color, -1)

        # Solid border
        cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)

        # Label
        idx = row.get("index", 0)
        label = f"R{idx} {row_type.upper()}"
        wc = row.get("word_count", 0)
        if wc:
            label = f"{label} ({wc}w)"
        cv2.putText(img, label, (x + 5, y + 18),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    # Blend overlay at 15% opacity
    cv2.addWeighted(overlay, 0.15, img, 0.85, 0, img)

    success, result_png = cv2.imencode(".png", img)
    if not success:
        raise HTTPException(status_code=500, detail="Failed to encode overlay image")

    return Response(content=result_png.tobytes(), media_type="image/png")


async def _get_words_overlay(session_id: str) -> Response:
    """Generate dewarped image with cell grid drawn on it."""
    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    word_result = session.get("word_result")
    if not word_result:
        raise HTTPException(status_code=404, detail="No word data available")

    # Support both new cell-based and legacy entry-based formats
    cells = word_result.get("cells")
    if not cells and not word_result.get("entries"):
        raise HTTPException(status_code=404, detail="No word data available")

    # Load dewarped image
    dewarped_png = await get_session_image(session_id, "dewarped")
    if not dewarped_png:
        raise HTTPException(status_code=404, detail="Dewarped image not available")

    arr = np.frombuffer(dewarped_png, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(status_code=500, detail="Failed to decode image")

    img_h, img_w = img.shape[:2]

    overlay = img.copy()

    if cells:
        # New cell-based overlay: color by column index
        col_palette = [
            (255, 180, 0),      # Blue (BGR)
            (0, 200, 0),        # Green
            (0, 140, 255),      # Orange
            (200, 100, 200),    # Purple
            (200, 200, 0),      # Cyan
            (100, 200, 200),    # Yellow-ish
        ]

        for cell in cells:
            bbox = cell.get("bbox_px", {})
            cx = bbox.get("x", 0)
            cy = bbox.get("y", 0)
            cw = bbox.get("w", 0)
            ch = bbox.get("h", 0)
            if cw <= 0 or ch <= 0:
                continue

            col_idx = cell.get("col_index", 0)
            color = col_palette[col_idx % len(col_palette)]

            # Cell rectangle border
            cv2.rectangle(img, (cx, cy), (cx + cw, cy + ch), color, 1)
            # Semi-transparent fill
            cv2.rectangle(overlay, (cx, cy), (cx + cw, cy + ch), color, -1)

            # Cell-ID label (top-left corner)
            cell_id = cell.get("cell_id", "")
            cv2.putText(img, cell_id, (cx + 2, cy + 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.28, color, 1)

            # Text label (bottom of cell)
            text = cell.get("text", "")
            if text:
                conf = cell.get("confidence", 0)
                if conf >= 70:
                    text_color = (0, 180, 0)
                elif conf >= 50:
                    text_color = (0, 180, 220)
                else:
                    text_color = (0, 0, 220)

                label = text.replace('\n', ' ')[:30]
                cv2.putText(img, label, (cx + 3, cy + ch - 4),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.35, text_color, 1)
    else:
        # Legacy fallback: entry-based overlay (for old sessions)
        column_result = session.get("column_result")
        row_result = session.get("row_result")
        col_colors = {
            "column_en": (255, 180, 0),
            "column_de": (0, 200, 0),
            "column_example": (0, 140, 255),
        }

        columns = []
        if column_result and column_result.get("columns"):
            columns = [c for c in column_result["columns"]
                       if c.get("type", "").startswith("column_")]

        content_rows_data = []
        if row_result and row_result.get("rows"):
            content_rows_data = [r for r in row_result["rows"]
                                 if r.get("row_type") == "content"]

        for col in columns:
            col_type = col.get("type", "")
            color = col_colors.get(col_type, (200, 200, 200))
            cx, cw = col["x"], col["width"]
            for row in content_rows_data:
                ry, rh = row["y"], row["height"]
                cv2.rectangle(img, (cx, ry), (cx + cw, ry + rh), color, 1)
                cv2.rectangle(overlay, (cx, ry), (cx + cw, ry + rh), color, -1)

        entries = word_result["entries"]
        entry_by_row: Dict[int, Dict] = {}
        for entry in entries:
            entry_by_row[entry.get("row_index", -1)] = entry

        for row_idx, row in enumerate(content_rows_data):
            entry = entry_by_row.get(row_idx)
            if not entry:
                continue
            conf = entry.get("confidence", 0)
            text_color = (0, 180, 0) if conf >= 70 else (0, 180, 220) if conf >= 50 else (0, 0, 220)
            ry, rh = row["y"], row["height"]
            for col in columns:
                col_type = col.get("type", "")
                cx, cw = col["x"], col["width"]
                field = {"column_en": "english", "column_de": "german", "column_example": "example"}.get(col_type, "")
                text = entry.get(field, "") if field else ""
                if text:
                    label = text.replace('\n', ' ')[:30]
                    cv2.putText(img, label, (cx + 3, ry + rh - 4),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.35, text_color, 1)

    # Blend overlay at 10% opacity
    cv2.addWeighted(overlay, 0.1, img, 0.9, 0, img)

    success, result_png = cv2.imencode(".png", img)
    if not success:
        raise HTTPException(status_code=500, detail="Failed to encode overlay image")

    return Response(content=result_png.tobytes(), media_type="image/png")


# ---------------------------------------------------------------------------
# Handwriting Removal Endpoint
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/remove-handwriting")
async def remove_handwriting_endpoint(session_id: str, req: RemoveHandwritingRequest):
    """
    Remove handwriting from a session image using inpainting.

    Steps:
    1. Load source image (auto → deskewed if available, else original)
    2. Detect handwriting mask (filtered by target_ink)
    3. Dilate mask to cover stroke edges
    4. Inpaint the image
    5. Store result as clean_png in the session

    Returns metadata including the URL to fetch the clean image.
    """
    import time as _time
    t0 = _time.monotonic()

    from services.handwriting_detection import detect_handwriting
    from services.inpainting_service import inpaint_image, dilate_mask as _dilate_mask, InpaintingMethod, image_to_png

    session = await get_session_db(session_id)
    if not session:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")

    # 1. Determine source image
    source = req.use_source
    if source == "auto":
        deskewed = await get_session_image(session_id, "deskewed")
        source = "deskewed" if deskewed else "original"

    image_bytes = await get_session_image(session_id, source)
    if not image_bytes:
        raise HTTPException(status_code=404, detail=f"Source image '{source}' not available")

    # 2. Detect handwriting mask
    detection = detect_handwriting(image_bytes, target_ink=req.target_ink)

    # 3. Convert mask to PNG bytes and dilate
    import io
    from PIL import Image as _PILImage
    mask_img = _PILImage.fromarray(detection.mask)
    mask_buf = io.BytesIO()
    mask_img.save(mask_buf, format="PNG")
    mask_bytes = mask_buf.getvalue()

    if req.dilation > 0:
        mask_bytes = _dilate_mask(mask_bytes, iterations=req.dilation)

    # 4. Inpaint
    method_map = {
        "telea": InpaintingMethod.OPENCV_TELEA,
        "ns": InpaintingMethod.OPENCV_NS,
        "auto": InpaintingMethod.AUTO,
    }
    inpaint_method = method_map.get(req.method, InpaintingMethod.AUTO)

    result = inpaint_image(image_bytes, mask_bytes, method=inpaint_method)
    if not result.success:
        raise HTTPException(status_code=500, detail="Inpainting failed")

    elapsed_ms = int((_time.monotonic() - t0) * 1000)

    meta = {
        "method_used": result.method_used.value if hasattr(result.method_used, "value") else str(result.method_used),
        "handwriting_ratio": round(detection.handwriting_ratio, 4),
        "detection_confidence": round(detection.confidence, 4),
        "target_ink": req.target_ink,
        "dilation": req.dilation,
        "source_image": source,
        "processing_time_ms": elapsed_ms,
    }

    # 5. Persist clean image (convert BGR ndarray → PNG bytes)
    clean_png_bytes = image_to_png(result.image)
    await update_session_db(session_id, clean_png=clean_png_bytes, handwriting_removal_meta=meta)

    return {
        **meta,
        "image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/clean",
        "session_id": session_id,
    }


# ---------------------------------------------------------------------------
# Auto-Mode Endpoint (Improvement 3)
# ---------------------------------------------------------------------------

class RunAutoRequest(BaseModel):
    from_step: int = 1          # 1=deskew, 2=dewarp, 3=columns, 4=rows, 5=words, 6=llm-review
    ocr_engine: str = "auto"    # "auto" | "rapid" | "tesseract"
    pronunciation: str = "british"
    skip_llm_review: bool = False
    dewarp_method: str = "ensemble"  # "ensemble" | "vlm" | "cv"


async def _auto_sse_event(step: str, status: str, data: Dict[str, Any]) -> str:
    """Format a single SSE event line."""
    import json as _json
    payload = {"step": step, "status": status, **data}
    return f"data: {_json.dumps(payload)}\n\n"


@router.post("/sessions/{session_id}/run-auto")
async def run_auto(session_id: str, req: RunAutoRequest, request: Request):
    """Run the full OCR pipeline automatically from a given step, streaming SSE progress.

    Steps:
        1. Deskew  — straighten the scan
        2. Dewarp  — correct vertical shear (ensemble CV or VLM)
        3. Columns — detect column layout
        4. Rows    — detect row layout
        5. Words   — OCR each cell
        6. LLM review — correct OCR errors (optional)

    Already-completed steps are skipped unless `from_step` forces a rerun.
    Yields SSE events of the form:
        data: {"step": "deskew", "status": "start"|"done"|"skipped"|"error", ...}

    Final event:
        data: {"step": "complete", "status": "done", "steps_run": [...], "steps_skipped": [...]}
    """
    if req.from_step < 1 or req.from_step > 6:
        raise HTTPException(status_code=400, detail="from_step must be 1-6")
    if req.dewarp_method not in ("ensemble", "vlm", "cv"):
        raise HTTPException(status_code=400, detail="dewarp_method must be: ensemble, vlm, cv")

    if session_id not in _cache:
        await _load_session_to_cache(session_id)

    async def _generate():
        steps_run: List[str] = []
        steps_skipped: List[str] = []
        error_step: Optional[str] = None

        session = await get_session_db(session_id)
        if not session:
            yield await _auto_sse_event("error", "error", {"message": f"Session {session_id} not found"})
            return

        cached = _get_cached(session_id)

        # -----------------------------------------------------------------
        # Step 1: Deskew
        # -----------------------------------------------------------------
        if req.from_step <= 1:
            yield await _auto_sse_event("deskew", "start", {})
            try:
                t0 = time.time()
                orig_bgr = cached.get("original_bgr")
                if orig_bgr is None:
                    raise ValueError("Original image not loaded")

                # Method 1: Hough lines
                try:
                    deskewed_hough, angle_hough = deskew_image(orig_bgr.copy())
                except Exception:
                    deskewed_hough, angle_hough = orig_bgr, 0.0

                # Method 2: Word alignment
                success_enc, png_orig = cv2.imencode(".png", orig_bgr)
                orig_bytes = png_orig.tobytes() if success_enc else b""
                try:
                    deskewed_wa_bytes, angle_wa = deskew_image_by_word_alignment(orig_bytes)
                except Exception:
                    deskewed_wa_bytes, angle_wa = orig_bytes, 0.0

                # Pick best method
                if abs(angle_wa) >= abs(angle_hough) or abs(angle_hough) < 0.1:
                    method_used = "word_alignment"
                    angle_applied = angle_wa
                    wa_arr = np.frombuffer(deskewed_wa_bytes, dtype=np.uint8)
                    deskewed_bgr = cv2.imdecode(wa_arr, cv2.IMREAD_COLOR)
                    if deskewed_bgr is None:
                        deskewed_bgr = deskewed_hough
                        method_used = "hough"
                        angle_applied = angle_hough
                else:
                    method_used = "hough"
                    angle_applied = angle_hough
                    deskewed_bgr = deskewed_hough

                success, png_buf = cv2.imencode(".png", deskewed_bgr)
                deskewed_png = png_buf.tobytes() if success else b""

                deskew_result = {
                    "method_used": method_used,
                    "rotation_degrees": round(float(angle_applied), 3),
                    "duration_seconds": round(time.time() - t0, 2),
                }

                cached["deskewed_bgr"] = deskewed_bgr
                cached["deskew_result"] = deskew_result
                await update_session_db(
                    session_id,
                    deskewed_png=deskewed_png,
                    deskew_result=deskew_result,
                    auto_rotation_degrees=float(angle_applied),
                    current_step=2,
                )
                session = await get_session_db(session_id)

                steps_run.append("deskew")
                yield await _auto_sse_event("deskew", "done", deskew_result)
            except Exception as e:
                logger.error(f"Auto-mode deskew failed for {session_id}: {e}")
                error_step = "deskew"
                yield await _auto_sse_event("deskew", "error", {"message": str(e)})
                yield await _auto_sse_event("complete", "error", {"error_step": error_step})
                return
        else:
            steps_skipped.append("deskew")
            yield await _auto_sse_event("deskew", "skipped", {"reason": "from_step > 1"})

        # -----------------------------------------------------------------
        # Step 2: Dewarp
        # -----------------------------------------------------------------
        if req.from_step <= 2:
            yield await _auto_sse_event("dewarp", "start", {"method": req.dewarp_method})
            try:
                t0 = time.time()
                deskewed_bgr = cached.get("deskewed_bgr")
                if deskewed_bgr is None:
                    raise ValueError("Deskewed image not available")

                if req.dewarp_method == "vlm":
                    success_enc, png_buf = cv2.imencode(".png", deskewed_bgr)
                    img_bytes = png_buf.tobytes() if success_enc else b""
                    vlm_det = await _detect_shear_with_vlm(img_bytes)
                    shear_deg = vlm_det["shear_degrees"]
                    if abs(shear_deg) >= 0.05 and vlm_det["confidence"] >= 0.3:
                        dewarped_bgr = _apply_shear(deskewed_bgr, -shear_deg)
                    else:
                        dewarped_bgr = deskewed_bgr
                    dewarp_info = {
                        "method": vlm_det["method"],
                        "shear_degrees": shear_deg,
                        "confidence": vlm_det["confidence"],
                        "detections": [vlm_det],
                    }
                else:
                    dewarped_bgr, dewarp_info = dewarp_image(deskewed_bgr)

                success_enc, png_buf = cv2.imencode(".png", dewarped_bgr)
                dewarped_png = png_buf.tobytes() if success_enc else b""

                dewarp_result = {
                    "method_used": dewarp_info["method"],
                    "shear_degrees": dewarp_info["shear_degrees"],
                    "confidence": dewarp_info["confidence"],
                    "duration_seconds": round(time.time() - t0, 2),
                    "detections": dewarp_info.get("detections", []),
                }

                cached["dewarped_bgr"] = dewarped_bgr
                cached["dewarp_result"] = dewarp_result
                await update_session_db(
                    session_id,
                    dewarped_png=dewarped_png,
                    dewarp_result=dewarp_result,
                    auto_shear_degrees=dewarp_info.get("shear_degrees", 0.0),
                    current_step=3,
                )
                session = await get_session_db(session_id)

                steps_run.append("dewarp")
                yield await _auto_sse_event("dewarp", "done", dewarp_result)
            except Exception as e:
                logger.error(f"Auto-mode dewarp failed for {session_id}: {e}")
                error_step = "dewarp"
                yield await _auto_sse_event("dewarp", "error", {"message": str(e)})
                yield await _auto_sse_event("complete", "error", {"error_step": error_step})
                return
        else:
            steps_skipped.append("dewarp")
            yield await _auto_sse_event("dewarp", "skipped", {"reason": "from_step > 2"})

        # -----------------------------------------------------------------
        # Step 3: Columns
        # -----------------------------------------------------------------
        if req.from_step <= 3:
            yield await _auto_sse_event("columns", "start", {})
            try:
                t0 = time.time()
                dewarped_bgr = cached.get("dewarped_bgr")
                if dewarped_bgr is None:
                    raise ValueError("Dewarped image not available")

                ocr_img = create_ocr_image(dewarped_bgr)
                h, w = ocr_img.shape[:2]

                geo_result = detect_column_geometry(ocr_img, dewarped_bgr)
                if geo_result is None:
                    layout_img = create_layout_image(dewarped_bgr)
                    regions = analyze_layout(layout_img, ocr_img)
                    cached["_word_dicts"] = None
                    cached["_inv"] = None
                    cached["_content_bounds"] = None
                else:
                    geometries, left_x, right_x, top_y, bottom_y, word_dicts, inv = geo_result
                    content_w = right_x - left_x
                    cached["_word_dicts"] = word_dicts
                    cached["_inv"] = inv
                    cached["_content_bounds"] = (left_x, right_x, top_y, bottom_y)

                    header_y, footer_y = _detect_header_footer_gaps(inv, w, h) if inv is not None else (None, None)
                    geometries = _detect_sub_columns(geometries, content_w, left_x=left_x,
                                                     top_y=top_y, header_y=header_y, footer_y=footer_y)
                    regions = classify_column_types(geometries, content_w, top_y, w, h, bottom_y,
                                                   left_x=left_x, right_x=right_x, inv=inv)

                columns = [asdict(r) for r in regions]
                column_result = {
                    "columns": columns,
                    "classification_methods": list({c.get("classification_method", "") for c in columns if c.get("classification_method")}),
                    "duration_seconds": round(time.time() - t0, 2),
                }

                cached["column_result"] = column_result
                await update_session_db(session_id, column_result=column_result,
                                        row_result=None, word_result=None, current_step=4)
                session = await get_session_db(session_id)

                steps_run.append("columns")
                yield await _auto_sse_event("columns", "done", {
                    "column_count": len(columns),
                    "duration_seconds": column_result["duration_seconds"],
                })
            except Exception as e:
                logger.error(f"Auto-mode columns failed for {session_id}: {e}")
                error_step = "columns"
                yield await _auto_sse_event("columns", "error", {"message": str(e)})
                yield await _auto_sse_event("complete", "error", {"error_step": error_step})
                return
        else:
            steps_skipped.append("columns")
            yield await _auto_sse_event("columns", "skipped", {"reason": "from_step > 3"})

        # -----------------------------------------------------------------
        # Step 4: Rows
        # -----------------------------------------------------------------
        if req.from_step <= 4:
            yield await _auto_sse_event("rows", "start", {})
            try:
                t0 = time.time()
                dewarped_bgr = cached.get("dewarped_bgr")
                session = await get_session_db(session_id)
                column_result = session.get("column_result") or cached.get("column_result")
                if not column_result or not column_result.get("columns"):
                    raise ValueError("Column detection must complete first")

                col_regions = [
                    PageRegion(
                        type=c["type"], x=c["x"], y=c["y"],
                        width=c["width"], height=c["height"],
                        classification_confidence=c.get("classification_confidence", 1.0),
                        classification_method=c.get("classification_method", ""),
                    )
                    for c in column_result["columns"]
                ]

                word_dicts = cached.get("_word_dicts")
                inv = cached.get("_inv")
                content_bounds = cached.get("_content_bounds")

                if word_dicts is None or inv is None or content_bounds is None:
                    ocr_img_tmp = create_ocr_image(dewarped_bgr)
                    geo_result = detect_column_geometry(ocr_img_tmp, dewarped_bgr)
                    if geo_result is None:
                        raise ValueError("Column geometry detection failed — cannot detect rows")
                    _g, lx, rx, ty, by, word_dicts, inv = geo_result
                    cached["_word_dicts"] = word_dicts
                    cached["_inv"] = inv
                    cached["_content_bounds"] = (lx, rx, ty, by)
                    content_bounds = (lx, rx, ty, by)

                left_x, right_x, top_y, bottom_y = content_bounds
                row_geoms = detect_row_geometry(inv, word_dicts, left_x, right_x, top_y, bottom_y)

                row_list = [
                    {
                        "index": r.index, "x": r.x, "y": r.y,
                        "width": r.width, "height": r.height,
                        "word_count": r.word_count,
                        "row_type": r.row_type,
                        "gap_before": r.gap_before,
                    }
                    for r in row_geoms
                ]
                row_result = {
                    "rows": row_list,
                    "row_count": len(row_list),
                    "content_rows": len([r for r in row_geoms if r.row_type == "content"]),
                    "duration_seconds": round(time.time() - t0, 2),
                }

                cached["row_result"] = row_result
                await update_session_db(session_id, row_result=row_result, current_step=5)
                session = await get_session_db(session_id)

                steps_run.append("rows")
                yield await _auto_sse_event("rows", "done", {
                    "row_count": len(row_list),
                    "content_rows": row_result["content_rows"],
                    "duration_seconds": row_result["duration_seconds"],
                })
            except Exception as e:
                logger.error(f"Auto-mode rows failed for {session_id}: {e}")
                error_step = "rows"
                yield await _auto_sse_event("rows", "error", {"message": str(e)})
                yield await _auto_sse_event("complete", "error", {"error_step": error_step})
                return
        else:
            steps_skipped.append("rows")
            yield await _auto_sse_event("rows", "skipped", {"reason": "from_step > 4"})

        # -----------------------------------------------------------------
        # Step 5: Words (OCR)
        # -----------------------------------------------------------------
        if req.from_step <= 5:
            yield await _auto_sse_event("words", "start", {"engine": req.ocr_engine})
            try:
                t0 = time.time()
                dewarped_bgr = cached.get("dewarped_bgr")
                session = await get_session_db(session_id)

                column_result = session.get("column_result") or cached.get("column_result")
                row_result = session.get("row_result") or cached.get("row_result")

                col_regions = [
                    PageRegion(
                        type=c["type"], x=c["x"], y=c["y"],
                        width=c["width"], height=c["height"],
                        classification_confidence=c.get("classification_confidence", 1.0),
                        classification_method=c.get("classification_method", ""),
                    )
                    for c in column_result["columns"]
                ]
                row_geoms = [
                    RowGeometry(
                        index=r["index"], x=r["x"], y=r["y"],
                        width=r["width"], height=r["height"],
                        word_count=r.get("word_count", 0), words=[],
                        row_type=r.get("row_type", "content"),
                        gap_before=r.get("gap_before", 0),
                    )
                    for r in row_result["rows"]
                ]

                word_dicts = cached.get("_word_dicts")
                if word_dicts is not None:
                    content_bounds = cached.get("_content_bounds")
                    top_y = content_bounds[2] if content_bounds else min(r.y for r in row_geoms)
                    for row in row_geoms:
                        row_y_rel = row.y - top_y
                        row_bottom_rel = row_y_rel + row.height
                        row.words = [
                            w for w in word_dicts
                            if row_y_rel <= w['top'] + w['height'] / 2 < row_bottom_rel
                        ]
                        row.word_count = len(row.words)

                ocr_img = create_ocr_image(dewarped_bgr)
                img_h, img_w = dewarped_bgr.shape[:2]

                cells, columns_meta = build_cell_grid(
                    ocr_img, col_regions, row_geoms, img_w, img_h,
                    ocr_engine=req.ocr_engine, img_bgr=dewarped_bgr,
                )
                duration = time.time() - t0

                col_types = {c['type'] for c in columns_meta}
                is_vocab = bool(col_types & {'column_en', 'column_de'})
                n_content_rows = len([r for r in row_geoms if r.row_type == 'content'])
                used_engine = cells[0].get("ocr_engine", "tesseract") if cells else req.ocr_engine

                word_result_data = {
                    "cells": cells,
                    "grid_shape": {
                        "rows": n_content_rows,
                        "cols": len(columns_meta),
                        "total_cells": len(cells),
                    },
                    "columns_used": columns_meta,
                    "layout": "vocab" if is_vocab else "generic",
                    "image_width": img_w,
                    "image_height": img_h,
                    "duration_seconds": round(duration, 2),
                    "ocr_engine": used_engine,
                    "summary": {
                        "total_cells": len(cells),
                        "non_empty_cells": sum(1 for c in cells if c.get("text")),
                        "low_confidence": sum(1 for c in cells if 0 < c.get("confidence", 0) < 50),
                    },
                }

                if is_vocab:
                    entries = _cells_to_vocab_entries(cells, columns_meta)
                    entries = _fix_character_confusion(entries)
                    entries = _fix_phonetic_brackets(entries, pronunciation=req.pronunciation)
                    word_result_data["vocab_entries"] = entries
                    word_result_data["entries"] = entries
                    word_result_data["entry_count"] = len(entries)
                    word_result_data["summary"]["total_entries"] = len(entries)

                await update_session_db(session_id, word_result=word_result_data, current_step=6)
                cached["word_result"] = word_result_data
                session = await get_session_db(session_id)

                steps_run.append("words")
                yield await _auto_sse_event("words", "done", {
                    "total_cells": len(cells),
                    "layout": word_result_data["layout"],
                    "duration_seconds": round(duration, 2),
                    "ocr_engine": used_engine,
                    "summary": word_result_data["summary"],
                })
            except Exception as e:
                logger.error(f"Auto-mode words failed for {session_id}: {e}")
                error_step = "words"
                yield await _auto_sse_event("words", "error", {"message": str(e)})
                yield await _auto_sse_event("complete", "error", {"error_step": error_step})
                return
        else:
            steps_skipped.append("words")
            yield await _auto_sse_event("words", "skipped", {"reason": "from_step > 5"})

        # -----------------------------------------------------------------
        # Step 6: LLM Review (optional)
        # -----------------------------------------------------------------
        if req.from_step <= 6 and not req.skip_llm_review:
            yield await _auto_sse_event("llm_review", "start", {"model": OLLAMA_REVIEW_MODEL})
            try:
                session = await get_session_db(session_id)
                word_result = session.get("word_result") or cached.get("word_result")
                entries = word_result.get("entries") or word_result.get("vocab_entries") or []

                if not entries:
                    yield await _auto_sse_event("llm_review", "skipped", {"reason": "no entries"})
                    steps_skipped.append("llm_review")
                else:
                    reviewed = await llm_review_entries(entries)

                    session = await get_session_db(session_id)
                    word_result_updated = dict(session.get("word_result") or {})
                    word_result_updated["entries"] = reviewed
                    word_result_updated["vocab_entries"] = reviewed
                    word_result_updated["llm_reviewed"] = True
                    word_result_updated["llm_model"] = OLLAMA_REVIEW_MODEL

                    await update_session_db(session_id, word_result=word_result_updated, current_step=7)
                    cached["word_result"] = word_result_updated

                    steps_run.append("llm_review")
                    yield await _auto_sse_event("llm_review", "done", {
                        "entries_reviewed": len(reviewed),
                        "model": OLLAMA_REVIEW_MODEL,
                    })
            except Exception as e:
                logger.warning(f"Auto-mode llm_review failed for {session_id} (non-fatal): {e}")
                yield await _auto_sse_event("llm_review", "error", {"message": str(e), "fatal": False})
                steps_skipped.append("llm_review")
        else:
            steps_skipped.append("llm_review")
            reason = "skipped by request" if req.skip_llm_review else "from_step > 6"
            yield await _auto_sse_event("llm_review", "skipped", {"reason": reason})

        # -----------------------------------------------------------------
        # Final event
        # -----------------------------------------------------------------
        yield await _auto_sse_event("complete", "done", {
            "steps_run": steps_run,
            "steps_skipped": steps_skipped,
        })

    return StreamingResponse(
        _generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )