"""Decisive probe: can an OPTIMIZED (PGD) perturbation move a vision encoder
enough to (a) push the decoy margin, (b) flip the local VLM, (c) survive JPEG,
(d) transfer to gpt-5.5 / gemini?

Two modes per image:
  box   -> unconstrained perturbation in a text-box region  (Nightshade ceiling)
  mask  -> perturbation confined to the decoy glyph pixels   (stealthy-text ceiling)

This gates the whole ensemble idea: if even the unconstrained optimized attack
can't flip the frontier after JPEG, more encoders won't help.
"""

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.fitness.embed import Embedder  # noqa: E402
from veil_pgd.fitness.semantic import embedding_distance  # noqa: E402
from veil_pgd.render.overlay import _load_font  # noqa: E402
from veil_pgd.robustness import scraper_sim  # noqa: E402
from veil_pgd.stealth.metrics import psnr, ssim  # noqa: E402
from veil_pgd.targets.base import LabelPrompt  # noqa: E402
from veil_pgd.targets.registry import Registry  # noqa: E402
from veil_pgd.targets.whitebox import WhiteBoxClient  # noqa: E402

from PIL import ImageDraw  # noqa: E402

ENC = "openclip:ViT-B-32"
FLIP_TAU = 0.5
DECOYS = {  # far-ish, reachable decoys per imagenette truth
    "cassette player": "jellyfish", "tench": "volcano", "church": "jellyfish",
    "chainsaw": "peacock", "English springer": "cassette player",
    "French horn": "jellyfish", "garbage truck": "flower", "gas pump": "banana",
    "golf ball": "volcano", "parachute": "octopus",
}


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


def rows(path, limit):
    r = []
    for line in Path(path).read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#"):
            p, t = line.split(",", 1)
            r.append((p.strip(), t.strip()))
    step = max(1, len(r) // limit)
    return r[::step][:limit]


def text_mask(img: Image.Image, decoy: str) -> tuple[Image.Image, list[int]]:
    """White decoy glyphs on black, bottom strip, at a readable size. Returns
    (mask, region_box). The mask marks the editable (glyph) pixels."""
    W, H = img.size
    m = Image.new("L", (W, H), 0)
    d = ImageDraw.Draw(m)
    px = max(14, int(H * 0.11))
    font = _load_font("DejaVuSans", px)
    l, t, r, b = d.multiline_textbbox((0, 0), decoy, font=font)
    tw, th = r - l, b - t
    x = max(2, (W - tw) // 2)
    y = H - th - max(2, int(H * 0.04))
    d.text((x, y), decoy, fill=255, font=font)
    return m, [0, int(H * 0.80), W, H]


def main():
    s = get_settings()
    reg = Registry(s)
    wb = WhiteBoxClient(s.klaus3_vision_service_url, timeout=180.0)
    emb = Embedder(reg.embeddings(), s.klaus3_vision_service_url)
    prompt = LabelPrompt()
    wb.load(ENC)

    import httpx
    surrogates = []
    for name, url in [("qwen-3.5-4b", s.klaus3_qwen_base_url),
                      ("gemma-4-4b", s.klaus3_gemma4b_base_url)]:
        try:
            httpx.get(url.rstrip("/") + "/models", timeout=3.0)
            surrogates.append(reg.surrogate(name))
        except Exception:
            pass
    log(f"surrogates: {[m.name for m in surrogates]}")

    imgs = rows("examples/testset.csv", 6)
    frontier_idx = {0, 2, 4}  # query paid models on 3 of the 6
    blackbox = reg.all_blackbox()
    results = []

    def surro_dist(image, truth):
        ds = []
        for m in surrogates:
            ds.append(embedding_distance(emb, m.label(image, prompt).parsed_label, truth))
        return (sum(ds) / len(ds)) if ds else 0.0

    for i, (path, truth) in enumerate(imgs):
        img = Image.open(path).convert("RGB")
        decoy = DECOYS.get(truth, "jellyfish")
        W, H = img.size
        mask, region = text_mask(img, decoy)
        row = {"image": Path(path).name, "truth": truth, "decoy": decoy, "modes": {}}
        log(f"[{i+1}/{len(imgs)}] {Path(path).name} truth={truth!r} decoy={decoy!r}")

        for mode in ("box", "mask"):
            kw = dict(model_id=ENC, eps=0.0627, steps=80, return_image=True)
            if mode == "box":
                out = wb.pgd_region(img, truth, decoy, region=region, **kw)
            else:
                out = wb.pgd_region(img, truth, decoy, region=region, mask=mask, **kw)
            adv = out["image"]
            adv_jpeg = scraper_sim(adv)
            # re-score margin after JPEG
            js = wb.score(adv_jpeg, truth, decoy, model_id=ENC, clean=img)
            rec = {
                "margin_before": round(out["margin_before"], 3),
                "margin_after": round(out["margin_after"], 3),
                "delta_margin": round(out["delta_margin"], 3),
                "margin_after_jpeg": round(js["margin"], 3),
                "psnr": round(psnr(img, adv), 1), "ssim": round(ssim(img, adv), 3),
                "editable_px": out["editable_px"],
                "surr_dist_clean": round(surro_dist(img, truth), 3),
                "surr_dist_adv": round(surro_dist(adv, truth), 3),
                "surr_dist_adv_jpeg": round(surro_dist(adv_jpeg, truth), 3),
            }
            rec["local_flip_jpeg"] = rec["surr_dist_adv_jpeg"] >= FLIP_TAU
            if i in frontier_idx:
                fr = {}
                for m in blackbox:
                    cp = m.label(img, prompt).parsed_label
                    ap = m.label(adv_jpeg, prompt).parsed_label
                    cd = embedding_distance(emb, cp, truth)
                    ad = embedding_distance(emb, ap, truth)
                    fr[m.name] = {"clean": cp, "adv": ap,
                                  "clean_dist": round(cd, 3), "adv_dist": round(ad, 3),
                                  "flip": ad >= FLIP_TAU and cd < FLIP_TAU}
                rec["frontier_jpeg"] = fr
            row["modes"][mode] = rec
            log(f"    {mode}: margin {rec['margin_before']}->{rec['margin_after']} "
                f"(jpeg {rec['margin_after_jpeg']}) psnr={rec['psnr']} "
                f"local_adv={rec['surr_dist_adv']} local_jpeg={rec['surr_dist_adv_jpeg']} "
                f"flip_jpeg={rec['local_flip_jpeg']}"
                + (f" frontier={ {k.split('/')[-1]: v['flip'] for k,v in rec['frontier_jpeg'].items()} }"
                   if 'frontier_jpeg' in rec else ""))
        results.append(row)

    Path("research/optimizer_probe.json").write_text(json.dumps(results, indent=2))

    print("\n============= OPTIMIZER PROBE SUMMARY =============")
    for mode in ("box", "mask"):
        deltas = [r["modes"][mode]["delta_margin"] for r in results]
        jdeltas = [r["modes"][mode]["margin_after_jpeg"] - r["modes"][mode]["margin_before"]
                   for r in results]
        psnrs = [r["modes"][mode]["psnr"] for r in results]
        lf = sum(r["modes"][mode]["local_flip_jpeg"] for r in results)
        print(f"{mode:>4}: mean margin push {sum(deltas)/len(deltas):+.3f} "
              f"(after JPEG {sum(jdeltas)/len(jdeltas):+.3f}); mean PSNR {sum(psnrs)/len(psnrs):.1f}dB; "
              f"local flip after JPEG {lf}/{len(results)}")
    for mode in ("box", "mask"):
        fl = {}
        for r in results:
            fr = r["modes"][mode].get("frontier_jpeg")
            if fr:
                for k, v in fr.items():
                    fl.setdefault(k, [0, 0]); fl[k][0] += v["flip"]; fl[k][1] += 1
        if fl:
            print(f"{mode:>4} frontier (after JPEG): "
                  + "; ".join(f"{k.split('/')[-1]} {v[0]}/{v[1]}" for k, v in fl.items()))
    print("==================================================")
    reg.close(); wb.close()


if __name__ == "__main__":
    main()