#!/usr/bin/env python3
"""Build launcher zip packages and the matching <gpu>.game.json manifest.

Build:
  python3 make_zips.py build <folder> --gpu all --out out [--limit-mb 250]
  (--gpu dxt / etc / pvr builds just one manifest)

Textures in the texdb folder are split per GPU: a file whose name contains
.dxt. / .pvr. / .etc. goes only into that GPU's texdb zips; files with no
marker (like mobile.txt) go into every GPU's zips.

Exception: texdb/player, texdb/playerhi and texdb/menu always use the .pvr
files (plus unmarked files) for every GPU. They are built once into
files/texdb/texdb-pvr*.zip and shared by all three manifests. Change the
list with --pvr-folders.

Verify the checksum method against an official zip (compare with its manifest):
  python3 make_zips.py verify anim.zip

<folder> is the game "files" folder, or its parent if it contains "files".

Output layout (upload `zip_data` to your server under /game/release/):
  out/zip_data/<gpu>.game.json
  out/zip_data/data/files/<group>/<name>.zip

Every zip entry is stored as files/<path> (like the official zips), so the
archives extract into the parent of the game's files folder.
"""
import argparse
import hashlib
import json
import os
import shutil
import sys
import zipfile
import zlib

MAX_ZIP = 499_999_999           # parser rejects zips of 500 MB or more
FIXED_TIME = (2026, 1, 1, 0, 0, 0)  # fixed so identical input gives identical zips
DEFAULT_EXCLUDES = {"samp/settings.json", "modloader", "cleo/cleo.log"}
GPU_FOLDERS = {"texdb"}         # these groups get a <gpu> subfolder (like official)
PVR_ONLY_DEFAULT = ["player", "playerhi", "menu"]  # texdb subfolders that always use .pvr


def excluded(rel, excludes):
    r = rel.lower()
    return any(r == e or r.startswith(e + "/") for e in excludes)


def collect(root, excludes):
    groups = {}
    for dp, dns, fns in os.walk(root):
        dns.sort()
        for fn in sorted(fns):
            ab = os.path.join(dp, fn)
            rel = os.path.relpath(ab, root).replace(os.sep, "/")
            if excluded(rel, excludes) or not os.path.isfile(ab):
                continue
            top = rel.split("/", 1)[0] if "/" in rel else ""
            groups.setdefault(top, []).append((ab, rel))
    return groups


def split(files, limit):
    parts, cur, size = [], [], 0
    for ab, rel in files:
        sz = os.path.getsize(ab)
        if cur and size + sz > limit:
            parts.append(cur)
            cur, size = [], 0
        cur.append((ab, rel))
        size += sz
    if cur:
        parts.append(cur)
    return parts


def write_zip(path, items):
    """Write files with files/<relative-path> names, without directory entries."""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"

    with zipfile.ZipFile(
        tmp, "w", zipfile.ZIP_DEFLATED, compresslevel=6
    ) as z:
        for ab, rel in items:
            # Keep the official files/ prefix on actual files, but do not add
            # separate empty directory records such as "files/" or "files/anim/".
            zi = zipfile.ZipInfo("files/" + rel, FIXED_TIME)
            zi.compress_type = zipfile.ZIP_DEFLATED
            zi.external_attr = 0o100644 << 16

            big = os.path.getsize(ab) >= (1 << 31) - 1
            with open(ab, "rb") as src, z.open(
                zi, "w", force_zip64=big
            ) as dst:
                shutil.copyfileobj(src, dst, 1 << 20)

    if os.path.exists(path):
        if checksums(path) != checksums(tmp):
            os.remove(tmp)
            sys.exit(
                f"Collision: {path} already exists with different content.\n"
                "Use a separate --out per GPU, or delete the old zip."
            )
        os.remove(tmp)
    else:
        os.replace(tmp, path)


def checksums(path):
    md5, crc = hashlib.md5(), 0
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            md5.update(chunk)
            crc = zlib.crc32(chunk, crc)
    return md5.hexdigest(), crc & 0xFFFFFFFF


def to_nodes(tree):
    kids = []
    for name, sub in tree["dirs"].items():
        data = to_nodes(sub)
        kids.append({"type": "dir", "name": name,
                     "size": sum(k["size"] for k in data), "data": data})
    kids.extend(tree["files"])
    kids.sort(key=lambda k: k["name"])
    return kids


MARKERS = {"dxt", "pvr", "etc"}


def gpu_ok(name, gpu):
    """Keep a file if it has this GPU's marker or no GPU marker at all.

    A marker is a dot-separated part of the file name, e.g. accessories.dxt.dat
    is a dxt file, mobile.txt is shared.
    """
    parts = set(name.lower().split("."))
    return not ((parts & MARKERS) - {gpu})


def build_group(zip_root, base, loc, items, limit):
    """Write the zips for one group and return (location dirs, leaf entries)."""
    parts = split(items, limit)
    leaves = []
    for i, part in enumerate(parts, 1):
        name = f"{base}.zip" if len(parts) == 1 else f"{base}-part-{i}.zip"
        rel_path = "/".join(["files"] + loc + [name])
        if len(rel_path) > 70:
            print(f"WARNING: long path may break package ids: {rel_path}")
        out_path = os.path.join(zip_root, "data", *rel_path.split("/"))
        print(f"zip {rel_path}  ({len(part)} files)", file=sys.stderr)
        write_zip(out_path, part)
        size = os.path.getsize(out_path)
        if not 1 <= size <= MAX_ZIP:
            sys.exit(f"{rel_path} is {size} bytes; must be under 500 MB. "
                     "Lower --limit-mb, or a single file is too big to fit.")
        md5, crc = checksums(out_path)
        leaves.append({"type": "file", "name": name,
                       "date_change": int(max(os.path.getmtime(ab) for ab, _ in part)),
                       "size": size, "hash": md5, "crc32": crc})
    return loc, leaves


def cmd_build(a):
    root = os.path.abspath(a.folder)
    if os.path.isdir(os.path.join(root, "files")):
        root = os.path.join(root, "files")
    if not os.path.isdir(root):
        sys.exit(f"Not a folder: {root}")
    excludes = set(a.exclude) | (set() if a.no_default_excludes else DEFAULT_EXCLUDES)
    limit = int(a.limit_mb * 1024 * 1024)
    zip_root = os.path.join(a.out, "zip_data")
    gpus = ["dxt", "etc", "pvr"] if a.gpu == "all" else [a.gpu]
    groups = collect(root, excludes)
    shared = {}  # zips that do not depend on the GPU are built once

    pvr_dirs = {d.lower() for d in a.pvr_folders}

    def sub_of(rel):
        segs = rel.split("/")
        return segs[1].lower() if len(segs) > 2 else ""

    for gpu in gpus:
        tree = {"dirs": {}, "files": []}

        def add(key, builder):
            if key not in shared:
                shared[key] = builder()
            loc, leaves = shared[key]
            node = tree
            for d in loc:
                node = node["dirs"].setdefault(d, {"dirs": {}, "files": []})
            node["files"].extend(leaves)

        for group in sorted(groups):
            items = groups[group]
            if group not in GPU_FOLDERS:
                if items:
                    base = "base" if group == "" else group
                    loc = [] if group == "" else [group]
                    add((group, None), lambda: build_group(
                        zip_root, base, loc, items, limit))
                continue
            # GPU folder (texdb): fixed-pvr subfolders vs per-GPU files
            fixed = [x for x in items if sub_of(x[1]) in pvr_dirs
                     and gpu_ok(os.path.basename(x[1]), "pvr")]
            rest = [x for x in items if sub_of(x[1]) not in pvr_dirs
                    and gpu_ok(os.path.basename(x[1]), gpu)]
            if gpu == gpus[0]:
                have = {x[1].lower() for x in items}
                for _, rel in items:
                    if sub_of(rel) not in pvr_dirs:
                        continue
                    parts = rel.split(".")
                    for m in MARKERS - {"pvr"}:
                        if m in [p.lower() for p in parts[1:]]:
                            twin = ".".join(m != p.lower() and p or "pvr" for p in parts)
                            if twin.lower() not in have:
                                print(f"WARNING: {rel} has no .pvr version "
                                      f"({twin}); it is dropped from every build",
                                      file=sys.stderr)
            if fixed:
                add((group, "pvr-fixed"), lambda: build_group(
                    zip_root, f"{group}-pvr", [group], fixed, limit))
            if rest:
                add((group, gpu), lambda: build_group(
                    zip_root, group, [group, gpu], rest, limit))

        data = to_nodes(tree)
        if not data:
            sys.exit("No files found.")
        manifest = {"data": {"data": [{"type": "dir", "name": "files",
                                       "size": sum(k["size"] for k in data),
                                       "data": data}], "obb": []}}
        mpath = os.path.join(zip_root, f"{gpu}.game.json")
        with open(mpath, "w", encoding="utf-8") as f:
            json.dump(manifest, f, indent=1, ensure_ascii=False)
        print(f"Done: {mpath}  total {manifest['data']['data'][0]['size']} bytes",
              file=sys.stderr)


def cmd_verify(a):
    md5, crc = checksums(a.zipfile)
    print(f"size  = {os.path.getsize(a.zipfile)}\nmd5   = {md5}\ncrc32 = {crc}")
    with zipfile.ZipFile(a.zipfile) as z:
        names = z.namelist()
    print(f"{len(names)} entries, first ones: {names[:6]}")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    b = sub.add_parser("build")
    b.add_argument("folder")
    b.add_argument("--gpu", required=True, choices=["dxt", "etc", "pvr", "all"],
                   help="one GPU, or all to write the three manifests in one run")
    b.add_argument("--out", default="out")
    b.add_argument("--limit-mb", type=float, default=250,
                   help="max uncompressed MB per zip part (default 250)")
    b.add_argument("--pvr-folders", nargs="*", default=PVR_ONLY_DEFAULT,
                   help="texdb subfolders that always use .pvr for every GPU")
    b.add_argument("--exclude", nargs="*", default=[])
    b.add_argument("--no-default-excludes", action="store_true")
    v = sub.add_parser("verify")
    v.add_argument("zipfile")
    a = ap.parse_args()
    {"build": cmd_build, "verify": cmd_verify}[a.cmd](a)


if __name__ == "__main__":
    main()