"""Extract and OCR every image referenced by an EPUB's reading spine.

Usage: python tools/ocr_epub_images.py "<Book Title>"

Writes extracted images to library/<Book Title>/images/ and OCR results to
library/<Book Title>/ocr-embedded-images.json.
"""
import json
import posixpath
import sys
import zipfile
from collections import Counter
from pathlib import Path
from xml.etree import ElementTree as ET

import cv2
import numpy as np
from lxml import html
from rapidocr_onnxruntime import RapidOCR


ROOT = Path(__file__).resolve().parent.parent / "library"
XLINK = "{http://www.w3.org/1999/xlink}href"


def epub_spine(epub):
    container = ET.fromstring(epub.read("META-INF/container.xml"))
    opf_path = container.find(".//{*}rootfile").attrib["full-path"]
    opf = ET.fromstring(epub.read(opf_path))
    base = posixpath.dirname(opf_path)
    manifest = {
        item.attrib["id"]: item.attrib
        for item in opf.findall(".//{*}manifest/{*}item")
    }
    return [
        posixpath.normpath(posixpath.join(base, manifest[item.attrib["idref"]]["href"]))
        for item in opf.findall(".//{*}spine/{*}itemref")
    ]


def image_refs(epub, documents):
    refs = []
    for document in documents:
        root = html.fromstring(epub.read(document))
        for element in root.xpath("//*[local-name()='img' or local-name()='image']"):
            href = element.get("src") or element.get(XLINK) or element.get("xlink:href") or element.get("href")
            if href:
                refs.append(posixpath.normpath(posixpath.join(posixpath.dirname(document), href)))
    return refs


def main(book):
    book_root = ROOT / book
    epub_path = book_root / "source-file.epub"
    output_path = book_root / "ocr-embedded-images.json"
    images_root = book_root / "images"
    images_root.mkdir(exist_ok=True)
    engine = RapidOCR()

    with zipfile.ZipFile(epub_path) as epub:
        documents = epub_spine(epub)
        refs = image_refs(epub, documents)
        counts = Counter(refs)
        unique_refs = list(dict.fromkeys(refs))
        recovered = (
            json.loads(output_path.read_text(encoding="utf-8"))
            if output_path.exists()
            else {}
        )
        print(f"EPUB images for {book}: {len(unique_refs)} unique, {len(refs)} references", flush=True)

        for index, ref in enumerate(unique_refs, start=1):
            cached = recovered.get(ref)
            if cached and (images_root / Path(cached["asset"]).name).exists():
                print(f"{index}/{len(unique_refs)} {ref}: cached", flush=True)
                continue
            data = epub.read(ref)
            asset_name = Path(ref).name
            (images_root / asset_name).write_bytes(data)
            image = cv2.imdecode(np.frombuffer(data, dtype=np.uint8), cv2.IMREAD_COLOR)
            if image is None:
                lines = []
            else:
                result, _ = engine(image)
                lines = [item[1].strip() for item in (result or []) if item[1].strip()]
            recovered[ref] = {
                "asset": f"/library/{book}/images/{asset_name}",
                "references": counts[ref],
                "text": "\n".join(lines),
            }
            print(f"{index}/{len(unique_refs)} {ref}: {len(lines)} OCR lines", flush=True)

    output_path.write_text(
        json.dumps(recovered, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
        newline="\n",
    )
    print(f"Wrote {output_path} and {images_root}", flush=True)


if __name__ == "__main__":
    main(sys.argv[1])