#!/usr/bin/env python3
"""Upscale QA: download before / after pairs (Kenerate upscale_image or upscale_video results), verify the new size,
measure sharpness gain, and write side-by-side + 100% crop comparison images and a report.

Examples
  python compare.py --pair https://.../orig.jpg::https://.../upscaled.png --pair a.jpg::a_4k.png
  python compare.py --manifest pairs.json     # [{"before": "...", "after": "...", "label": "hero"}]
  python compare.py --pair clip.mp4::clip_1080p.mp4   # videos: compares a frame from the middle
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import kmedia as K  # noqa: E402


def as_image(p: Path, tmp: Path, tag: str):
    if K.kind_of(str(p)) == "video":
        info = K.probe(p)
        fr = K.extract_frame(p, info["duration"] / 2, tmp / f"{tag}_mid.png")
        return K.open_image(fr).convert("RGB"), info
    im = K.open_image(p).convert("RGB")
    return im, {"width": im.width, "height": im.height}


def sharp(im) -> float:
    from PIL import ImageFilter, ImageStat
    g = im.convert("L")
    g.thumbnail((1024, 1024))
    return ImageStat.Stat(g.filter(ImageFilter.FIND_EDGES)).var[0]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--pair", action="append", default=[], help="BEFORE::AFTER (URLs or files)")
    ap.add_argument("--manifest", help='[{"before": "...", "after": "...", "label": "..."}]')
    ap.add_argument("--download-only", action="store_true", help="just download the 'after' files with clean names")
    ap.add_argument("--name", default="upscale")
    ap.add_argument("--out-dir", default="kenerate-out/upscale-batch")
    a = ap.parse_args()
    pairs = []
    if a.manifest:
        pairs += [(d["before"], d["after"], d.get("label") or f"{i:02d}") for i, d in enumerate(json.loads(Path(a.manifest).read_text(encoding="utf-8")), 1)]
    for i, p in enumerate(a.pair, len(pairs) + 1):
        b, _, af = p.partition("::")
        pairs.append((b.strip(), af.strip(), f"{i:02d}"))
    if not pairs:
        K.fail("No pairs.", "Pass --pair BEFORE::AFTER (repeat) or --manifest pairs.json.")
    st = K.Status(3, "upscale-batch")
    K.require(status=st)
    from PIL import Image, ImageDraw
    tmp = K.workdir()
    out = Path(a.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    st.step(f"Fetching {len(pairs)} pair(s)")
    got = []
    for b, af, lab in pairs:
        bp = K.fetch(b, tmp, f"before_{K.safe_name(lab)}", st) if b else None
        ap_ = K.fetch(af, out, f"{K.safe_name(a.name)}_{K.safe_name(lab)}", st)
        got.append((bp, ap_, lab))

    st.step("Comparing")
    rows = []
    for i, (bp, ap_, lab) in enumerate(got, 1):
        after, ainfo = as_image(ap_, tmp, f"a{i}")
        row = {"label": lab, "after": str(Path(ap_).resolve()), "after_size": f"{ainfo['width']}x{ainfo['height']}"}
        if bp is not None and not a.download_only:
            before, binfo = as_image(bp, tmp, f"b{i}")
            scale = round(after.width / max(before.width, 1), 2)
            up = before.resize(after.size, Image.BICUBIC)
            gain = round(sharp(after) / max(sharp(up), 1e-6), 2)
            row.update({"before_size": f"{binfo['width']}x{binfo['height']}", "scale": scale, "sharpness_vs_bicubic": gain,
                        "ok": scale > 1.05})
            if scale <= 1.05:
                st.warn(f"{lab}: the 'after' is not bigger ({row['before_size']} -> {row['after_size']})")
            # side by side at display size + 100% centre crop
            h = 540
            l = before.resize((int(before.width * h / before.height), h), Image.BICUBIC)
            r = after.resize((int(after.width * h / after.height), h), Image.LANCZOS)
            cs = min(400, after.width // 2, after.height // 2)
            cx, cy = after.width // 2, after.height // 2
            ca = after.crop((cx - cs // 2, cy - cs // 2, cx + cs // 2, cy + cs // 2))
            cb = up.crop((cx - cs // 2, cy - cs // 2, cx + cs // 2, cy + cs // 2))
            W = l.width + r.width + 30
            sheet = Image.new("RGB", (max(W, cs * 2 + 30), h + cs + 80), (20, 20, 20))
            sheet.paste(l, (10, 30))
            sheet.paste(r, (l.width + 20, 30))
            sheet.paste(cb, (10, h + 60))
            sheet.paste(ca, (cs + 20, h + 60))
            d = ImageDraw.Draw(sheet)
            f = K.load_font(18, "regular")
            d.text((10, 6), f"before {row['before_size']}", font=f, fill=(230, 230, 230))
            d.text((l.width + 20, 6), f"after {row['after_size']}  x{scale}", font=f, fill=(230, 230, 230))
            d.text((10, h + 36), "100% crop: bicubic", font=f, fill=(200, 200, 200))
            d.text((cs + 20, h + 36), f"100% crop: upscaled (sharpness x{gain})", font=f, fill=(200, 200, 200))
            cp = out / f"{K.safe_name(a.name)}_{K.safe_name(lab)}_compare.jpg"
            sheet.save(cp, quality=90)
            row["compare"] = str(cp.resolve())
        rows.append(row)
        st.progress(i * 100 / len(got), lab)

    st.step("Report")
    (out / f"{K.safe_name(a.name)}_report.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    K.finish({"pairs": rows, "elapsed_s": st.elapsed()})


if __name__ == "__main__":
    K.main_guard(main)