#!/usr/bin/env python3
"""
为已合并的 OCR 书籍 md 补插图：从 _chunks/*.jsonl 提取每页 markdown.images，
下载到 <md_dir>/images/，并在每页 '<!-- page N -->' 标记后插入图片引用。

用法: python ocr_book_insert_images.py <book.md>
（jsonl 缓存需在 md 同目录 _chunks/ 下，即 ocr_book_to_md.py 的产出布局）
"""
import json
import glob
import os
import re
import sys
import requests

PAGE_RE = re.compile(r"<!-- page (\d+) -->")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    md_path = os.path.abspath(sys.argv[1])
    md_dir = os.path.dirname(md_path)
    chunks_dir = os.path.join(md_dir, "_chunks")
    img_root = os.path.join(md_dir, "images")

    # 1) 按页序遍历所有 lpr，收集每页 images（全局游标推进）
    page_imgs = []          # list of dict(key->url)，与页序对齐
    for jf in sorted(glob.glob(os.path.join(chunks_dir, "chunk_*.jsonl"))):
        for line in open(jf, encoding="utf-8"):
            line = line.strip()
            if not line:
                continue
            try:
                o = json.loads(line)
            except json.JSONDecodeError:
                continue
            for res in (o.get("result") or {}).get("layoutParsingResults") or []:
                imgs = (res.get("markdown") or {}).get("images") or {}
                page_imgs.append(imgs)
    print(f"pages from jsonl: {len(page_imgs)}")

    # 2) 下载（同 URL 去重），保存到 images/<key>
    os.makedirs(img_root, exist_ok=True)
    url_to_key = {}
    n_dl = n_skip = n_fail = 0
    for pidx, imgs in enumerate(page_imgs, 1):
        for key, url in imgs.items():
            dst = os.path.join(img_root, key)
            if url in url_to_key:
                n_skip += 1
                continue
            if os.path.exists(dst) and os.path.getsize(dst) > 0:
                url_to_key[url] = key  # 已下载过，仅登记去重
                n_skip += 1
                continue
            try:
                os.makedirs(os.path.dirname(dst), exist_ok=True)
                r = requests.get(url, timeout=60)
                if r.status_code != 200:
                    raise RuntimeError(f"HTTP {r.status_code}")
                with open(dst, "wb") as f:
                    f.write(r.content)
                url_to_key[url] = key
                n_dl += 1
                print(f"  [page {pidx}] dl {key} ({len(r.content)} B)")
            except Exception as e:
                n_fail += 1
                print(f"  [page {pidx}] FAIL {key}: {e}")
    print(f"downloaded: {n_dl}, dedup-skip: {n_skip}, failed: {n_fail}")

    # 3) 读 md，按页补齐图片引用（幂等）：
    #    VL 的 markdown.text 自带 HTML <img src="imgs/xxx.jpg">，路径需修为 images/imgs/；
    #    仅对 text 中未出现的图（images 字段有但 text 无引用）插入 markdown 引用，避免双图。
    with open(md_path, encoding="utf-8") as f:
        md_text = f.read()
    # 修 VL 原生 HTML 路径（幂等：已修的 images/imgs/ 不再变）
    md_text = re.sub(r'(<img src=")imgs/', r"\1images/imgs/", md_text)
    # 已有引用集合（HTML 或 markdown 形式，均以 imgs/<key> 结尾）
    have = set(re.findall(r"images/imgs/([^\")\s]+\.jpg)", md_text))
    lines = md_text.split("\n")
    out = []
    page_no = 0
    inserted = 0
    for line in lines:
        out.append(line)
        m = PAGE_RE.search(line)
        if m:
            page_no = int(m.group(1))
            imgs = page_imgs[page_no - 1] if page_no - 1 < len(page_imgs) else {}
            for i, (key, _url) in enumerate(imgs.items(), 1):
                rel = key.replace("\\", "/")
                if os.path.basename(key) in have:  # key 形如 imgs/xxx.jpg，have 存 basename
                    continue  # text 已有 HTML 引用，不重复
                out.append(f"![p{page_no}-img{i}](images/{rel})")
                have.add(os.path.basename(key))
                inserted += 1
    with open(md_path, "w", encoding="utf-8") as f:
        f.write("\n".join(out))
    print(f"inserted {inserted} missing image refs into {md_path}")


if __name__ == "__main__":
    main()