#!/usr/bin/env python3
"""
Upscale + clean up the qt images in public/assets/qts/.

If the anime ESRGAN model (scripts/anime_x4.onnx) is present, it does a real
AI 4x upscale (crisp line art, defined detail). Otherwise it falls back to a
LANCZOS upscale + sharpen. Either way it then cleans the alpha edge and trims.

Overwrites files in place. Run:  python scripts/enhance_qts.py

Model: RealESR anime x4 (ONNX), BSD-3. Runs on CPU via onnxruntime.
"""

import sys
from pathlib import Path
from PIL import Image, ImageFilter

sys.path.insert(0, str(Path(__file__).resolve().parent))
import fetch_qts as f  # reuse refine() + trim()

HERE = Path(__file__).resolve().parent
QDIR = HERE.parent / "public" / "assets" / "qts"
MODEL_PATH = HERE / "anime_x4.onnx"
TARGET = 900        # px on the long side of the final image
TILE, PAD = 256, 16

# --- try to load the AI upscaler ---
_sess = None
try:
    import numpy as np
    import onnxruntime as ort
    if MODEL_PATH.exists():
        _sess = ort.InferenceSession(str(MODEL_PATH), providers=["CPUExecutionProvider"])
        _in = _sess.get_inputs()[0].name
except Exception:
    _sess = None


def esrgan(img_rgb, scale=4):
    """AI 4x upscale of an RGB PIL image, tiled to keep memory/seams in check."""
    arr = np.asarray(img_rgb).astype(np.float32) / 255.0
    h, w, _ = arr.shape
    out = np.zeros((h * scale, w * scale, 3), np.float32)
    for y in range(0, h, TILE):
        for x in range(0, w, TILE):
            y0, x0 = max(0, y - PAD), max(0, x - PAD)
            y1, x1 = min(h, y + TILE + PAD), min(w, x + TILE + PAD)
            chunk = arr[y0:y1, x0:x1, :].transpose(2, 0, 1)[None]
            res = _sess.run(None, {_in: chunk})[0][0]
            res = np.clip(res, 0, 1).transpose(1, 2, 0)
            ty0, tx0 = (y - y0) * scale, (x - x0) * scale
            oh, ow = min(TILE, h - y), min(TILE, w - x)
            out[y*scale:(y+oh)*scale, x*scale:(x+ow)*scale, :] = \
                res[ty0:ty0+oh*scale, tx0:tx0+ow*scale, :]
    return Image.fromarray((out * 255).round().astype("uint8"), "RGB")


def upscale(img):
    """Return an RGBA image upscaled (AI if available, else LANCZOS+sharpen)."""
    rgb, alpha = img.convert("RGB"), img.getchannel("A")
    if _sess is not None:
        big = esrgan(rgb)
    else:
        w, h = rgb.size
        big = rgb.resize((w * 2, h * 2), Image.LANCZOS)
        big = big.filter(ImageFilter.UnsharpMask(radius=2, percent=120, threshold=2))
    alpha = alpha.resize(big.size, Image.LANCZOS)
    return Image.merge("RGBA", (*big.split(), alpha))


def cap(img, target=TARGET):
    w, h = img.size
    if max(w, h) <= target:
        return img
    s = target / max(w, h)
    return img.resize((round(w * s), round(h * s)), Image.LANCZOS)


def main():
    pngs = sorted(QDIR.glob("*.png"))
    if not pngs:
        sys.exit("No images in the folder.")
    print("mode:", "AI (anime ESRGAN x4)" if _sess else "LANCZOS fallback", "\n")
    for p in pngs:
        img = Image.open(p).convert("RGBA")
        w, h = img.size
        out = upscale(img)
        out = cap(out)
        out = f.refine(out)   # defringe + clean alpha
        out = f.trim(out)     # crop transparent margins
        out.save(p)
        print(f"{p.name}: {w}x{h} -> {out.size[0]}x{out.size[1]}")
    print("\nDone. Hard-refresh the page.")


if __name__ == "__main__":
    main()