import json, glob, os
import fitz
from PIL import Image, ImageDraw, ImageFont

BASE = os.path.dirname(os.path.abspath(__file__))
try:
    FONT = ImageFont.truetype("C:/Windows/Fonts/arialbd.ttf", 30)
except Exception:
    FONT = ImageFont.load_default()

def splice_level(level_dir, label=None):
    made = 0
    LVL = label or os.path.basename(os.path.normpath(level_dir))
    for mf in sorted(glob.glob(os.path.join(level_dir, "*", "manifest.json"))):
        folder = os.path.dirname(mf)
        pdf = os.path.join(folder, "paper.pdf")
        out = os.path.join(folder, "_figures.png")
        if not os.path.exists(pdf):
            continue
        if os.path.exists(out):
            continue  # idempotent
        try:
            m = json.load(open(mf, encoding="utf-8"))
        except Exception:
            continue
        if "questions" not in m:
            continue
        pages = {}
        for q in m["questions"]:
            if q.get("type_id") == 0:
                continue
            if not q.get("image_needed"):
                continue
            pg = q.get("image_page")
            if not pg:
                continue
            pages.setdefault(int(pg), []).append(q.get("n"))
        if not pages:
            # fallback: figures flagged but no image_page recorded → render all pages
            if not any(q.get("image_needed") and q.get("type_id") != 0 for q in m["questions"]):
                continue
            try:
                _d = fitz.open(pdf); _n = _d.page_count; _d.close()
            except Exception:
                continue
            pages = {p: ["?"] for p in range(1, _n + 1)}
        try:
            doc = fitz.open(pdf)
        except Exception:
            continue
        imgs = []
        for pg in sorted(pages):
            if pg < 1 or pg > doc.page_count:
                continue
            pix = doc[pg - 1].get_pixmap(matrix=fitz.Matrix(2, 2), colorspace=fitz.csRGB)
            im = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
            label = "page %d  ->  Q%s" % (pg, ", Q".join(str(x) for x in pages[pg]))
            band = Image.new("RGB", (im.width, 48), (255, 232, 120))
            ImageDraw.Draw(band).text((10, 9), label, fill=(0, 0, 0), font=FONT)
            imgs.append(band)
            imgs.append(im)
        doc.close()
        if not imgs:
            continue
        W = max(i.width for i in imgs)
        H = sum(i.height for i in imgs) + 12 * len(imgs)
        sheet = Image.new("RGB", (W, H), (255, 255, 255))
        y = 0
        for i in imgs:
            sheet.paste(i, (0, y))
            y += i.height + 12
        sheet.save(out)
        # also write to the central "Old Images" folder named <LEVEL>.<paper>.png
        lvl = LVL
        short = os.path.basename(folder)
        if short.startswith(lvl.lower() + "-"):
            short = short[len(lvl) + 1:]
        oimg = os.path.join(BASE, "Old Images")
        os.makedirs(oimg, exist_ok=True)
        sheet.save(os.path.join(oimg, "%s.%s.png" % (lvl, short)))
        made += 1
        print("wrote", os.path.relpath(out, BASE), "pages", sorted(pages))
    return made

if __name__ == "__main__":
    import sys
    lvl = sys.argv[1] if len(sys.argv) > 1 else "P2"
    if lvl == "P1":
        n = splice_level(os.path.join(BASE, "P1", "_ingest"), label="P1")
    else:
        n = splice_level(os.path.join(BASE, lvl))
    print("spliced %d papers in %s" % (n, lvl))