#!/usr/bin/env python3
# Packs Taisei's installed game files (data/<content id>) into the zips the
# play page downloads: start, stages23 and stages456 (see build.sh). Each zip
# entry is named by its content id, as the game's resource index asks for it.
# Usage: pack.py <res-index.inc.h> <install/data> <pack-start.txt> <out dir>
import hashlib, os, re, sys, zipfile

index, data, start_list, out = sys.argv[1:5]
paths = {}  # content id -> resource paths
for line in open(index, encoding="utf-8"):
    m = re.match(r'FILE\(\s*\d+, "([0-9a-f]+)", "[^"]*", "([^"]*)"\)', line)
    if m:
        # The source path; the resource path is what follows the package folder or build dir.
        src = m[2]
        rel = src.split(".pkgdir/", 1)[1] if ".pkgdir/" in src else re.sub(r"^.*/resources/[^/]+/", "", src)
        paths.setdefault(m[1], []).append(rel)

start = set(l.strip() for l in open(start_list, encoding="utf-8") if l.strip())

def group(rel):
    if rel in start:
        return "start"
    m = re.search(r"stage(\d|ex)", rel)
    if m and m[1] != "1":
        return "stages23" if m[1] in "23" else "stages456"
    if rel.startswith("gfx/credits") or rel in ("bgm/credits.opus", "bgm/ending.opus", "bgm/bonus0.opus", "bgm/scuttle.opus"):
        return "stages456"
    return "start"

order = ["start", "stages23", "stages456"]
packs = {k: [] for k in order}
for cid, rels in sorted(paths.items()):
    packs[min((group(r) for r in rels), key=order.index)].append(cid)

os.makedirs(out, exist_ok=True)
for name in order:
    tmp = os.path.join(out, name + ".zip")
    with zipfile.ZipFile(tmp, "w") as z:
        for cid in packs[name]:
            info = zipfile.ZipInfo(cid, (1980, 1, 1, 0, 0, 0))
            raw = open(os.path.join(data, cid), "rb").read()
            # Pictures, sound and .zst files are compressed already.
            stored = all(r.endswith((".basis", ".opus", ".zst", ".webp")) for r in paths[cid])
            info.compress_type = zipfile.ZIP_STORED if stored else zipfile.ZIP_DEFLATED
            z.writestr(info, raw, compresslevel=9)
    h = hashlib.sha256(open(tmp, "rb").read()).hexdigest()[:12]
    os.replace(tmp, os.path.join(out, f"{name}-{h}.zip"))
    print(name, len(packs[name]), "files")