"""
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)
ARGS = sys.argv[1:]


def _opt(name, default):
    return ARGS[ARGS.index(name) + 1] if name in ARGS else default


# A second pack keeps to itself:  --src kit_export_lite --prefix SciFI_Lite_Assets --json proxies_lite.json
# reads that folder's textures.json, writes textures/SciFI_Lite_Assets/** and textures/proxies_lite.json.
SRC = os.path.join(ROOT, _opt("--src", "kit_export"))
DST = os.path.join(ROOT, "textures")
PREFIX = _opt("--prefix", "")
JSON_NAME = _opt("--json", "proxies.json")
MAX_SIDE = int(ARGS[0]) if ARGS and ARGS[0].isdigit() 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.join(PREFIX, 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)
    # The full-size export is deleted once the copies are made (it is big). Run without it, this script would write
    # every texture into the list as missing, and Blender would lose them all: so it stops instead.
    gone = [p for p, rec in tex.items() if rec.get("file") and not rec.get("failed") and not rec.get("skipped")
            and not os.path.isfile(os.path.join(SRC, rec["file"]))]
    if gone and "--force" not in ARGS:
        print("STOPPED: %d of %d full-size textures are not in %s (first: %s). Bring them back with "
              "build\\ue_export_kit.ps1 -Do textures, or pass --force." % (len(gone), len(tex), SRC, gone[0]))
        sys.exit(2)
    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, JSON_NAME), "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()