"""Splits SuperTuxKart's data folder into the zips the play page downloads.

start: everything the menus need, every track's description (its .xml and
.music files and screenshot, so all tracks, grand prix and challenges are
listed from the start) and the story mode's own tracks in full;
tracks-<n>: the other tracks' models and textures, in grand prix order;
music: the tracks' music (the menu, results and story mode tunes stay in
start); replays: the ghost replays.
Writes <out>/<name>.zip and <out>/packs.json (which tracks each zip holds).
Usage: python3 pack.py <data dir> <out dir>
"""
import json, os, re, sys, zipfile

data, out = sys.argv[1], sys.argv[2]
os.makedirs(out, exist_ok=True)
STORY = {"overworld", "tutorial", "introcutscene", "introcutscene2", "endcutscene", "featunlocked", "gpwin", "gplose"}
# Music played away from the tracks in the zips.
START_MUSIC = {"menutheme.ogg", "race_summary.ogg", "race_win_theme.ogg", "lose_theme.ogg", "overworld.ogg"}
GROUPS = [
    ["sandtrack", "scotland", "abyss", "volcano_island", "hacienda"],
    ["cornfield_crossing", "snowtuxpeak", "ravenbridge_mansion", "zengarden", "cocoa_temple"],
    ["olivermath", "gran_paradiso_island", "lighthouse", "candela_city", "snowmountain"],
    ["minigolf", "mines", "stk_enterprise", "xr591", "black_forest"],
]
tracks = sorted(os.listdir(os.path.join(data, "tracks")))
grouped = {t for g in GROUPS for t in g}
GROUPS.append([t for t in tracks if t not in STORY and t not in grouped])

packs = {"start": [], "music": [], "replays": []}
packs.update({f"tracks-{i + 1}": [] for i in range(len(GROUPS))})
group_of = {t: f"tracks-{i + 1}" for i, g in enumerate(GROUPS) for t in g}

def screenshot(track):
    xml = open(os.path.join(data, "tracks", track, "track.xml"), encoding="utf-8", errors="replace").read()
    m = re.search(r'screenshot\s*=\s*"([^"]+)"', xml)
    return m.group(1) if m else None

shots = {t: screenshot(t) for t in tracks}
for root, _, files in os.walk(data):
    for f in files:
        path = os.path.relpath(os.path.join(root, f), data)
        parts = path.split(os.sep)
        if parts[0] == "tracks" and len(parts) > 2 and parts[1] not in STORY:
            t = parts[1]
            name = "/".join(parts[2:])
            light = len(parts) == 3 and (f.endswith((".xml", ".music")) or f == shots[t])
            packs["start" if light else group_of[t]].append(path)
        elif parts[0] == "music" and f.endswith(".ogg") and f not in START_MUSIC:
            packs["music"].append(path)
        elif parts[0] == "replay":
            packs["replays"].append(path)
        else:
            packs["start"].append(path)

STORED = (".png", ".jpg", ".ogg", ".ttf", ".otf", ".spm", ".b3dz")
for name, paths in packs.items():
    with zipfile.ZipFile(os.path.join(out, name + ".zip"), "w") as z:
        for p in sorted(paths):
            info = zipfile.ZipInfo(p.replace(os.sep, "/"), (2020, 1, 1, 0, 0, 0))
            comp = zipfile.ZIP_STORED if p.lower().endswith(STORED) else zipfile.ZIP_DEFLATED
            with open(os.path.join(data, p), "rb") as fh:
                z.writestr(info, fh.read(), compress_type=comp, compresslevel=9)
    print(name, len(paths), os.path.getsize(os.path.join(out, name + ".zip")))
json.dump({f"tracks-{i + 1}": g for i, g in enumerate(GROUPS)}, open(os.path.join(out, "packs.json"), "w"), indent=1)