#!/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

DEFAULT_SPOOL = "/home/ubuntu/w3cs-lab/bundle-spool"
DEFAULT_REPLAYS_JSON = "/home/ubuntu/w3cs-lab/app/replays.json"
DEFAULT_CATALOG_JSON = "/home/ubuntu/w3cs-lab/catalog-serving.json"
DEFAULT_REPLAY_ROOT = "/home/ubuntu/w3cs-lab/replays"
PACK_PLAIN_LIMIT = 48 * 1024 * 1024
MIN_BUNDLE_BYTES = 1024
SHA1 = re.compile(r"[0-9a-f]{40}\Z")
PACK_OBJECT = re.compile(r"[0-9a-f]{64}\.zpack\Z")
BUNDLE_DIGEST = re.compile(r"[0-9a-f]{64}\Z")


def is_webp_variant_bundle(path, size):
    """Identify a standalone encoded WebP DEFINE_BLOB bundle.

    WebP quality variants use the browser's loose content digest. They must
    not join the generic replay/map pack, which would prefetch raw plus q95,
    q80, and q40 copies. This check also filters variants logged by an older
    relay before that policy was enforced at capture time.
    """
    if size < 44:
        return False
    try:
        with open(path, "rb") as source:
            header = source.read(44)
    except OSError:
        return False
    if len(header) != 44 or header[0] != 7:
        return False
    payload_size = struct.unpack_from("<I", header, 4)[0]
    encoded_size = struct.unpack_from("<I", header, 12)[0]
    return (size == payload_size + 8
            and payload_size == encoded_size + 24
            and header[32:36] == b"RIFF"
            and header[40:44] == b"WEBP")


def load_json(path, default):
    try:
        with open(path, encoding="utf-8") as source:
            return json.load(source)
    except (OSError, ValueError, TypeError):
        return default


def replay_sources(replays_json, catalog_json):
    """Return the deploy-time aliases and the full on-demand catalog.

    app/replays.json is only a small list of replays staged when the worker
    was deployed. The serving catalog is the authority for on-demand SHA-1
    replays, which can first appear long after that deploy.
    """
    rows = load_json(replays_json, [])
    aliases = {str(row.get("id")): row for row in rows
               if isinstance(row, dict) and row.get("id")}
    catalog = load_json(catalog_json, {})
    catalog_rows = catalog.get("replays", {}) if isinstance(catalog, dict) else {}
    if not isinstance(catalog_rows, dict):
        catalog_rows = {}
    return aliases, catalog_rows


def replay_metadata(replay, aliases, catalog_rows, replay_root):
    """Resolve the map/profile that produced a digest manifest.

    The materialized descriptor is the closest authority to the running
    session. The compact catalog covers on-demand replays that have not been
    copied into app/replays.json. The deploy-time alias remains a fallback for
    hand-staged names such as patch14.
    """
    alias = aliases.get(replay, {})
    catalog = catalog_rows.get(replay, {})
    descriptor = load_json(os.path.join(replay_root, replay + ".json"), {})
    sources = tuple(source for source in (descriptor, catalog, alias)
                    if isinstance(source, dict))

    def first(*names):
        for source in sources:
            for name in names:
                value = source.get(name)
                if value not in (None, ""):
                    return str(value)
        return ""

    return {
        "profile": first("engineProfile", "profile"),
        "mapPath": first("mapPath", "map"),
        "mapSha1": first("mapContentSha1", "mapSha1").lower(),
    }


def pack_key(profile, metadata):
    """Key packs by the actual map content, never only its basename."""
    if not re.fullmatch(r"[A-Za-z0-9._-]+", profile):
        return ""
    map_sha1 = metadata.get("mapSha1", "")
    if SHA1.fullmatch(map_sha1):
        return f"{profile}__map-{map_sha1}"
    map_path = metadata.get("mapPath", "").replace("/", "\\").lower()
    if not map_path:
        return ""
    # Legacy hand-staged aliases can lack a map content hash. Hash the full
    # normalized path so Maps\\SW\\Arena and Maps\\FrozenThrone\\Arena do not
    # collapse into one oversized pack merely because their basenames match.
    path_identity = hashlib.sha256(map_path.encode("utf-8")).hexdigest()
    return f"{profile}__path-{path_identity}"


def prior_state_is_usable(previous, out_dir):
    """Return true only when a saved state references intact local packs."""
    if not isinstance(previous, dict):
        return False
    fingerprint = previous.get("fingerprint")
    if (not isinstance(fingerprint, str)
            or not re.fullmatch(r"[0-9a-f]{64}", fingerprint)):
        return False
    for name in ("entryCount", "plainTotal"):
        value = previous.get(name)
        if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
            return False
    packs = previous.get("packs")
    if not isinstance(packs, list) or not packs:
        return False
    for pack in packs:
        if not isinstance(pack, dict):
            return False
        object_name = pack.get("object")
        if not isinstance(object_name, str) or not PACK_OBJECT.fullmatch(
                object_name):
            return False
        for name in ("plainBytes", "packedBytes", "entries"):
            value = pack.get(name)
            if (isinstance(value, bool) or not isinstance(value, int)
                    or value <= 0):
                return False
        pack_path = os.path.join(out_dir, object_name)
        if not os.path.isfile(pack_path):
            return False
        try:
            actual_size = os.path.getsize(pack_path)
        except OSError:
            return False
        if actual_size != pack["packedBytes"]:
            return False
    return True


def saved_built_digests(previous):
    """Return the exact digest set represented by saved packs, if recorded."""
    values = previous.get("digests") if isinstance(previous, dict) else None
    if not isinstance(values, list) or len(values) != previous.get("entryCount"):
        return None
    digests = set(values)
    if len(digests) != len(values) or any(
            not isinstance(value, str) or not BUNDLE_DIGEST.fullmatch(value)
            for value in values):
        return None
    return digests


def write_json_atomic(path, value):
    temp = path + ".tmp"
    with open(temp, "w", encoding="utf-8") as target:
        json.dump(value, target)
    os.replace(temp, path)


def replay_entry_counts(replay_digests, built_digests):
    return {replay: len(digests & built_digests)
            for replay, digests in replay_digests.items()}


def main(*, spool=None, replays_json=None, catalog_json=None,
         replay_root=None):
    spool = spool or os.environ.get("W3CS_BUNDLE_SPOOL", DEFAULT_SPOOL)
    replays_json = replays_json or os.environ.get(
        "W3CS_REPLAYS_JSON", DEFAULT_REPLAYS_JSON)
    catalog_json = catalog_json or os.environ.get(
        "W3CS_CATALOG_JSON", DEFAULT_CATALOG_JSON)
    replay_root = replay_root or os.environ.get(
        "W3CS_REPLAY_ROOT", DEFAULT_REPLAY_ROOT)
    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)

    aliases, catalog_rows = replay_sources(replays_json, catalog_json)

    # Union session logs into per-(profile, map) digest sets; remember
    # which replays feed each key.
    unions = {}
    key_replays = {}
    key_replay_digests = {}
    generated_manifests = set()
    skipped_metadata = []
    skipped_profile = []
    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("__")
        if not replay or not re.fullmatch(r"[A-Za-z0-9._-]+", replay):
            continue
        metadata = replay_metadata(
            replay, aliases, catalog_rows, replay_root)
        expected_profile = metadata.get("profile")
        if expected_profile and expected_profile != profile:
            # Catalog reroutes leave old digest logs behind. Do not combine
            # resources captured by another engine build with the current one.
            skipped_profile.append(name)
            continue
        key = pack_key(profile, metadata)
        if not key:
            skipped_metadata.append(name)
            continue
        digests = unions.setdefault(key, set())
        key_replays.setdefault(key, set()).add(replay)
        replay_digests = key_replay_digests.setdefault(
            key, {}).setdefault(replay, set())
        with open(os.path.join(manifests_dir, name)) as handle:
            for line in handle:
                line = line.strip()
                if BUNDLE_DIGEST.fullmatch(line):
                    digests.add(line)
                    replay_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
                    and not is_webp_variant_bundle(path, size)):
                entries.append((digest, size))
        if not entries:
            continue
        eligible_digests = {digest for digest, _ in entries}
        fingerprint = hashlib.sha256(
            "\n".join(d for d, _ in entries).encode()).hexdigest()
        marker = os.path.join(state_dir, key + ".built")
        previous = None
        try:
            with open(marker, encoding="utf-8") as source:
                previous = json.load(source)
        except (OSError, ValueError, TypeError):
            pass
        if not prior_state_is_usable(previous, out_dir):
            previous = None
        previous_digests = saved_built_digests(previous)
        if previous and previous.get("fingerprint") == fingerprint:
            # A matching fingerprint proves the current digest set is exactly
            # what these packs contain. Upgrade legacy state so later
            # hysteresis can report exact replay coverage without rebuilding.
            if previous_digests != eligible_digests:
                previous["digests"] = sorted(eligible_digests)
                write_json_atomic(marker, previous)
            write_manifests(out_dir, key_replays[key], key, previous["packs"],
                            replay_entry_counts(key_replay_digests[key],
                                                eligible_digests),
                            generated_manifests)
            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 and previous_digests is not None:
            prev_entries = previous.get("entryCount", 0)
            prev_plain = previous.get("plainTotal", 0)
            removed_digests = previous_digests - eligible_digests
            # A removed or newly ineligible digest must leave the next pack.
            # Growth hysteresis is safe only for additions; applying it to a
            # shrink would retain retired WebP variants forever.
            if (not removed_digests
                    and len(entries) < prev_entries * 1.10
                    and total_plain < prev_plain + 4 * 1024 * 1024):
                write_manifests(out_dir, key_replays[key], key,
                                previous["packs"], replay_entry_counts(
                                    key_replay_digests[key], previous_digests),
                                generated_manifests)
                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:
                with open(os.path.join(spool, digest + ".bundle"),
                          "rb") as source:
                    plain += source.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)
            with open(packed_path, "rb") as source:
                packed = source.read()
            os.remove(raw_path)
            os.remove(packed_path)
            object_name = hashlib.sha256(packed).hexdigest() + ".zpack"
            final = os.path.join(out_dir, object_name)
            try:
                existing_size = os.path.getsize(final)
            except OSError:
                existing_size = -1
            if existing_size != len(packed):
                temp = final + ".tmp"
                with open(temp, "wb") as handle:
                    handle.write(packed)
                os.replace(temp, final)
            packs.append({"object": object_name,
                          "plainBytes": len(plain),
                          "packedBytes": len(packed),
                          "entries": len(chunk)})
        write_json_atomic(marker, {
            "fingerprint": fingerprint, "packs": packs,
            "entryCount": len(entries), "plainTotal": total_plain,
            "digests": sorted(eligible_digests),
        })
        write_manifests(out_dir, key_replays[key], key, packs,
                        replay_entry_counts(key_replay_digests[key],
                                            eligible_digests),
                        generated_manifests)
        print(f"built {key}: {len(entries)} bundles -> "
              f"{len(packs)} pack(s)")

    prune_stale_manifests(out_dir, generated_manifests)
    collect_garbage(state_dir, out_dir)
    if skipped_metadata:
        print(f"skipped {len(skipped_metadata)} digest manifest(s) without "
              "replay map metadata", file=sys.stderr)
    if skipped_profile:
        print(f"skipped {len(skipped_profile)} stale digest manifest(s) from "
              "another engine profile", file=sys.stderr)


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


def prune_stale_manifests(out_dir, generated_manifests):
    """Remove pointers not produced from a current digest manifest.

    The uploader syncs only this worker's private pointer namespace, so this
    cannot delete another worker's manifests or shared immutable pack files.
    """
    pattern = re.compile(r"replay-[A-Za-z0-9._-]+\.json\Z")
    for entry in os.scandir(out_dir):
        if (entry.is_file(follow_symlinks=False)
                and pattern.fullmatch(entry.name)
                and entry.name not in generated_manifests):
            os.remove(entry.path)
            print(f"pruned stale replay manifest {entry.name}")


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). This
    only reclaims local storage; the publisher retains remote immutable
    pack objects so one worker cannot delete another worker's generation."""
    import time
    referenced = set()
    for name in os.listdir(state_dir):
        if not name.endswith(".built"):
            continue
        try:
            with open(os.path.join(state_dir, name),
                      encoding="utf-8") as source:
                state = json.load(source)
        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())