#!/usr/bin/env python3
"""Builds hd_ui.pak: the 40250 UI textures at twice their resolution, for src/host/hd_ui.cpp.

The phone lays the 40250 UI out 600 logical pixels high and magnifies it about 1.8x onto the screen, so
the 1x art (32x32 icons, board frames) looks soft. This reads every UI image the client would load
(d:/ymir work/ui/, icon/, locale/<locale>/) from the registered packs (first registered pack wins, as
CEterPackManager), upscales it 4x with Real-ESRGAN (realesrgan-x4plus, ncnn-vulkan build) and
Lanczos-reduces it to exactly 2x. The engine keeps the 1x size for layout and UVs; the host only
samples finer texels, so nothing about the 40250 UI changes but the texel density.

  scripts/build_hd_ui.py --esrgan DIR [--client build/zh-locale/Client] [--locale zh] [--out FILE]

DIR holds the realesrgan-ncnn-vulkan binary and its models/ (github.com/xinntao/Real-ESRGAN releases,
v0.2.5.0). Neither the tool nor the pak enters git: the output is Ymir art (see .gitignore).

Grain: realesrgan-x4plus smooths the stone and paper fills of the boards; two rounds of back-projection
(the source minus the result box-reduced to 1x, enlarged and added) put that texture back.

Edges: each image is padded before upscaling so the network sees context at the border. An image whose
opposite edges continue each other (a tiling board fill) is padded by wrapping, so its repeats meet
without a seam; others by reflection. Colours under fully transparent pixels are bled in from their
opaque neighbours first, so the upscaled alpha edges carry no dark fringe.

Pak format (little endian): "MTHD", u32 version (1), u32 count, then count x {u16 name length, name
bytes (the pack path, lowercase), u32 offset, u32 size}, then the PNG payloads at their offsets.
"""

from __future__ import annotations

import argparse
import io
import json
import os
import shutil
import struct
import subprocess
import sys
import tempfile
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

import numpy as np
from PIL import Image, ImageFilter

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO / "tools/epk_scan"))
sys.path.insert(0, str(REPO / ".agents/skills/metin2-40250-parity-audit/scripts"))
import epk_scan  # noqa: E402
import refroot  # noqa: E402

IMAGE_EXTS = (".tga", ".dds", ".png", ".jpg", ".bmp")
MAX_SIDE = 1024  # larger images (none in the UI dirs today) are left at 1x
PAD = 8          # source pixels of context around each image
MODEL = "realesrgan-x4plus"
BACK_PROJECT_ROUNDS = 2
VERSION = 1


def pack_entries(client: Path):
    """{pack path: reader} for the first registered pack holding each name."""
    _, keys, order, packs, _ = epk_scan.load(client, refroot.reference_root())
    seen = {}
    for name, _folder in order:
        p = packs[name]
        if p["missing"]:
            continue
        for e in p["entries"]:
            seen.setdefault(e["name"], (p["epk"], e))
    return keys, seen


def read_entry(keys, epk: Path, e: dict) -> bytes:
    with open(epk, "rb") as f:
        f.seek(e["pos"])
        stored = f.read(e["size"])
        if e["type"] == 0:
            return stored
        f.seek(e["pos"])
        obj = f.read(epk_scan.lz_object_len(stored[:16]))
    return epk_scan.lz_object(obj, keys["s_adwEterPackSecurityKey"] if e["type"] == 2 else None)


def wanted(name: str, locale: str) -> bool:
    if not name.endswith(IMAGE_EXTS):
        return False
    return (name.startswith("d:/ymir work/ui/") or name.startswith("icon/")
            or name.startswith(f"locale/{locale}/"))


def tiles(a: np.ndarray, axis: int) -> bool:
    """Whether the image continues across the edges of `axis` (1: left/right, 0: top/bottom)."""
    if a.shape[axis] < 4 or a[..., 3].min() < 255:
        return False
    rgb = a[..., :3].astype(np.int32)
    first, last = np.take(rgb, 0, axis=axis), np.take(rgb, -1, axis=axis)
    wrap = np.abs(first - last).mean()
    inner = np.abs(np.diff(rgb, axis=axis)).mean()
    return wrap <= max(2.0, 1.5 * inner)


def bleed(im: Image.Image) -> Image.Image:
    """Fills the colour of fully transparent pixels from nearby opaque ones (premultiplied blur)."""
    a = np.asarray(im).astype(np.float32)
    alpha = a[..., 3:4] / 255.0
    if alpha.min() > 0:
        return im
    fill = a[..., :3].copy()
    for radius in (2, 6, 16, 48):
        pre = Image.fromarray(np.clip(a[..., :3] * alpha, 0, 255).astype(np.uint8))
        w = Image.fromarray((alpha[..., 0] * 255).astype(np.uint8))
        pre_b = np.asarray(pre.filter(ImageFilter.GaussianBlur(radius))).astype(np.float32)
        w_b = np.asarray(w.filter(ImageFilter.GaussianBlur(radius))).astype(np.float32)[..., None]
        ok = (w_b > 0.5) & (alpha == 0)
        est = pre_b / np.maximum(w_b / 255.0, 1e-6)
        fill = np.where(ok, est, fill)
        alpha_known = (alpha > 0) | ok
        if alpha_known.all():
            break
    out = a.copy()
    out[..., :3] = np.where(alpha == 0, fill, a[..., :3])
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8), "RGBA")


def prepare(job):
    """Decodes one pack image into a padded RGBA PNG for the upscaler; returns its metadata or None."""
    idx, name, data, work = job
    try:
        im = Image.open(io.BytesIO(data))
        im.load()
    except Exception:
        return None
    im = im.convert("RGBA")
    w, h = im.size
    if max(w, h) > MAX_SIDE or min(w, h) < 2:
        return None
    a = np.asarray(im)
    tile_x, tile_y = tiles(a, 1), tiles(a, 0)
    src = np.asarray(bleed(im))
    px, py = min(PAD, w), min(PAD, h)
    src = np.pad(src, ((0, 0), (px, px), (0, 0)), mode="wrap" if tile_x else "reflect")
    src = np.pad(src, ((py, py), (0, 0), (0, 0)), mode="wrap" if tile_y else "reflect")
    Image.fromarray(src, "RGBA").save(work / "in" / f"{idx}.png", compress_level=1)
    return {"idx": idx, "name": name, "w": w, "h": h, "px": px, "py": py,
            "alpha": bool(a[..., 3].min() < 255)}


def original(meta, work: Path) -> Image.Image:
    src = Image.open(work / "in" / f"{meta['idx']}.png").convert("RGBA")
    return src.crop((meta["px"], meta["py"], meta["px"] + meta["w"], meta["py"] + meta["h"]))


def back_project(hd: Image.Image, src: Image.Image, rounds: int = BACK_PROJECT_ROUNDS) -> Image.Image:
    """Adds back what the upscale lost at the source scale, so the result box-reduces to the source.

    realesrgan-x4plus denoises, and the fine stone and paper grain of the 40250 boards is what it takes
    for noise: the frames come out sharp but the fills flat. The residual (source minus the result
    reduced to 1x), bicubic-enlarged and added, restores that grain while the network's edges stay."""
    a = np.asarray(hd).astype(np.float32)
    ref = np.asarray(src).astype(np.float32)
    for _ in range(rounds):
        down = np.asarray(Image.fromarray(np.clip(a + 0.5, 0, 255).astype(np.uint8), "RGBA").resize(src.size, Image.BOX))
        residual = ref - down.astype(np.float32)
        a += np.stack([np.asarray(Image.fromarray(residual[..., c]).resize(hd.size, Image.BICUBIC)) for c in range(4)], -1)
    return Image.fromarray(np.clip(a + 0.5, 0, 255).astype(np.uint8), "RGBA")


def finish(job):
    """Crops the padding off the 4x result and reduces it to 2x; returns (name, png bytes) or None."""
    meta, work = job
    path = work / "out" / f"{meta['idx']}.png"
    if not path.exists():
        return None
    up = Image.open(path).convert("RGBA")
    s = 4
    up = up.crop((meta["px"] * s, meta["py"] * s, (meta["px"] + meta["w"]) * s, (meta["py"] + meta["h"]) * s))
    hd = back_project(up.resize((meta["w"] * 2, meta["h"] * 2), Image.LANCZOS), original(meta, work))
    if not meta["alpha"]:
        hd.putalpha(255)
    buf = io.BytesIO()
    hd.save(buf, "PNG", compress_level=6)
    return meta["name"], buf.getvalue()


def write_pak(out: Path, items: list[tuple[str, bytes]]) -> None:
    items = sorted(items)
    offset = 12 + sum(2 + len(n.encode("latin-1")) + 8 for n, _ in items)
    header = [b"MTHD", struct.pack("<II", VERSION, len(items))]
    entries = []
    for name, data in items:
        raw = name.encode("latin-1")
        entries.append(struct.pack("<H", len(raw)) + raw + struct.pack("<II", offset, len(data)))
        offset += len(data)
    tmp = out.with_suffix(".tmp")
    with open(tmp, "wb") as f:
        f.write(b"".join(header) + b"".join(entries))
        for _, data in items:
            f.write(data)
    os.replace(tmp, out)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--esrgan", type=Path, default=os.environ.get("MT_ESRGAN"), required="MT_ESRGAN" not in os.environ,
                    help="directory with realesrgan-ncnn-vulkan and models/ (or MT_ESRGAN)")
    ap.add_argument("--client", type=Path, default=REPO / "build/zh-locale/Client")
    ap.add_argument("--locale", default="zh")
    ap.add_argument("--out", type=Path, help="default: <client>/hd_ui.pak")
    ap.add_argument("--work", type=Path, help="keep intermediates here (default: a temporary directory)")
    ap.add_argument("--reuse", action="store_true",
                    help="keep the upscaled images already in --work and only redo the 2x reduction and the pak")
    ap.add_argument("--limit", type=int, default=0, help="only the first N images (a quick look)")
    args = ap.parse_args()
    out = args.out or args.client / "hd_ui.pak"
    binary = args.esrgan / "realesrgan-ncnn-vulkan"
    if not binary.exists():
        sys.exit(f"{binary} not found")

    keys, seen = pack_entries(args.client)
    names = sorted(n for n in seen if wanted(n, args.locale))
    if args.limit:
        names = names[:args.limit]
    print(f"{len(names)} UI images in {args.client}/pack")

    if args.reuse and not (args.work and (args.work / "manifest.json").exists()):
        sys.exit("--reuse needs a --work directory from an earlier run")
    work = args.work or Path(tempfile.mkdtemp(prefix="hd_ui_"))
    try:
        if args.reuse:
            metas = json.loads((work / "manifest.json").read_text())
        else:
            for sub in ("in", "out"):
                shutil.rmtree(work / sub, ignore_errors=True)
                (work / sub).mkdir(parents=True)
            jobs = [(i, n, read_entry(keys, *seen[n]), work) for i, n in enumerate(names)]
            with ProcessPoolExecutor() as pool:
                metas = [m for m in pool.map(prepare, jobs, chunksize=16) if m]
            print(f"prepared {len(metas)} (skipped {len(names) - len(metas)} undecodable or oversized)")
            (work / "manifest.json").write_text(json.dumps(metas))
            subprocess.run([str(binary), "-i", str(work / "in"), "-o", str(work / "out"), "-n", MODEL, "-s", "4",
                            "-f", "png"], cwd=args.esrgan, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        with ProcessPoolExecutor() as pool:
            items = [r for r in pool.map(finish, [(m, work) for m in metas], chunksize=16) if r]
        write_pak(out, items)
        print(f"wrote {out}: {len(items)} images, {out.stat().st_size / 2**20:.1f} MB")
    finally:
        if not args.work:
            shutil.rmtree(work, ignore_errors=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())