# worker_parse.py — manga parse (paged manga only). FastAPI :8009. GPU, session-guarded ("magi").
# Magi v2 chapter-wide pass: panel detection + reading order + OCR in one shot. Replaces the
# crop+ocr stages for paged manga; downstream ocr stage no-ops because rows are pre-populated.
# webtoons do NOT come here — they use worker_crop /crop/webtoon. See manga-two-repo-split memory.
import os, uuid
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import cv2
import numpy as np
import transport

app = FastAPI()
SHM = "/dev/shm"
MAGI_MODEL = "ragavsachdeva/magiv2"
_model = None


def _load_magi():
    global _model
    if _model is None:
        import torch
        from transformers import AutoModel
        _model = AutoModel.from_pretrained(MAGI_MODEL, trust_remote_code=True).cuda().eval()
        _model._torch = torch
    return _model


class ParseInput(BaseModel):
    page_uris: list        # all pages of the chapter, in order
    manga_id: str
    chapter_id: str
    session_id: str = ""   # magi GPU lease (opened by orchestrator)
    job_id: str = ""


def _center_in(box, panel) -> bool:
    x1, y1, x2, y2 = box
    px1, py1, px2, py2 = panel
    cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
    return px1 <= cx <= px2 and py1 <= cy <= py2


def assemble_panels(pages, results, manga_id, chapter_id, put):
    """Flatten Magi's per-page output into chapter-order panels with their OCR.
    `pages`: RGB np arrays. `results`: per-page dicts (Magi keys). `put(np_crop, uri)` uploads.
    Text is assigned to the panel whose box contains the text-box center; SFX (non-essential)
    is dropped so narration isn't polluted. bbox converted [x1,y1,x2,y2] -> [x,y,w,h]."""
    out, gidx = [], 0
    for img, res in zip(pages, results):
        panels = res.get("panels", [])
        texts = res.get("texts", [])
        ocr = res.get("ocr", [])
        essential = res.get("is_essential_text", [True] * len(texts))
        for p in panels:
            x1, y1, x2, y2 = (int(v) for v in p)
            uri = f"s3://manga/{manga_id}/{chapter_id}/panels/p{gidx:03d}.png"
            if not transport.exists(uri):        # deterministic per gidx -> resumable
                put(img[y1:y2, x1:x2], uri)
            ocr_texts = []
            for ti, tb in enumerate(texts):
                if ti < len(ocr) and essential[ti] and _center_in(tb, p):
                    tx1, ty1, tx2, ty2 = (int(v) for v in tb)
                    ocr_texts.append({"text_id": f"t{ti}", "content": ocr[ti],
                                      "bbox": [tx1, ty1, tx2 - tx1, ty2 - ty1], "confidence": 1.0})
            out.append({"panel_index": gidx, "uri": uri,
                        "bbox": [x1, y1, x2 - x1, y2 - y1], "ocr": ocr_texts})
            gidx += 1
    return out


def _put_crop(np_rgb, uri):
    tmp = f"{SHM}/parse_{uuid.uuid4().hex[:8]}.png"
    cv2.imwrite(tmp, cv2.cvtColor(np_rgb, cv2.COLOR_RGB2BGR))
    transport.put(tmp, uri)
    os.remove(tmp)


@app.post("/parse")
async def parse(data: ParseInput):
    tag = uuid.uuid4().hex[:8]
    pages = []
    for i, u in enumerate(data.page_uris):
        local = transport.get(u, f"{SHM}/pp_{tag}_{i:03d}.png")
        img = cv2.imread(local)
        if img is None:
            raise HTTPException(400, f"page not readable: {u}")
        pages.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        os.remove(local)
    model = _load_magi()
    # ponytail: empty character bank in v1 — identity stays with the downstream siglip stage;
    #   feed a real bank (known-char ref crops + names) here to get Magi speaker association.
    bank = {"images": [], "names": []}
    with model._torch.no_grad():
        results = model.do_chapter_wide_prediction(pages, bank, use_tqdm=False, do_ocr=True)
    panels = assemble_panels(pages, results, data.manga_id, data.chapter_id, _put_crop)
    return {"panels": panels}


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


if __name__ == "__main__":
    # self-check: model-free. Fake a 2-page Magi result and assert panel flattening,
    # chapter-order indexing, text->panel containment, SFX drop, and bbox conversion.
    transport.exists = lambda *a, **k: False   # no minio in self-check
    stored = {}
    pages = [np.zeros((100, 100, 3), np.uint8), np.zeros((100, 100, 3), np.uint8)]
    results = [
        {"panels": [[0, 0, 50, 100], [50, 0, 100, 100]],           # page 0: two panels
         "texts": [[10, 10, 20, 20], [60, 10, 70, 20]],             # one text in each
         "ocr": ["HELLO", "BOOM"], "is_essential_text": [True, False]},  # BOOM = SFX, dropped
        {"panels": [[0, 0, 100, 100]],                              # page 1: one panel
         "texts": [[5, 5, 15, 15]], "ocr": ["WORLD"], "is_essential_text": [True]},
    ]
    panels = assemble_panels(pages, results, "m", "c", lambda img, uri: stored.__setitem__(uri, img.shape))
    assert [p["panel_index"] for p in panels] == [0, 1, 2], "chapter-order index"
    assert panels[0]["ocr"][0]["content"] == "HELLO"
    assert panels[1]["ocr"] == [], "SFX text dropped from panel 1"
    assert panels[2]["ocr"][0]["content"] == "WORLD"
    assert panels[0]["bbox"] == [0, 0, 50, 100], "xyxy->xywh"
    assert panels[0]["ocr"][0]["bbox"] == [10, 10, 10, 10]
    assert stored, "crops uploaded via put"
    print("worker_parse self-check ok")