#!/usr/bin/env python3
"""Element-head teacher data, step 2: Claude Haiku 5.5 picks how each image should render in dark mode.

Reads the contact sheets from elem_sheets.py and asks, per row, for the best rendering (A keep, B invert, C adapt,
D remove), every acceptable one, and the image's role. Batch API (half price) with a cost ledger and a hard budget;
reads ANTHROPIC_API_KEY from the environment. Output: OUT/labels.jsonl, one line per sheet.
"""
import argparse, base64, json, os, random, sys, time, urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

MODEL = "claude-haiku-5-5"
PRICE_IN, PRICE_OUT = 0.10 / 1e6, 0.50 / 1e6
PROXY = os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
HDRS = {"x-api-key": os.environ.get("ANTHROPIC_API_KEY", ""), "anthropic-version": "2023-06-01", "content-type": "application/json"}
ROLES = ["photo", "illustration", "diagram", "logo", "icon", "text", "backdrop", "avatar", "ui"]
OPTS = ["A", "B", "C", "D"]

SYSTEM = """You are a senior product designer reviewing how a dark-mode browser extension renders the images of light websites. Each numbered row of the sheet is one image from a real web page.

Left: the image in its original light page (inside the red box), or on its light background when the page area was not captured. Then the candidate renderings on the dark surface that the dark theme puts behind that image. These are the exact pixels the browser would show:
- A keep: unchanged.
- B invert: hue-preserving inversion. Dark becomes light and light becomes dark while hues stay, so yellows, pastels and white areas turn dark.
- C adapt: only dark neutral ink (black or gray text, lines, glyphs) becomes light, and a white background box is removed; colours are kept and lifted just enough to read on dark.
- D remove: the background image is hidden so the dark surface shows through (offered only for large CSS backgrounds; the sample text shows how content on top of it will look in dark mode).
"n/a" means that rendering is not possible for this image.

Pick the rendering a careful designer would ship in a polished dark theme:
- Photos, video thumbnails, artwork, product shots, screenshots, maps, avatars, flags, emoji, ads and colourful illustrations must look exactly as their author intended: A. Never alter the colours of real content.
- Everything must stay clearly visible. Dark icons, arrows, glyphs, text rendered as an image, math formulas and line art that would vanish or be hard to see on the dark surface must become light: B or C, whichever looks more natural.
- Logos must stay recognizable and true to the brand. Keep brand colours: A when the logo already reads well on the dark surface, C when it mixes brand colours with dark text or sits in a white box, B only if the inverted logo still looks like the brand.
- Diagrams, charts and documents on white: A unless the inverted or adapted version stays faithful and clearly readable.
- Large light textures, pale washes, subtle patterns and plain light panels behind content: D. Photographic or illustrated backgrounds: A.
- A white box around an image is acceptable for content (photos, screenshots), not for icons and logos when a cleaner option exists.

For every row return its number k, the image's role, the best rendering, and every rendering you would consider acceptable (always including the best). Roles: photo, illustration, diagram (charts, maps, documents, screenshots of text), logo, icon (UI glyphs and symbols), text (text or math rendered as an image), backdrop (decorative background behind content), avatar (people, profile pictures), ui (other interface decoration: dividers, bullets, shapes, badges)."""

SCHEMA = {
    "type": "object",
    "properties": {"rows": {"type": "array", "items": {"type": "object", "properties": {
        "k": {"type": "integer"}, "role": {"type": "string", "enum": ROLES}, "best": {"type": "string", "enum": OPTS},
        "ok": {"type": "array", "items": {"type": "string", "enum": OPTS}}}, "required": ["k", "role", "best", "ok"], "additionalProperties": False}}},
    "required": ["rows"], "additionalProperties": False,
}


def http(method, path, body=None, timeout=300):
    req = urllib.request.Request(PROXY + path, data=json.dumps(body).encode() if body is not None else None, headers=HDRS, method=method)
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read()


def as_jpeg(png):
    from PIL import Image
    import io
    buf = io.BytesIO()
    Image.open(io.BytesIO(png)).convert("RGB").save(buf, "JPEG", quality=88)
    return buf.getvalue()


def params(sheet_png, n_rows, args):
    content = [{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(as_jpeg(sheet_png)).decode()}},
               {"type": "text", "text": f"This sheet has {n_rows} rows (#1-#{n_rows}). Return one entry per row."}]
    return {"model": MODEL, "max_tokens": args.max_out, "system": SYSTEM, "thinking": {"type": "adaptive"},
            "output_config": {"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
            "messages": [{"role": "user", "content": content}]}


def parse(msg, rows):
    txt = "".join(b.get("text", "") for b in msg.get("content", []) if b.get("type") == "text")
    ans = json.loads(txt)["rows"]
    out = {}
    for a in ans:
        k = int(a["k"])
        if not 1 <= k <= len(rows) or k in out:
            continue
        offer = rows[k - 1]["offer"]
        best = a["best"] if a["best"] in offer else None
        ok = sorted(set(x for x in a["ok"] if x in offer) | ({best} if best else set()))
        if best:
            out[k] = {"role": a["role"], "best": best, "ok": ok}
    return out


def usage_cost(u, batch=True):
    c = u.get("input_tokens", 0) * PRICE_IN + u.get("output_tokens", 0) * PRICE_OUT
    return c * (0.5 if batch else 1.0)


class Ledger:
    def __init__(self, path):
        self.path, self.cost, self.n = path, 0.0, 0
        if path.exists():
            for l in path.open():
                r = json.loads(l)
                self.cost += r["cost"]
                self.n += 1

    def add(self, rec):
        self.cost += rec["cost"]
        self.n += 1
        with self.path.open("a") as f:
            f.write(json.dumps(rec) + "\n")


def load_sheets(sheet_dir):
    sheets = {}
    for mp in sorted(Path(sheet_dir).glob("meta.*.jsonl")):
        for l in mp.open():
            if l.strip():
                s = json.loads(l)
                sheets[s["sheet"]] = s
    return sheets


def write_label(out_f, sid, meta, labels, rec, dup=False):
    rows = []
    for r in meta["rows"]:
        lab = labels.get(r["k"])
        if lab:
            rows.append(dict(r, **lab))
    out_f.write(json.dumps({"sheet": sid, "dup": dup, "type": meta["type"], "rows": rows, "usage": rec}) + "\n")
    out_f.flush()


def pilot(args, sheets, ledger, out_dir):
    ids = list(sheets)[: args.pilot]
    with (out_dir / "pilot.jsonl").open("a") as f:
        def one(sid):
            meta = sheets[sid]
            png = (Path(args.sheets) / "sheets" / f"{sid}.png").read_bytes()
            t0 = time.time()
            m = json.loads(http("POST", "/v1/messages", params(png, len(meta["rows"]), args), timeout=600))
            u = m.get("usage", {})
            rec = {"id": sid, "cost": usage_cost(u, batch=False), "in": u.get("input_tokens"), "out": u.get("output_tokens"), "secs": round(time.time() - t0, 1)}
            ledger.add(rec)
            write_label(f, sid, meta, parse(m, meta["rows"]), rec)
            return rec
        with ThreadPoolExecutor(6) as ex:
            for rec in ex.map(one, ids):
                print(rec, flush=True)
    print(f"pilot spent ${ledger.cost:.4f} on {ledger.n}", flush=True)


def batch(args, sheets, ledger, out_dir):
    state_p = out_dir / "batches.jsonl"
    batches = {}
    if state_p.exists():
        for l in state_p.open():
            b = json.loads(l)
            batches.setdefault(b["id"], {}).update(b)
    submitted = set(s for b in batches.values() for s in b.get("ids", []))
    labels_f = (out_dir / "labels.jsonl").open("a")
    est = args.est
    rnd = random.Random(11)
    while True:
        pending = 0.0
        for bid, b in list(batches.items()):
            if b.get("collected"):
                continue
            try:
                info = json.loads(http("GET", f"/v1/messages/batches/{bid}"))
            except Exception as e:
                print("poll error", str(e)[:200], flush=True)
                pending += len(b["ids"]) * est
                continue
            if info["processing_status"] != "ended":
                pending += len(b["ids"]) * est
                continue
            raw = http("GET", f"/v1/messages/batches/{bid}/results", timeout=900).decode()
            ok = bad = 0
            for line in raw.splitlines():
                if not line.strip():
                    continue
                r = json.loads(line)
                cid = r["custom_id"]
                sid, dup = cid.split("_")[0], cid.endswith("_dup")
                res = r["result"]
                if res["type"] != "succeeded":
                    bad += 1
                    continue
                m = res["message"]
                u = m.get("usage", {})
                rec = {"id": cid, "cost": usage_cost(u), "in": u.get("input_tokens"), "out": u.get("output_tokens"), "batch": bid}
                ledger.add(rec)
                try:
                    write_label(labels_f, sid, sheets[sid], parse(m, sheets[sid]["rows"]), rec, dup)
                    ok += 1
                except Exception:
                    bad += 1
            b["collected"] = True
            with state_p.open("a") as f:
                f.write(json.dumps({"id": bid, "collected": True}) + "\n")
            print(f"[{time.strftime('%H:%M:%S')}] collected {bid}: {ok} ok, {bad} failed", flush=True)
        if ledger.n > 100:
            est = max(0.0005, ledger.cost / ledger.n)
        if args.follow:
            sheets.update(load_sheets(args.sheets))
        new = [s for s in sheets if s not in submitted]
        room = args.budget - ledger.cost - pending
        can = int(max(0, room) / est)
        if new and can > 0:
            chunk = new[: min(args.batch_size, can)]
            reqs = []
            for sid in chunk:
                png = (Path(args.sheets) / "sheets" / f"{sid}.png").read_bytes()
                p = params(png, len(sheets[sid]["rows"]), args)
                reqs.append({"custom_id": sid, "params": p})
                if rnd.random() < args.dup:  # independent second opinion: measures label agreement
                    reqs.append({"custom_id": sid + "_dup", "params": p})
            try:
                info = json.loads(http("POST", "/v1/messages/batches", {"requests": reqs}, timeout=900))
                b = {"id": info["id"], "ids": chunk, "t": time.time()}
                batches[info["id"]] = b
                with state_p.open("a") as f:
                    f.write(json.dumps(b) + "\n")
                submitted.update(chunk)
                print(f"[{time.strftime('%H:%M:%S')}] submitted {info['id']} with {len(reqs)} (spent ${ledger.cost:.2f}, pending ~${pending:.2f}, est/sheet ${est:.4f})", flush=True)
                continue
            except Exception as e:
                print("submit error", str(e)[:300], flush=True)
                time.sleep(60)
        open_b = [b for b in batches.values() if not b.get("collected")]
        print(f"[{time.strftime('%H:%M:%S')}] spent ${ledger.cost:.3f} on {ledger.n} | open batches {len(open_b)} | unsubmitted {len(new)}", flush=True)
        if not open_b and (not new or can <= 0) and not args.follow:
            break
        if room <= 0 and not open_b:
            print("budget reached", flush=True)
            break
        time.sleep(args.poll)
    print(f"DONE ${ledger.cost:.3f} {ledger.n}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sheets", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--budget", type=float, default=25.0)
    ap.add_argument("--effort", default="high")
    ap.add_argument("--max-out", type=int, default=16000)
    ap.add_argument("--pilot", type=int, default=0, help="label N sheets with the realtime API and stop")
    ap.add_argument("--batch-size", type=int, default=300)
    ap.add_argument("--dup", type=float, default=0.02)
    ap.add_argument("--est", type=float, default=0.002)
    ap.add_argument("--poll", type=int, default=60)
    ap.add_argument("--follow", action="store_true", help="keep picking up new sheets")
    args = ap.parse_args()
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    sheets = load_sheets(args.sheets)
    ledger = Ledger(out_dir / ("ledger_pilot.jsonl" if args.pilot else "ledger.jsonl"))
    print(len(sheets), "sheets", flush=True)
    if args.pilot:
        return pilot(args, sheets, ledger, out_dir)
    batch(args, sheets, ledger, out_dir)


if __name__ == "__main__":
    main()