#!/usr/bin/env python3
"""Build per-(profile, map) resource-bundle packs for the CDN.

A cold browser currently resolves thousands of small content-addressed
bundles; a pack is a handful of large zstd files it prefetches and
unpacks into IndexedDB during the replay load screen instead.

Inputs (written by the relays into the shared spool):
  <spool>/<digest>.bundle                    bundle bodies
  <spool>/manifests/<profile>__<replay>.digests   per-session digest logs

Outputs (uploaded by w3cs-bundle-upload.sh afterwards):
  <spool>/pack-out/<sha256>.zpack            immutable pack objects
  <spool>/pack-out/replay-<replay>.json      per-replay manifest pointers

Pack binary layout (after zstd decompression):
  "W3PK" | u32 version=1 | u32 count |
  count x (32-byte digest | u32 length) | bodies concatenated in order.

Packs rebuild only when a key's digest union changes; per-replay
manifests point at the current pack set for that replay's (profile,
map). Everything is a pure function of the digest union, so reruns are
idempotent.
"""
import hashlib
import json
import os
import re
import struct
import subprocess
import sys

SPOOL = "/home/ubuntu/w3cs-lab/bundle-spool"
REPLAYS_JSON = "/home/ubuntu/w3cs-lab/app/replays.json"
PACK_PLAIN_LIMIT = 48 * 1024 * 1024
MIN_BUNDLE_BYTES = 1024


def map_slug(map_path):
    base = map_path.replace("\\", "/").rsplit("/", 1)[-1] or "nomap"
    return re.sub(r"[^A-Za-z0-9._-]", "_", base)[:80]


def main():
    manifests_dir = os.path.join(SPOOL, "manifests")
    state_dir = os.path.join(SPOOL, "pack-state")
    out_dir = os.path.join(SPOOL, "pack-out")
    os.makedirs(state_dir, exist_ok=True)
    os.makedirs(out_dir, exist_ok=True)

    try:
        rows = json.load(open(REPLAYS_JSON))
    except OSError:
        rows = []
    replay_info = {row["id"]: (row.get("profile", "auto"),
                               map_slug(row.get("map", "")))
                   for row in rows}

    # Union session logs into per-(profile, map) digest sets; remember
    # which replays feed each key.
    unions = {}
    key_replays = {}
    if not os.path.isdir(manifests_dir):
        return
    for name in os.listdir(manifests_dir):
        if not name.endswith(".digests"):
            continue
        profile, _, replay = name[:-len(".digests")].partition("__")
        info = replay_info.get(replay)
        if not info:
            continue
        key = f"{info[0]}__{info[1]}"
        digests = unions.setdefault(key, set())
        key_replays.setdefault(key, set()).add(replay)
        with open(os.path.join(manifests_dir, name)) as handle:
            for line in handle:
                line = line.strip()
                if len(line) == 64:
                    digests.add(line)

    for key, digests in sorted(unions.items()):
        # Only digests whose bodies exist and clear the size floor: tiny
        # bundles cost more index than they save.
        entries = []
        for digest in sorted(digests):
            path = os.path.join(SPOOL, digest + ".bundle")
            try:
                size = os.path.getsize(path)
            except OSError:
                continue
            if size >= MIN_BUNDLE_BYTES:
                entries.append((digest, size))
        if not entries:
            continue
        fingerprint = hashlib.sha256(
            "\n".join(d for d, _ in entries).encode()).hexdigest()
        marker = os.path.join(state_dir, key + ".built")
        previous = None
        try:
            previous = json.load(open(marker))
        except (OSError, ValueError):
            pass
        if previous and previous.get("fingerprint") == fingerprint:
            write_manifests(out_dir, key_replays[key], key, previous["packs"])
            continue
        # Rebuild hysteresis. During a map's warm-up every session grows
        # the union a little, and rebuilding each time churns out a new
        # immutable generation that returning viewers re-download almost
        # entirely. Rebuild only on meaningful growth (>=10% entries or
        # >=4 MiB plain); until then the previous generation keeps
        # serving and the handful of new digests ride the per-bundle
        # tier, which references cover regardless.
        total_plain = sum(size for _, size in entries)
        if previous:
            prev_entries = previous.get("entryCount", 0)
            prev_plain = previous.get("plainTotal", 0)
            if (len(entries) < prev_entries * 1.10
                    and total_plain < prev_plain + 4 * 1024 * 1024):
                write_manifests(out_dir, key_replays[key], key,
                                previous["packs"])
                continue

        # Chunk into packs and build each.
        packs = []
        chunk, chunk_bytes = [], 0
        chunks = []
        for digest, size in entries:
            if chunk and chunk_bytes + size > PACK_PLAIN_LIMIT:
                chunks.append(chunk)
                chunk, chunk_bytes = [], 0
            chunk.append((digest, size))
            chunk_bytes += size
        if chunk:
            chunks.append(chunk)
        for chunk in chunks:
            plain = bytearray()
            plain += b"W3PK" + struct.pack("<II", 1, len(chunk))
            for digest, size in chunk:
                plain += bytes.fromhex(digest) + struct.pack("<I", size)
            for digest, size in chunk:
                plain += open(os.path.join(SPOOL, digest + ".bundle"),
                              "rb").read()
            raw_path = os.path.join(out_dir, f".building-{key}.plain")
            with open(raw_path, "wb") as handle:
                handle.write(plain)
            packed_path = raw_path + ".zst"
            # Level 10 on two threads: level 19 across every core burned a full
            # vCPU per pack for a few percent of size, and packs share the
            # box with live game seats. Uploaded once, cached forever.
            subprocess.run(["zstd", "-10", "-T2", "-q", "-f", raw_path,
                            "-o", packed_path], check=True)
            packed = open(packed_path, "rb").read()
            os.remove(raw_path)
            os.remove(packed_path)
            object_name = hashlib.sha256(packed).hexdigest() + ".zpack"
            final = os.path.join(out_dir, object_name)
            if not os.path.exists(final):
                temp = final + ".tmp"
                with open(temp, "wb") as handle:
                    handle.write(packed)
                os.rename(temp, final)
            packs.append({"object": object_name,
                          "plainBytes": len(plain),
                          "packedBytes": len(packed),
                          "entries": len(chunk)})
        json.dump({"fingerprint": fingerprint, "packs": packs,
                   "entryCount": len(entries),
                   "plainTotal": total_plain},
                  open(marker, "w"))
        write_manifests(out_dir, key_replays[key], key, packs)
        print(f"built {key}: {len(entries)} bundles -> "
              f"{len(packs)} pack(s)")

    collect_garbage(state_dir, out_dir)


def write_manifests(out_dir, replays, key, packs):
    manifest = {"version": 1, "key": key, "packs": packs}
    encoded = json.dumps(manifest, separators=(",", ":"))
    for replay in sorted(replays):
        if not re.fullmatch(r"[A-Za-z0-9._-]+", replay):
            continue
        path = os.path.join(out_dir, f"replay-{replay}.json")
        try:
            if open(path).read() == encoded:
                continue
        except OSError:
            pass
        temp = path + ".tmp"
        with open(temp, "w") as handle:
            handle.write(encoded)
        os.rename(temp, path)


def collect_garbage(state_dir, out_dir):
    """Superseded pack generations: every rebuild writes a new immutable
    object and orphans the previous one. Keep packs referenced by any
    current key state, plus anything younger than 24 h (a client can hold
    a cached manifest pointing at an old pack for minutes, not days).
    The uploader syncs these deletions to R2."""
    import time
    referenced = set()
    for name in os.listdir(state_dir):
        if not name.endswith(".built"):
            continue
        try:
            state = json.load(open(os.path.join(state_dir, name)))
        except (OSError, ValueError):
            continue
        for pack in state.get("packs", []):
            referenced.add(pack.get("object"))
    cutoff = time.time() - 24 * 3600
    for name in os.listdir(out_dir):
        if not name.endswith(".zpack") or name in referenced:
            continue
        path = os.path.join(out_dir, name)
        try:
            if os.path.getmtime(path) < cutoff:
                os.remove(path)
                print(f"pruned superseded pack {name}")
        except OSError:
            pass


if __name__ == "__main__":
    sys.exit(main())