"""OCR the Level-4 books with the local Dolf OCR Platform (http://localhost:10010).

Each PDF is split into chunks of CHUNK pages (uploads stay small, progress is resumable).
Output: Book/L4/<slug>/pNNN-MMM.md with '# Page N' headings using the book's own page numbers.
The API key is read from LOCALAPIKEY in the OCR project's .env and never printed.

  python tools/ocr_level4.py [slug ...] [--max-chunks N]
"""
import os, sys, time, tempfile, concurrent.futures as cf
import fitz, requests

SRC = r"D:/OneDrive - DolfTech/SABIS-EBooks/EBooks/French"
OCR_ENV = r"F:/Solutions/AI/pdf_datasets_ocr-enhanced/.env"
BASE = "http://localhost:10010"
OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "Book", "L4")
CHUNK = 20
BOOKS = {  # slug -> folder name
    "conjugaison4": "Exercices de Conjugaison 4",
    "grammaire4": "GRAMMAIRE 4",
    "parlez3": "À vos marques prêts parlez 3",
    "parlez5": "À vos marques prêts parlez 5",
}
OPTIONS = {"lang_mode": "fr", "handwriting": "false", "extract_images": "true", "ocr_image_captions": "false", "use_cache": "true", "max_tokens": "3000"}


def api_key():
    for line in open(OCR_ENV, encoding="utf-8"):
        if line.startswith("LOCALAPIKEY="):
            return line.split("=", 1)[1].strip().strip('"').strip("'")
    sys.exit("LOCALAPIKEY not found")


def book_pdf(slug):
    """The folder names on disk may use a different Unicode form (À as one or two code points)."""
    import unicodedata as u
    want = u.normalize("NFC", BOOKS[slug])
    folder = next(d for d in os.listdir(SRC) if u.normalize("NFC", d) == want)
    pdf = next(f for f in os.listdir(os.path.join(SRC, folder)) if f.lower().endswith(".pdf"))
    return os.path.join(SRC, folder, pdf)


def chunks(slug):
    pdf = book_pdf(slug)
    doc = fitz.open(pdf)
    for start in range(0, doc.page_count, CHUNK):
        end = min(start + CHUNK, doc.page_count)
        yield doc, start, end, os.path.join(OUT, slug, f"p{start + 1:03d}-{end:03d}.md")


def ocr_chunk(key, doc, start, end, dst):
    tmp = os.path.join(tempfile.gettempdir(), f"l4_{os.getpid()}_{start}.pdf")
    part = fitz.open()
    part.insert_pdf(doc, from_page=start, to_page=end - 1)
    part.save(tmp, garbage=3, deflate=True)
    part.close()
    s = requests.Session()
    t0 = time.time()
    for attempt in range(4):
        with open(tmp, "rb") as fh:
            r = s.post(f"{BASE}/api/ocr", headers={"X-API-Key": key}, files={"file": (os.path.basename(dst)[:-3] + ".pdf", fh)}, data=OPTIONS, timeout=180)
        if r.status_code == 429:
            time.sleep(60); continue
        r.raise_for_status(); break
    d = r.json()
    os.remove(tmp)
    if d.get("cached"):
        pages = d["result"]
    else:
        tid = d["task_id"]
        while True:
            time.sleep(4)
            st = s.get(f"{BASE}/api/status/{tid}", headers={"X-API-Key": key}, timeout=60).json()
            if st.get("status") == "done":
                pages = st.get("result") or []; break
            if st.get("status") not in ("processing", None):
                raise RuntimeError(f"{dst}: {st.get('status')} {st.get('error')}")
            if time.time() - t0 > 3600:
                raise TimeoutError(dst)
    parts = [f"# Page {start + p.get('page', i + 1)}\n\n{p.get('text') or ''}" for i, p in enumerate(sorted(pages, key=lambda p: p.get("page", 0)))]
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    with open(dst + ".tmp", "w", encoding="utf-8") as f:
        f.write("\n\n---\n\n".join(parts) + "\n")
    os.replace(dst + ".tmp", dst)
    return f"{os.path.relpath(dst, OUT)}: {len(pages)} pages in {time.time() - t0:.0f}s"


def main():
    for st in (sys.stdout, sys.stderr):
        st.reconfigure(encoding="utf-8", errors="replace")
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    maxc = int(sys.argv[sys.argv.index("--max-chunks") + 1]) if "--max-chunks" in sys.argv else 10 ** 6
    if maxc != 10 ** 6:
        args = [a for a in args if a != str(maxc)]
    key = api_key()
    todo = [c for slug in (args or BOOKS) for c in chunks(slug) if not os.path.exists(c[3])][:maxc]
    print(f"{len(todo)} chunk(s) to OCR", flush=True)
    with cf.ThreadPoolExecutor(max_workers=2) as ex:   # the platform runs 2 GPU workers
        futs = []
        for c in todo:
            futs.append(ex.submit(ocr_chunk, key, *c)); time.sleep(7)   # stay under 10 submissions/minute
        for f in cf.as_completed(futs):
            try:
                print(f.result(), flush=True)
            except Exception as e:
                print("FAILED", e, flush=True)


if __name__ == "__main__":
    main()