"""
Makes the 1K working copies of the kit's textures for Blender.

    python build/make_proxies.py [max_side]

Reads  kit_export/textures.json + kit_export/textures/**  (full-size PNG/TGA straight out of Unreal)
Writes textures/**/<same name>.png  (longest side <= max_side, default 1024) and textures/proxies.json.

Normal maps (Unreal compression TC_NORMALMAP) get their green channel flipped here, once, so the Blender
materials can use them as they are: Unreal normals are DirectX style (green down), Blender wants green up.
The pack's own files in the Unreal project are never touched.
"""
import json
import os
import sys
from multiprocessing import Pool

from PIL import Image

Image.MAX_IMAGE_PIXELS = None
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
SRC = os.path.join(ROOT, "kit_export")
DST = os.path.join(ROOT, "textures")
MAX_SIDE = int(sys.argv[1]) if len(sys.argv) > 1 else 1024


def work(item):
    path, rec = item
    rel = rec.get("file")
    if not rel or rec.get("failed") or rec.get("skipped"):
        return path, None
    src = os.path.join(SRC, rel)
    out_rel = os.path.splitext(rel[len("textures/"):])[0] + ".png"
    dst = os.path.join(DST, out_rel)
    is_normal = "NORMALMAP" in rec.get("compression", "")
    flip = is_normal and not rec.get("flip_green")
    info = {"file": out_rel.replace("\\", "/"), "srgb": bool(rec.get("srgb")), "normal": is_normal,
            "green_flipped": flip, "source_size": rec.get("size")}
    try:
        if os.path.isfile(dst) and os.path.getmtime(dst) >= os.path.getmtime(src):
            with Image.open(dst) as im:
                info["size"] = list(im.size)
                info["mode"] = im.mode
            return path, info
        with Image.open(src) as im:
            im.load()
            if im.mode in ("I;16", "I;16B", "I"):
                im = im.point(lambda v: v / 256).convert("L")
            elif im.mode not in ("L", "LA", "RGB", "RGBA"):
                im = im.convert("RGBA" if "A" in im.getbands() else "RGB")
            # an alpha channel that is fully opaque carries nothing: drop it (smaller file, no surprise in Blender)
            if im.mode in ("RGBA", "LA"):
                lo, hi = im.getchannel("A").getextrema()
                if lo == 255:
                    im = im.convert("RGB" if im.mode == "RGBA" else "L")
            w, h = im.size
            s = min(1.0, MAX_SIDE / float(max(w, h)))
            if s < 1.0:
                im = im.resize((max(1, round(w * s)), max(1, round(h * s))), Image.LANCZOS)
            if flip:
                if im.mode != "RGB":
                    im = im.convert("RGB")
                r, g, b = im.split()
                g = g.point(lambda v: 255 - v)
                im = Image.merge("RGB", (r, g, b))
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            im.save(dst, optimize=False, compress_level=6)
            info["size"] = list(im.size)
            info["mode"] = im.mode
        return path, info
    except Exception as e:  # noqa: BLE001
        return path, {"error": "%s: %s" % (type(e).__name__, e), "file": None}


def main():
    with open(os.path.join(SRC, "textures.json"), "r", encoding="utf-8") as f:
        tex = json.load(f)
    out = {}
    errors = 0
    with Pool(max(2, (os.cpu_count() or 4) - 2)) as pool:
        for n, (path, info) in enumerate(pool.imap_unordered(work, sorted(tex.items()), chunksize=4)):
            if info is None:
                continue
            out[path] = info
            if info.get("error"):
                errors += 1
                print("ERROR", path, info["error"])
            if (n + 1) % 100 == 0:
                print("%d / %d" % (n + 1, len(tex)), flush=True)
    os.makedirs(DST, exist_ok=True)
    with open(os.path.join(DST, "proxies.json"), "w", encoding="utf-8") as f:
        json.dump(out, f, indent=1, sort_keys=True)
    total = 0
    for dp, _, files in os.walk(DST):
        total += sum(os.path.getsize(os.path.join(dp, x)) for x in files)
    print("proxies: %d written, %d errors, %.0f MB in %s" % (len(out) - errors, errors, total / 1e6, DST))


if __name__ == "__main__":
    main()