"""Review sheets for the mentor: per task one row with the student's attempts side by side and the browser report.
usage: mentor_sheet.py --tag mtr_dm --tasks data/visual/mentor_tasks.jsonl --out results/mentor/review --per 5
Rows: '<id>  <first words of the task>' and under each screenshot 'ok' or the problems found (blank, errors, static).
"""
import argparse
import json
import os

from PIL import Image, ImageDraw

ap = argparse.ArgumentParser()
ap.add_argument("--tag", required=True)
ap.add_argument("--tasks", default="data/visual/mentor_tasks.jsonl")
ap.add_argument("--out", required=True)
ap.add_argument("--per", type=int, default=5)
ap.add_argument("--samples", type=int, default=2)
ap.add_argument("--ids", nargs="*")
ap.add_argument("--w", type=int, default=400)
a = ap.parse_args()

tasks = [json.loads(l) for l in open(a.tasks, encoding="utf-8")]
if a.ids:
    tasks = [t for t in tasks if t["id"].split("/")[-1] in a.ids]
rep = {}
for r in json.load(open(os.path.join("results", a.tag, "visual.json"))):
    rep[(r["id"].split("/")[-1], r["sample"])] = r


def flags(r, task):
    if r is None:
        return "no answer"
    out = []
    if r.get("parse_error"):
        out.append("bad XML")
    if r["errors"]:
        out.append("JS: " + r["errors"][0][:60])
    if r["failed"]:
        out.append("load failed")
    if r["blank"]:
        out.append("blank")
    if task["animated"] and not r["animated"] and not r["blank"]:
        out.append("static")
    return "ok" if not out else "; ".join(out)


os.makedirs(a.out, exist_ok=True)
w, h, head, foot = a.w, a.w * 3 // 4, 20, 18
for k in range(0, len(tasks), a.per):
    rows = tasks[k:k + a.per]
    sheet = Image.new("RGB", (a.samples * (w + 6) + 6, len(rows) * (h + head + foot + 8) + 6), "#202124")
    d = ImageDraw.Draw(sheet)
    for i, t in enumerate(rows):
        tid = t["id"].split("/")[-1]
        y = 6 + i * (h + head + foot + 8)
        d.text((8, y + 3), f"{tid}  {t['prompt'][:95]}", fill="#e8eaed")
        for s in range(a.samples):
            x = 6 + s * (w + 6)
            p = os.path.join("results", a.tag, "png", f"{tid}_{s}.png")
            if os.path.exists(p):
                sheet.paste(Image.open(p).convert("RGB").resize((w, h)), (x, y + head))
            else:
                d.rectangle([x, y + head, x + w, y + head + h], outline="#5f6368")
            d.text((x + 2, y + head + h + 3), f"s{s}: {flags(rep.get((tid, s)), t)}", fill="#fdd663")
    sheet.save(os.path.join(a.out, f"sheet_{k // a.per:02d}.png"))
print(f"{(len(tasks) + a.per - 1) // a.per} sheets in {a.out}")