import argparse
import json
from pathlib import Path

import pypdfium2 as pdfium
from PIL import Image, ImageDraw


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "research" / "page_images"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("source")
    parser.add_argument("pages", help="Comma separated 1-based PDF pages")
    parser.add_argument("--scale", type=float, default=1.6)
    args = parser.parse_args()
    manifest = json.loads((ROOT / "research" / "local_sources.json").read_text(encoding="utf-8"))
    public = ROOT / "research" / "public_downloads.json"
    if public.exists():
        manifest.extend(json.loads(public.read_text(encoding="utf-8")))
    source = next(row for row in manifest if row["id"] == args.source)
    document = pdfium.PdfDocument(ROOT / source["file"])
    OUT.mkdir(parents=True, exist_ok=True)
    thumbs = []
    for number in map(int, args.pages.split(",")):
        page = document[number - 1]
        image = page.render(scale=args.scale).to_pil()
        image.save(OUT / f"{args.source}-p{number}.png")
        image.thumbnail((850, 620))
        thumb = Image.new("RGB", (870, 660), "white")
        thumb.paste(image, (10, 30))
        ImageDraw.Draw(thumb).text((10, 8), f"{args.source} | PDF page {number}", fill="black")
        thumbs.append(thumb)
    if len(thumbs) > 1:
        columns = min(2, len(thumbs))
        rows = (len(thumbs) + columns - 1) // columns
        sheet = Image.new("RGB", (870 * columns, 660 * rows), "#e9edef")
        for index, thumb in enumerate(thumbs):
            sheet.paste(thumb, (index % columns * 870, index // columns * 660))
        sheet.save(OUT / f"{args.source}-contact.png")
    print(source["name"], args.pages)


if __name__ == "__main__":
    main()