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

BASE = os.path.dirname(os.path.abspath(__file__))

def slugify(s):
    return re.sub(r'[^a-z0-9]+', '-', (s or '').lower()).strip('-')
try:
    FONT = ImageFont.truetype("C:/Windows/Fonts/arialbd.ttf", 26)
except Exception:
    FONT = ImageFont.load_default()

def crop_level(level_dir, label=None):
    LVL = label or os.path.basename(os.path.normpath(level_dir))
    oimg = os.path.join(BASE, "Old Images"); os.makedirs(oimg, exist_ok=True)
    onlyimg = os.path.join(BASE, "Only Images"); os.makedirs(onlyimg, exist_ok=True)
    made = 0
    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")
        if not os.path.exists(pdf):
            continue
        try:
            m = json.load(open(mf, encoding="utf-8"))
        except Exception:
            continue
        if "questions" not in m:
            continue
        figs = [q for q in m["questions"]
                if q.get("image_needed") and q.get("type_id") != 0
                and q.get("image_bbox") and q.get("image_page")]
        if not figs:
            continue
        try:
            doc = fitz.open(pdf)
        except Exception:
            continue
        short = os.path.basename(folder)
        if short.startswith(LVL.lower() + "-"):
            short = short[len(LVL) + 1:]
        slug = slugify((m.get("paper") or {}).get("source_prefix") or short)
        crops = []
        for q in figs:
            pg = int(q["image_page"]); bb = q["image_bbox"]
            if pg < 1 or pg > doc.page_count or not (isinstance(bb, list) and len(bb) == 4):
                continue
            page = doc[pg - 1]; R = page.rect; W = R.width; H = R.height
            x0, y0, x1, y1 = bb
            clip = fitz.Rect(x0 * W, y0 * H, x1 * W, y1 * H)
            if clip.is_empty or clip.width < 4 or clip.height < 4:
                continue
            pix = page.get_pixmap(matrix=fitz.Matrix(3, 3), clip=clip, colorspace=fitz.csRGB)
            im = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
            im.save(os.path.join(folder, "fig-q%s.png" % q.get("n")))
            # per-question crop straight into Only Images, named to match the DB
            # path pN/images/<slug>-q<n>.png — ready to colourise, no montage split
            im.save(os.path.join(onlyimg, "%s-q%s.png" % (slug, q.get("n"))))
            crops.append((q.get("n"), im))
        doc.close()
        if not crops:
            continue
        pad = 10; lblh = 32
        Wm = max(c.width for _, c in crops)
        Hm = sum(c.height + lblh + pad for _, c in crops) + 10
        sheet = Image.new("RGB", (Wm + 20, Hm), (255, 255, 255)); d = ImageDraw.Draw(sheet); y = 5
        for n, c in crops:
            d.rectangle([0, y, Wm + 20, y + lblh - 4], fill=(255, 232, 120))
            d.text((8, y + 5), "Q%s" % n, fill=(0, 0, 0), font=FONT); y += lblh
            sheet.paste(c, (10, y)); y += c.height + pad
        sheet.save(os.path.join(oimg, "%s.%s.png" % (LVL, short)))
        made += 1
        print("cropped", short, "->", len(crops), "figs")
    print("tight-cropped %d papers in %s" % (made, LVL))

if __name__ == "__main__":
    lvl = sys.argv[1] if len(sys.argv) > 1 else "P3"
    if lvl == "P1":
        crop_level(os.path.join(BASE, "P1", "_ingest"), label="P1")
    else:
        crop_level(os.path.join(BASE, lvl))