# worker_ocr.py — stage 3 text extraction. FastAPI :8001. cpu (easyocr), no session.
# panel in -> text blocks with bboxes + confidence out. orchestrator persists to sqlite.
# easyocr (not tesseract): it reads stylized manga lettering far better -- recovers whole
# lines tesseract garbles or drops. runs on GPU (~0.4s/page warm) by default; the OCR stage
# runs before any LLM session opens so it doesn't contend with the resident model. set
# OCR_GPU=0 to force CPU (~3s/page). GPU needs MIOPEN_FIND_MODE=FAST in the env or the first
# ROCm run spends ~60s in MIOpen's exhaustive kernel search -- the launcher sets it.
import os, uuid
from fastapi import FastAPI
from pydantic import BaseModel
import transport

app = FastAPI()
SHM = "/dev/shm"
MIN_CONF = 0.3      # easyocr line confidence floor
OCR_GPU = os.environ.get("OCR_GPU", "1") == "1"

_reader = None      # easyocr.Reader, lazy-loaded on first request


def _get_reader():
    global _reader
    if _reader is None:
        import easyocr
        _reader = easyocr.Reader(["en"], gpu=OCR_GPU, verbose=False)
    return _reader


def _detections_to_texts(detections):
    """easyocr readtext output [(box_pts, text, conf)] -> our text blocks with xywh bboxes.
    box_pts is 4 corner [x,y] points. drops low-confidence and art-noise (<2 letters).
    casing is left as-is (mixed) -- the vision stage re-cases from the image anyway."""
    texts = []
    for i, (box, txt, conf) in enumerate(detections):
        txt = txt.strip()
        if conf < MIN_CONF or sum(c.isalpha() for c in txt) < 2:
            continue
        xs = [p[0] for p in box]; ys = [p[1] for p in box]
        x, y = int(min(xs)), int(min(ys))
        texts.append({
            "id": f"t{i+1:03d}",
            "content": txt,
            "bbox": [x, y, int(max(xs)) - x, int(max(ys)) - y],
            "confidence": round(float(conf), 3),
        })
    return texts


def ocr_image(path: str):
    return _detections_to_texts(_get_reader().readtext(path, detail=1, paragraph=False))


class OCRInput(BaseModel):
    panel_uri: str
    job_id: str = ""
    panel_id: str = ""


@app.post("/ocr")
async def ocr(data: OCRInput):
    local = transport.get(data.panel_uri, f"{SHM}/ocr_{uuid.uuid4().hex[:8]}.png")
    texts = ocr_image(local)
    os.remove(local)
    return {"panel_id": data.panel_id, "texts": texts}


@app.post("/unload")
async def unload():
    """free the resident easyocr reader (~1-2GB) once the OCR stage is done, before gemma4 loads.
    ocr isn't session-managed, so the orchestrator calls this at stage end."""
    global _reader
    was = _reader is not None
    _reader = None
    import gc; gc.collect()
    try:
        import torch; torch.cuda.empty_cache()
    except Exception:
        pass
    return {"ok": True, "unloaded": was}


@app.get("/health")
async def health():
    return {"status": "ok"}


if __name__ == "__main__":
    # self-check: detection->text-block conversion (pure, no model needed).
    dets = [
        ([[10, 10], [110, 10], [110, 40], [10, 40]], "HELLO", 0.9),   # kept
        ([[10, 200], [70, 200], [70, 230], [10, 230]], "WORLD", 0.8),  # kept
        ([[5, 5], [13, 5], [13, 13], [5, 13]], "=", 0.9),              # art-noise: <2 letters
        ([[0, 0], [50, 0], [50, 20], [0, 20]], "REAL", 0.1),           # below MIN_CONF
    ]
    texts = _detections_to_texts(dets)
    assert [t["content"] for t in texts] == ["HELLO", "WORLD"], texts
    assert texts[0]["bbox"] == [10, 10, 100, 30], texts[0]["bbox"]
    assert texts[1]["confidence"] == 0.8
    print("worker_ocr self-check ok")