"""Replace (dummy) teacher distributions with real Clef / Clef-Flash outputs via Cloudflare Workers AI.
BILLABLE: each request is charged per input token (see Cloudflare pricing). Default is --dry-run (no requests).
Env: CLOUDFLARE_ACCOUNT_ID, CLOUDFLARE_API_TOKEN (Workers AI permission).
Usage: python scripts/label_workers_ai.py in.jsonl out.jsonl --model clef-flash --run"""
import argparse, base64, io, json, os, sys, time
from pathlib import Path
import requests
from PIL import Image
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from glint.data import load_jsonl, teacher_questions

MAX_BYTES = 190_000   # third-party observation: larger images overflowed on Workers AI

def data_url(path, max_side=1024):
    img = Image.open(path).convert("RGB")
    while max_side >= 256:
        im = img.copy(); im.thumbnail((max_side, max_side)); buf = io.BytesIO(); im.save(buf, "JPEG", quality=85)
        if buf.tell() <= MAX_BYTES:
            return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
        max_side = int(max_side * 0.8)
    raise ValueError(f"{path}: too large")

def answers_to_teacher(answers):
    t = {}
    for qid, a in answers.items():
        if a.get("type") == "noul" or "noul" in a:
            p = float(a["noul"]); t[qid] = {"true": p, "false": 1 - p}
        else:
            t[qid] = {str(k): float(v) for k, v in a["probabilities"].items()}
    return t

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("inp"); ap.add_argument("out")
    ap.add_argument("--model", default="clef-flash", choices=["clef-flash", "clef"])
    ap.add_argument("--run", action="store_true", help="actually send billable requests (default: dry run)")
    ap.add_argument("--sleep", type=float, default=0.0)
    args = ap.parse_args()
    rows = load_jsonl(args.inp)
    acct, tok = os.environ.get("CLOUDFLARE_ACCOUNT_ID"), os.environ.get("CLOUDFLARE_API_TOKEN")
    url = f"https://api.cloudflare.com/client/v4/accounts/{acct}/ai/run/@cf/cloudflare/{args.model}"
    with open(args.out, "w", encoding="utf-8") as f:
        for ex in rows:
            base = Path(ex.pop("_base"))
            qs = teacher_questions(ex["questions"])            # `box` questions: gold comes from ground truth, not Clef
            if not qs:
                if args.run:
                    f.write(json.dumps(ex, ensure_ascii=False) + "\n")
                continue
            body = {"model": args.model, "state": ex.get("state", ""), "questions": qs,
                    "images": [data_url(p if Path(p).is_absolute() else base / p) for p in ex.get("images", [])][:4]}
            if not args.run:
                print(f"[dry-run] {ex['id']}: body {len(json.dumps(body))/1024:.0f} KB"); continue
            if not (acct and tok):
                sys.exit("set CLOUDFLARE_ACCOUNT_ID and CLOUDFLARE_API_TOKEN")
            r = requests.post(url, headers={"Authorization": f"Bearer {tok}"}, json=body, timeout=120)
            r.raise_for_status(); res = r.json()["result"]
            ex["teacher"] = answers_to_teacher(res["answers"])
            ex.setdefault("meta", {}).update({"teacher_model": f"@cf/cloudflare/{args.model}", "teacher_is_dummy": False,
                                              "teacher_usage": res.get("usage")})
            f.write(json.dumps(ex, ensure_ascii=False) + "\n"); time.sleep(args.sleep)

if __name__ == "__main__":
    main()