#!/bin/bash
# Push new content-addressed resource bundles from the shared spool to R2
# (served through cdn.war3replays.com). Bundles are immutable: uploaded
# once, cached forever. A marker in <spool>/uploaded/ tells the relays a
# digest is CDN-backed, which lets them send 36-byte references instead
# of inline bundle bytes even to browsers that never saw the content.
#
# Runs from w3cs-bundle-upload.timer. Credentials: /etc/w3cs/rclone.conf
# (root-only; object-scoped R2 keys).
set -euo pipefail

# Environment overrides keep the production defaults explicit and let the
# uploader run against isolated local fixtures in its regression tests.
W3CS_UPLOAD_SPOOL=${W3CS_BUNDLE_SPOOL:-/home/ubuntu/w3cs-lab/bundle-spool}
W3CS_UPLOAD_MARK="$W3CS_UPLOAD_SPOOL/uploaded"
W3CS_UPLOAD_REMOTE=${W3CS_BUNDLE_REMOTE:-r2:war3replays/resource-bundle}
W3CS_UPLOAD_PACK_REMOTE=${W3CS_BUNDLE_PACK_REMOTE:-r2:war3replays/resource-bundle-packs/}
W3CS_UPLOAD_RCLONE=${W3CS_RCLONE:-rclone}
W3CS_UPLOAD_RCLONE_CONFIG=${W3CS_RCLONE_CONFIG:-/etc/w3cs/rclone.conf}
W3CS_UPLOAD_PACK_BUILDER=${W3CS_PACK_BUILDER:-/usr/local/sbin/w3cs-pack-build.py}
W3CS_UPLOAD_PACK_PUBLISHER=${W3CS_PACK_PUBLISHER_ID:-}
W3CS_UPLOAD_RECONCILE_SECONDS=${W3CS_BUNDLE_RECONCILE_SECONDS:-86400}
W3CS_UPLOAD_TMP_ROOT=${W3CS_BUNDLE_TMP_ROOT:-/tmp}
W3CS_UPLOAD_RECONCILE_STAMP="$W3CS_UPLOAD_MARK/.last-r2-reconcile"

mkdir -p "$W3CS_UPLOAD_MARK"
W3CS_UPLOAD_WORK=$(mktemp -d "$W3CS_UPLOAD_TMP_ROOT/w3cs-bundle-upload.XXXXXX")
trap 'rm -rf "$W3CS_UPLOAD_WORK"' EXIT
W3CS_UPLOAD_PENDING="$W3CS_UPLOAD_WORK/pending"
W3CS_UPLOAD_STALE="$W3CS_UPLOAD_WORK/stale"
W3CS_UPLOAD_CONFIRMED="$W3CS_UPLOAD_WORK/confirmed"
: > "$W3CS_UPLOAD_PENDING"
: > "$W3CS_UPLOAD_STALE"
: > "$W3CS_UPLOAD_CONFIRMED"

# A marker is an optimization, not permanent proof. Once a day, compare the
# complete marker set with R2's authoritative object listing. Missing objects
# lose their markers before upload, so a failed repair remains retryable and
# relays stop treating the affected bundle as CDN-backed after a restart.
W3CS_UPLOAD_NOW=${W3CS_NOW_EPOCH:-$(date +%s)}
W3CS_UPLOAD_LAST_RECONCILE=0
if [ -e "$W3CS_UPLOAD_RECONCILE_STAMP" ]; then
  W3CS_UPLOAD_LAST_RECONCILE=$(python3 - "$W3CS_UPLOAD_RECONCILE_STAMP" <<'PY'
import os
import sys

try:
    print(int(os.path.getmtime(sys.argv[1])))
except OSError:
    print(0)
PY
)
  W3CS_UPLOAD_LAST_RECONCILE=${W3CS_UPLOAD_LAST_RECONCILE:-0}
fi
W3CS_UPLOAD_DO_RECONCILE=0
if (( W3CS_UPLOAD_RECONCILE_SECONDS <= 0 ||
      W3CS_UPLOAD_NOW - W3CS_UPLOAD_LAST_RECONCILE >= W3CS_UPLOAD_RECONCILE_SECONDS )); then
  W3CS_UPLOAD_DO_RECONCILE=1
fi

if (( W3CS_UPLOAD_DO_RECONCILE )); then
  W3CS_UPLOAD_REMOTE_LIST="$W3CS_UPLOAD_WORK/remote"
  "$W3CS_UPLOAD_RCLONE" --config "$W3CS_UPLOAD_RCLONE_CONFIG" \
    lsf "$W3CS_UPLOAD_REMOTE" --files-only --include "*.bin" \
    --format ps --separator $'\t' \
    > "$W3CS_UPLOAD_REMOTE_LIST"
  python3 - "$W3CS_UPLOAD_SPOOL" "$W3CS_UPLOAD_MARK" \
      "$W3CS_UPLOAD_REMOTE_LIST" "$W3CS_UPLOAD_PENDING" \
      "$W3CS_UPLOAD_STALE" "$W3CS_UPLOAD_CONFIRMED" <<'PY'
import hashlib
import os
import re
import sys

spool, marker_dir, remote_path, pending_path, stale_path, confirmed_path = sys.argv[1:]
digest_pattern = re.compile(r"[0-9a-f]{64}\Z")

remote = {}
with open(remote_path, encoding="utf-8") as source:
    for raw in source:
        try:
            name, size = raw.rstrip("\n").rsplit("\t", 1)
            digest = name[:-4]
            if name.endswith(".bin") and digest_pattern.fullmatch(digest):
                remote[digest] = int(size)
        except ValueError:
            continue

pending = []
stale = set()
for entry in os.scandir(marker_dir):
    if not (entry.is_file(follow_symlinks=False)
            and digest_pattern.fullmatch(entry.name)):
        continue
    bundle_path = os.path.join(spool, entry.name + ".bundle")
    try:
        local_size = os.path.getsize(bundle_path)
    except OSError:
        local_size = None
    if entry.name not in remote or (local_size is not None
                                    and remote[entry.name] != local_size):
        stale.add(entry.name)
confirmed = []
for entry in os.scandir(spool):
    if not entry.is_file(follow_symlinks=False) or not entry.name.endswith(".bundle"):
        continue
    digest = entry.name[:-7]
    if not digest_pattern.fullmatch(digest):
        continue
    marked = os.path.exists(os.path.join(marker_dir, digest))
    local_size = entry.stat(follow_symlinks=False).st_size
    present = digest in remote and remote[digest] == local_size
    if not present:
        hasher = hashlib.sha256()
        with open(entry.path, "rb") as source:
            for block in iter(lambda: source.read(1024 * 1024), b""):
                hasher.update(block)
        if hasher.hexdigest() != digest:
            raise SystemExit(
                f"local bundle digest mismatch: {entry.path}")
        pending.append(entry.path)
    elif not marked:
        confirmed.append(os.path.join(marker_dir, digest))

stale_paths = [os.path.join(marker_dir, digest) for digest in sorted(stale)]
for path, rows in ((pending_path, pending), (stale_path, stale_paths),
                   (confirmed_path, confirmed)):
    with open(path, "w", encoding="utf-8") as target:
        target.writelines(row + "\n" for row in rows)

print(f"reconciled bundles: remote={len(remote)} pending={len(pending)} "
      f"stale_markers={len(stale_paths)} confirmed_unmarked={len(confirmed)}")
PY
  while IFS= read -r W3CS_UPLOAD_PATH; do
    [ -n "$W3CS_UPLOAD_PATH" ] && rm -f -- "$W3CS_UPLOAD_PATH"
  done < "$W3CS_UPLOAD_STALE"
  while IFS= read -r W3CS_UPLOAD_PATH; do
    [ -n "$W3CS_UPLOAD_PATH" ] && touch -- "$W3CS_UPLOAD_PATH"
  done < "$W3CS_UPLOAD_CONFIRMED"
else
  # The frequent path is local-only. os.scandir avoids one basename process
  # per bundle; that old loop consumed minutes of CPU at 200k+ files.
  python3 - "$W3CS_UPLOAD_SPOOL" "$W3CS_UPLOAD_MARK" \
      "$W3CS_UPLOAD_PENDING" <<'PY'
import os
import re
import sys

spool, marker_dir, pending_path = sys.argv[1:]
digest_pattern = re.compile(r"[0-9a-f]{64}\Z")
with open(pending_path, "w", encoding="utf-8") as target:
    for entry in os.scandir(spool):
        if not entry.is_file(follow_symlinks=False) or not entry.name.endswith(".bundle"):
            continue
        digest = entry.name[:-7]
        if (digest_pattern.fullmatch(digest)
                and not os.path.exists(os.path.join(marker_dir, digest))):
            target.write(entry.path + "\n")
PY
fi

# Hard-link the immutable snapshot into a staging directory with the served
# object names (<digest>.bin) so one rclone invocation moves the whole batch.
W3CS_UPLOAD_STAGING="$W3CS_UPLOAD_WORK/staging"
mkdir "$W3CS_UPLOAD_STAGING"
W3CS_UPLOAD_COUNT=0
while IFS= read -r W3CS_UPLOAD_FILE; do
  [ -n "$W3CS_UPLOAD_FILE" ] || continue
  W3CS_UPLOAD_NAME=${W3CS_UPLOAD_FILE##*/}
  W3CS_UPLOAD_DIGEST=${W3CS_UPLOAD_NAME%.bundle}
  ln "$W3CS_UPLOAD_FILE" "$W3CS_UPLOAD_STAGING/$W3CS_UPLOAD_DIGEST.bin"
  W3CS_UPLOAD_COUNT=$((W3CS_UPLOAD_COUNT + 1))
done < "$W3CS_UPLOAD_PENDING"

# --s3-no-head: rclone 1.60 (the Ubuntu 24.04 package) verifies an upload
# with HEAD <key>?versionId=<id from the PUT>, and R2 answers that with
# 501 Not Implemented. Content-MD5 on the PUT still guards the bytes.
W3CS_UPLOAD_R2_OPTS=(--s3-no-head)

if (( W3CS_UPLOAD_COUNT > 0 )); then
  "$W3CS_UPLOAD_RCLONE" --config "$W3CS_UPLOAD_RCLONE_CONFIG" \
    copy "$W3CS_UPLOAD_STAGING" "$W3CS_UPLOAD_REMOTE" \
    "${W3CS_UPLOAD_R2_OPTS[@]}" \
    --no-traverse --size-only --transfers 8 --retries 2 \
    --header-upload "Cache-Control: public, max-age=31536000, immutable"

  # rclone exited 0: every staged object either existed or was uploaded.
  # Re-create repaired stale markers only after that successful result.
  while IFS= read -r W3CS_UPLOAD_FILE; do
    [ -n "$W3CS_UPLOAD_FILE" ] || continue
    W3CS_UPLOAD_NAME=${W3CS_UPLOAD_FILE##*/}
    touch -- "$W3CS_UPLOAD_MARK/${W3CS_UPLOAD_NAME%.bundle}"
  done < "$W3CS_UPLOAD_PENDING"
  echo "uploaded $W3CS_UPLOAD_COUNT bundle(s)"
fi

if (( W3CS_UPLOAD_DO_RECONCILE )); then
  python3 - "$W3CS_UPLOAD_RECONCILE_STAMP" "$W3CS_UPLOAD_NOW" <<'PY'
import os
import sys

path, timestamp = sys.argv[1], int(sys.argv[2])
with open(path, "a", encoding="utf-8"):
    pass
os.utime(path, (timestamp, timestamp))
PY
fi

# Pack work must run even when every loose bundle already has a marker. A new
# replay can add a digest manifest without creating any new bundle bodies.
"$W3CS_UPLOAD_PACK_BUILDER"
W3CS_UPLOAD_PACK_OUT="$W3CS_UPLOAD_SPOOL/pack-out"
if [ -d "$W3CS_UPLOAD_PACK_OUT" ]; then
  # Multiple workers publish into this shared prefix. Never sync a worker's
  # partial local view: sync lets one worker delete another worker's packs.
  "$W3CS_UPLOAD_RCLONE" --config "$W3CS_UPLOAD_RCLONE_CONFIG" \
    copy "$W3CS_UPLOAD_PACK_OUT" "$W3CS_UPLOAD_PACK_REMOTE" \
    "${W3CS_UPLOAD_R2_OPTS[@]}" \
    --include "*.zpack" --size-only --transfers 4 --retries 2 \
    --header-upload "Cache-Control: public, max-age=31536000, immutable"

  # Replay manifests are mutable pointers. Each worker owns a separate R2
  # namespace, so two boxes can never overwrite or downgrade one another.
  # The browser reads every worker shard and selects the richest valid one.
  if [[ ! "$W3CS_UPLOAD_PACK_PUBLISHER" =~ ^[A-Za-z0-9.-]+$ ]]; then
    echo "invalid or missing W3CS_PACK_PUBLISHER_ID" >&2
    exit 2
  fi
  W3CS_UPLOAD_PACK_SHARD_REMOTE="${W3CS_UPLOAD_PACK_REMOTE%/}/by-worker/$W3CS_UPLOAD_PACK_PUBLISHER/"
  W3CS_UPLOAD_PUBLISH_MANIFESTS="$W3CS_UPLOAD_WORK/publish-pack-manifests"
  mkdir "$W3CS_UPLOAD_PUBLISH_MANIFESTS"
  python3 - "$W3CS_UPLOAD_PACK_OUT" \
      "$W3CS_UPLOAD_PUBLISH_MANIFESTS" <<'PY'
import json
import os
import re
import shutil
import sys

local_dir, publish_dir = sys.argv[1:]
manifest_pattern = re.compile(r"replay-[A-Za-z0-9._-]+\.json\Z")
pack_pattern = re.compile(r"[0-9a-f]{64}\.zpack\Z")


def load_manifest(path):
    try:
        with open(path, encoding="utf-8") as source:
            manifest = json.load(source)
        if manifest.get("version") != 1 or not isinstance(manifest.get("key"), str):
            return None
        packs = manifest.get("packs")
        replay_entries = manifest.get("replayEntries")
        if (not isinstance(packs, list) or not packs
                or isinstance(replay_entries, bool)
                or not isinstance(replay_entries, int)
                or replay_entries <= 0):
            return None
        pack_entries = 0
        for pack in packs:
            if not isinstance(pack, dict):
                return None
            name = pack.get("object")
            packed_bytes = pack.get("packedBytes")
            plain_bytes = pack.get("plainBytes")
            entries = pack.get("entries")
            if (not isinstance(name, str) or not pack_pattern.fullmatch(name)
                    or not isinstance(packed_bytes, int) or packed_bytes <= 0
                    or not isinstance(plain_bytes, int) or plain_bytes <= 0
                    or not isinstance(entries, int) or entries <= 0):
                return None
            pack_entries += entries
            pack_path = os.path.join(local_dir, name)
            if (not os.path.isfile(pack_path)
                    or os.path.getsize(pack_path) != packed_bytes):
                return None
        if replay_entries > pack_entries:
            return None
        return manifest
    except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError):
        return None


published = 0
rejected = 0
for entry in os.scandir(local_dir):
    if not entry.is_file(follow_symlinks=False) or not manifest_pattern.fullmatch(entry.name):
        continue
    if load_manifest(entry.path) is None:
        rejected += 1
        print(f"rejected invalid pack manifest: {entry.name}", file=sys.stderr)
        continue
    shutil.copy2(entry.path, os.path.join(publish_dir, entry.name))
    published += 1

print(f"pack shard manifests: publish={published} rejected={rejected}")
PY
  # This namespace has exactly one writer, so syncing derived pointers is
  # safe. It removes retired replay manifests without touching the shared
  # immutable zpack root or another worker's shard.
  "$W3CS_UPLOAD_RCLONE" --config "$W3CS_UPLOAD_RCLONE_CONFIG" \
    sync "$W3CS_UPLOAD_PUBLISH_MANIFESTS" "$W3CS_UPLOAD_PACK_SHARD_REMOTE" \
    "${W3CS_UPLOAD_R2_OPTS[@]}" --delete-after \
    --include "replay-*.json" --transfers 4 --retries 2 \
    --header-upload "Cache-Control: public, max-age=300"
fi