"""Re-query GPT-5.5 only (uncapped output) on the v021_base run and rescore.

GPT-5.5 is a reasoning model; the old 24-token cap sometimes truncated it before it
emitted a visible label, showing up as an empty answer. We now leave max_tokens unset.
This refreshes GPT-5.5's clean labels (cache) and adv labels (frontier_raw.json) for
runs/v021_base while leaving Gemini's cached results untouched, then rescoring.
"""
from __future__ import annotations

import io
import json
import sys
import time
from pathlib import Path

from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from veil_pgd.config import get_settings  # noqa: E402
from veil_pgd.targets.base import LabelPrompt  # noqa: E402
from veil_pgd.targets.registry import Registry  # noqa: E402

RUN = Path("runs/v021_base")
IMAGES = Path("examples/testset60")
CACHE = Path("runs/frontier_clean.json")


def jpeg(im: Image.Image, q: int = 85) -> Image.Image:
    b = io.BytesIO()
    im.convert("RGB").save(b, "JPEG", quality=q)
    b.seek(0)
    return Image.open(b).convert("RGB")


def log(m: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)


def main() -> None:
    s = get_settings()
    reg = Registry(s)
    prompt = LabelPrompt()  # max_tokens now None -> uncapped
    gpt = next(m for m in reg.all_blackbox() if "gpt-5.5" in m.name)
    name = gpt.name

    cache = json.loads(CACHE.read_text())
    raw = json.loads((RUN / "frontier_raw.json").read_text())

    for fname in list(cache):
        clean_img = jpeg(Image.open(IMAGES / fname).convert("RGB"))
        lbl = gpt.label(clean_img, prompt).parsed_label or ""
        cache[fname]["models"][name] = lbl
        log(f"clean {fname}: {lbl!r}")
    CACHE.write_text(json.dumps(cache, indent=2))

    for rec in raw:
        fname = rec["image"]
        stem = Path(fname).stem
        adv = jpeg(Image.open(RUN / "adv" / f"{stem}.png").convert("RGB"))
        adv_lbl = gpt.label(adv, prompt).parsed_label or ""
        clean_lbl = cache[fname]["models"][name]
        rec["models"][name] = {"clean_label": clean_lbl, "adv_label": adv_lbl}
        log(f"adv {fname}: {clean_lbl!r}->{adv_lbl!r}")
    (RUN / "frontier_raw.json").write_text(json.dumps(raw, indent=2))
    reg.close()
    log("done; run: python scripts/eval_frontier_v02.py --rescore --runs runs/v021_base")


if __name__ == "__main__":
    main()