"""How far apart are the three pictures of one design?

For a sample of tasks, render the ground truth three ways and compare them:

    skia     the server renderer: what every reward is computed from
    browser  the editor's DOM (Chromium, or --browser): what a computer-use agent sees
    crello   Crello's own shipped preview: what the design actually looked like

skia-vs-browser is the number that matters for the GUI surface (an agent must
not be misled by its own screen). skia-vs-crello says how faithful the
renderer is to the source data. Run against a live server:

    python scripts/parity.py --server http://127.0.0.1:8007 --n 40 --out parity.json
"""

from __future__ import annotations

import argparse
import io
import json
import statistics
import sys
import urllib.request
from pathlib import Path

import numpy as np
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from design_canvas.core.expert import html_page  # noqa: E402
from design_canvas.core.pack import load_pack  # noqa: E402
from design_canvas.core.score import ssim  # noqa: E402

SIZE = 512


def post(url, data: bytes, ctype):
    req = urllib.request.Request(url, data=data, headers={"content-type": ctype}, method="POST")
    with urllib.request.urlopen(req, timeout=20) as r:
        return json.loads(r.read())


def compare(a: Image.Image, b: Image.Image) -> dict:
    b = b.resize(a.size, Image.LANCZOS)
    x = np.asarray(a.convert("RGB"), dtype=np.float64)
    y = np.asarray(b.convert("RGB"), dtype=np.float64)
    return {"mae": float(np.abs(x - y).mean()), "ssim": ssim(x, y)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", default="http://127.0.0.1:8007")
    ap.add_argument("--split", default="test")
    ap.add_argument("--n", type=int, default=40)
    ap.add_argument("--out", default=None)
    ap.add_argument("--browser", default="chromium", choices=["chromium", "firefox", "webkit"])
    ap.add_argument("--save-examples", default=None, help="directory for side-by-side PNGs")
    args = ap.parse_args(argv)

    from playwright.sync_api import sync_playwright

    pack = load_pack()
    n = min(args.n, pack.count(args.split))
    step = max(1, pack.count(args.split) // n)
    rows = []
    with sync_playwright() as p:
        flags = ["--font-render-hinting=none", "--disable-lcd-text"] if args.browser == "chromium" else []
        browser = getattr(p, args.browser).launch(args=flags)
        page = browser.new_page(viewport={"width": 1600, "height": 1000}, device_scale_factor=1)
        for i in range(0, step * n, step):
            task = pack.at(args.split, i)
            gt = pack.gt_doc(task)
            ep = post(f"{args.server}/api/episodes",
                      json.dumps({"task_id": task["task_id"]}).encode(), "application/json")
            eid = ep["episode_id"]
            post(f"{args.server}/api/episodes/{eid}/html", html_page(gt).encode(), "text/plain")
            page.goto(f"{args.server}/editor/?episode={eid}&skin=figma")
            page.wait_for_function("window.designEditor && window.designEditor.ready")
            scale = SIZE / max(gt.width, gt.height)
            page.evaluate(f"""async () => {{
                const a = window.__app;
                a.store.select([]);
                a.stage.setZoom({scale});
                await document.fonts.ready;
                await a.store.whenIdle(3000);
                document.querySelector('.overlay').style.display = 'none';
                document.querySelector('.artboard').style.boxShadow = 'none';
            }}""")
            page.wait_for_timeout(150)
            dom = Image.open(io.BytesIO(page.locator(".artboard").screenshot())).convert("RGB")
            skia = Image.open(io.BytesIO(pack.target_png(task, SIZE))).convert("RGB")
            crello = Image.open(pack.preview_path(task)).convert("RGB")
            row = {"task_id": task["task_id"], "n_text": task["stats"]["n_text"],
                   "browser_vs_skia": compare(skia, dom), "crello_vs_skia": compare(skia, crello)}
            rows.append(row)
            print(f"{task['task_id']}  browser/skia ssim {row['browser_vs_skia']['ssim']:.3f} "
                  f"mae {row['browser_vs_skia']['mae']:.2f}   crello/skia ssim "
                  f"{row['crello_vs_skia']['ssim']:.3f} mae {row['crello_vs_skia']['mae']:.2f}")
            if args.save_examples:
                out = Path(args.save_examples)
                out.mkdir(parents=True, exist_ok=True)
                w, h = skia.size
                sheet = Image.new("RGB", (w * 3 + 16, h), "white")
                for k, im in enumerate((crello.resize((w, h)), skia, dom.resize((w, h)))):
                    sheet.paste(im, (k * (w + 8), 0))
                sheet.save(out / f"{task['task_id']}.png")
        browser.close()

    def summ(key, metric):
        vals = [r[key][metric] for r in rows]
        return {"mean": statistics.mean(vals), "min": min(vals), "max": max(vals),
                "median": statistics.median(vals)}

    summary = {
        "n": len(rows), "size": SIZE,
        "browser_vs_skia": {"ssim": summ("browser_vs_skia", "ssim"), "mae": summ("browser_vs_skia", "mae")},
        "crello_vs_skia": {"ssim": summ("crello_vs_skia", "ssim"), "mae": summ("crello_vs_skia", "mae")},
    }
    print(json.dumps(summary, indent=2))
    if args.out:
        Path(args.out).write_text(json.dumps({"summary": summary, "rows": rows}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())