#!/usr/bin/env python3
"""
书籍扫描件批量 OCR -> 单本 Markdown 全文版（PaddleOCR-VL-1.6 整页 Markdown）。

背景：PaddleOCR v2 jobs API 单文件 100 页截断，整本书需按 100 页拆片提交；
VL-1.6 模型返回每页 layoutParsingResults[].markdown.text（整页 Markdown，
含表格/版面还原）。本脚本把全书合并为一份 .md 存档（raw/书籍/_ocr-md/）。

Usage:
    python ocr_book_to_md.py <input.pdf> <output.md> [--pages 100] [--keep-chunks]

流程：
  1. fitz 拆片（默认每 100 页一片，临时 PDF 存 <out_dir>/_chunks/）
  2. 每片提交 PaddleOCR-VL-1.6 job 并轮询（失败自动重试 2 次）
  3. done 后下载 jsonl 缓存到 _chunks/chunk_XX.jsonl（已存在则跳过=断点续跑）
  4. 解析 jsonl 中 markdown 文本，按页序合并为 output.md
"""
import json
import os
import sys
import time
import requests
import fitz  # PyMuPDF

JOB_URL = "https://paddleocr.aistudio-app.com/api/v2/ocr/jobs"
MODEL = "PaddleOCR-VL-1.6"
OPTIONAL = {"useDocOrientationClassify": False, "useDocUnwarping": False, "useTextlineOrientation": False}


def load_token():
    tok = os.environ.get("PADDLEOCR_ACCESS_TOKEN", "").strip()
    if tok:
        return tok
    env_path = os.path.join(
        os.path.expanduser("~"), "AppData", "Local",
        "Hermes Agent CN Desktop", "data", "hermes-home", ".env",
    )
    try:
        with open(env_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line.startswith("PADDLEOCR_ACCESS_TOKEN="):
                    return line.split("=", 1)[1].strip()
    except OSError:
        pass
    return ""


TOKEN = load_token()
if not TOKEN:
    print("FATAL: PADDLEOCR_ACCESS_TOKEN not found")
    sys.exit(1)


def submit_and_wait(pdf_path: str, headers: dict, retries: int = 2):
    """提交一片 PDF，轮询至 done，返回 jsonl 文本。失败重试。"""
    data = {"model": MODEL, "optionalPayload": json.dumps(OPTIONAL)}
    for attempt in range(1, retries + 2):
        try:
            with open(pdf_path, "rb") as f:
                resp = requests.post(JOB_URL, headers=headers, data=data, files={"file": f}, timeout=120)
            if resp.status_code != 200:
                print(f"  submit HTTP {resp.status_code}: {resp.text[:300]}")
                raise RuntimeError(f"submit failed: {resp.status_code}")
            job_id = resp.json()["data"]["jobId"]
            print(f"  job={job_id} attempt={attempt}")
            while True:
                time.sleep(8)
                r = requests.get(f"{JOB_URL}/{job_id}", headers=headers, timeout=60)
                st = r.json()["data"]["state"]
                if st == "running":
                    try:
                        ep = r.json()["data"]["extractProgress"]["extractedPages"]
                        tp = r.json()["data"]["extractProgress"]["totalPages"]
                        print(f"  running {ep}/{tp}")
                    except KeyError:
                        print("  running...")
                elif st == "done":
                    jurl = r.json()["data"]["resultUrl"]["jsonUrl"]
                    jr = requests.get(jurl, timeout=120)
                    jr.raise_for_status()
                    print(f"  done ({len(jr.text)} bytes jsonl)")
                    return jr.text
                elif st == "failed":
                    msg = r.json()["data"].get("errorMsg", "?")
                    print(f"  FAILED: {msg}")
                    raise RuntimeError(f"job failed: {msg}")
        except (RuntimeError, requests.RequestException) as e:
            print(f"  attempt {attempt} error: {e}")
            if attempt <= retries:
                print("  retrying in 20s...")
                time.sleep(20)
            else:
                raise


def parse_jsonl(jsonl_text: str):
    """解析 jsonl -> list[page_md]（按行序）。"""
    pages = []
    for line in jsonl_text.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        result = obj.get("result") or {}
        # VL 模型 jsonl 分块返回：每行含 N 个 layoutParsingResults（=N 页），
        # 每个 lpr 的 markdown.text 即一页；dataInfo.numPages 佐证块内页数。
        lpr = result.get("layoutParsingResults") or []
        for res in lpr:
            md = res.get("markdown") or {}
            t = md.get("text")
            if t:
                pages.append(t)
    return pages


def main():
    if len(sys.argv) < 3:
        print(__doc__)
        sys.exit(1)
    pdf_path, out_md = sys.argv[1], sys.argv[2]
    pages_per_chunk = 100
    keep_chunks = False
    if "--pages" in sys.argv:
        pages_per_chunk = int(sys.argv[sys.argv.index("--pages") + 1])
    if "--keep-chunks" in sys.argv:
        keep_chunks = True

    base = os.path.dirname(os.path.abspath(out_md))
    chunks_dir = os.path.join(base, "_chunks")
    os.makedirs(chunks_dir, exist_ok=True)
    headers = {"Authorization": f"bearer {TOKEN}"}

    doc = fitz.open(pdf_path)
    total = doc.page_count
    print(f"PDF pages: {total}, chunks of {pages_per_chunk}")

    all_pages = []
    chunk_idx = 0
    for start in range(0, total, pages_per_chunk):
        end = min(start + pages_per_chunk, total)
        chunk_idx += 1
        chunk_pdf = os.path.join(chunks_dir, f"chunk_{chunk_idx:02d}_{start+1}-{end}.pdf")
        jsonl_cache = chunk_pdf.replace(".pdf", ".jsonl")
        print(f"\n=== chunk {chunk_idx}: pages {start+1}-{end} ===")
        if os.path.exists(jsonl_cache):
            print(f"  cached jsonl exists, skip submit: {os.path.basename(jsonl_cache)}")
            with open(jsonl_cache, encoding="utf-8") as f:
                jsonl_text = f.read()
        else:
            if not os.path.exists(chunk_pdf):
                sub = fitz.open()
                sub.insert_pdf(doc, from_page=start, to_page=end - 1)
                sub.save(chunk_pdf, garbage=3)
                sub.close()
                print(f"  split -> {os.path.basename(chunk_pdf)}")
            jsonl_text = submit_and_wait(chunk_pdf, headers)
            with open(jsonl_cache, "w", encoding="utf-8") as f:
                f.write(jsonl_text)
        pages = parse_jsonl(jsonl_text)
        print(f"  parsed {len(pages)} pages")
        all_pages.extend(pages)
        if not keep_chunks:
            try:
                os.remove(chunk_pdf)
            except OSError:
                pass
    doc.close()

    with open(out_md, "w", encoding="utf-8") as f:
        f.write(f"<!-- OCR: {os.path.basename(pdf_path)} | {MODEL} | {total} pages | {time.strftime('%Y-%m-%d %H:%M')} -->\n\n")
        for i, pg in enumerate(all_pages, 1):
            f.write(f"\n\n<!-- page {i} -->\n\n{pg.strip()}")
    print(f"\nDONE: {out_md} ({len(all_pages)} pages, {os.path.getsize(out_md)} bytes)")


if __name__ == "__main__":
    main()