#!/usr/bin/env bash
set -euo pipefail

CLASSIC_ROOT=${W3_CLASSIC_ROOT:-/opt/war3-runtime/classic}
CLASSIC_131_ROOT=${W3_CLASSIC_131_ROOT:-/opt/war3-runtime/classic-131}
PREFIX=${WINEPREFIX:-"$CLASSIC_ROOT/prefix"}
DISPLAY_NUMBER=${W3_D3D9_DISPLAY:-:11}
REPLAY_ROOT=${W3CS_REPLAY_ROOT:-/home/ubuntu/w3cs-lab/replays}
REPLAY_ID=${W3CS_REPLAY_ID:-default}
GAME_PROFILE=${W3CS_GAME_PROFILE:-auto}
PROFILE_ROOT=${W3CS_ENGINE_PROFILE_ROOT:-/home/ubuntu/w3cs-lab/engine-profiles}
REPLAY_DIRECTORY=Replays
REPLAY_WINDOWS='C:\users\ubuntu\Documents\Warcraft III\Replays\trace.w3g'
SEAT_STATE=${W3CS_SEAT_STATE:-/tmp/w3cs-seat-state.json}
WINE_BIN=/opt/wine-staging/bin/wine
WINESERVER_BIN=/opt/wine-staging/bin/wineserver
# The viewer selects the session's authoritative resolution. NOTE: only the
# 1.28.5 engine honors it (registry stamp); command-line -width/-height are
# NOT understood by the older engines - unknown arguments break their
# -loadfile parsing and strand the game at the main menu (verified live).
# Older lanes render at their desktop-derived size until a safe per-lane
# mechanism exists (candidate: resize the X window after launch).
SOURCE_WIDTH=${W3CS_SOURCE_WIDTH:-${W3_D3D9_SOURCE_WIDTH:-1024}}
SOURCE_HEIGHT=${W3CS_SOURCE_HEIGHT:-${W3_D3D9_SOURCE_HEIGHT:-768}}
if [[ ! "$SOURCE_WIDTH" =~ ^[0-9]{3,4}$ || ! "$SOURCE_HEIGHT" =~ ^[0-9]{3,4}$ ]]; then
  SOURCE_WIDTH=1024
  SOURCE_HEIGHT=768
fi
GAME_ARGS=(-window)

if [[ ! "$REPLAY_ID" =~ ^[A-Za-z0-9._-]{1,96}$ ]]; then
  echo "invalid replay id" >&2
  exit 2
fi
REPLAY_SOURCE="$REPLAY_ROOT/$REPLAY_ID.w3g"
REPLAY_METADATA="$REPLAY_ROOT/$REPLAY_ID.json"
if [[ ! -f "$REPLAY_SOURCE" ]]; then
  if [[ "$REPLAY_ID" =~ ^[0-9a-f]{40}$ &&
        -f "${W3CS_CATALOG_INDEX:-}" &&
        ( -n "${W3CS_CATALOG_BUCKET:-}" || -n "${W3CS_CATALOG_BASE_URL:-}" ) ]]; then
    # R2 mirror through the public CDN (no AWS credentials on CPU
    # workers) or the AWS bucket (GPU workers); the index schema matches.
    if [[ -n "${W3CS_CATALOG_BASE_URL:-}" ]]; then
      python3 /home/ubuntu/w3cs-lab/materialize-catalog-replay.py \
        --index "$W3CS_CATALOG_INDEX" --source-sha1 "$REPLAY_ID" \
        --output-root "$REPLAY_ROOT" --base-url "$W3CS_CATALOG_BASE_URL"
    else
      python3 /home/ubuntu/w3cs-lab/materialize-catalog-replay.py \
        --index "$W3CS_CATALOG_INDEX" --source-sha1 "$REPLAY_ID" \
        --output-root "$REPLAY_ROOT" --bucket "$W3CS_CATALOG_BUCKET" \
        --region "${W3CS_CATALOG_REGION:-us-east-1}"
    fi
  else
    echo "replay not found: $REPLAY_ID" >&2
    exit 3
  fi
fi

if [[ "$GAME_PROFILE" == auto ]]; then
  [[ -f "$REPLAY_METADATA" ]] || {
    echo "automatic engine selection needs replay metadata: $REPLAY_ID" >&2
    exit 4
  }
  GAME_PROFILE=$(python3 - "$REPLAY_METADATA" <<'PY'
import json
import sys
with open(sys.argv[1], encoding="utf-8") as source:
    value = json.load(source).get("engineProfile")
if not isinstance(value, str):
    raise SystemExit("replay metadata has no engineProfile")
print(value)
PY
  )
fi

case "$GAME_PROFILE" in
  native-1285)
    BASE_GAME_DIR="$CLASSIC_ROOT/war3-1285/Warcraft III"
    GAME_EXE='Warcraft III.exe'
    ;;
  native-124ab|native-124cde|native-1220|native-1120|native-1110|native-1100)
    BASE_GAME_DIR="$CLASSIC_ROOT/war3rk/Warcraft III"
    GAME_EXE='war3.exe'
    ENGINE_OVERLAY="$PROFILE_ROOT/${GAME_PROFILE/native-/classic-}/files"
    ;;
  replaykit|replaykit-124ab)
    BASE_GAME_DIR="$CLASSIC_ROOT/war3rk/Warcraft III"
    GAME_EXE='ReplayKit.exe'
    ;;
  native-1140)
    BASE_GAME_DIR="$CLASSIC_ROOT/war3-1140/Warcraft III"
    GAME_EXE='war3.exe'
    ;;
  native-1311-d3d9)
    BASE_GAME_DIR="$CLASSIC_131_ROOT/source/Warcraft III"
    GAME_EXE='x86/Warcraft III.exe'
    PREFIX="${W3CS_PREFIX64:-$CLASSIC_131_ROOT/prefix64}"
    REPLAY_DIRECTORY=Replay
    REPLAY_WINDOWS='C:\users\ubuntu\Documents\Warcraft III\Replay\trace.w3g'
    WINE_BIN=/usr/bin/wine
    WINESERVER_BIN=/usr/bin/wineserver
    GAME_ARGS=(-graphicsapi Direct3D9 -windowmode windowed \
      -width "$SOURCE_WIDTH" -height "$SOURCE_HEIGHT")
    ;;
  *)
    echo "unsupported game profile: $GAME_PROFILE" >&2
    exit 4
    ;;
esac
if [[ "$GAME_PROFILE" != native-124ab && \
      "$GAME_PROFILE" != native-124cde && \
      ! -f "$BASE_GAME_DIR/$GAME_EXE" ]]; then
  echo "game profile is not installed: $GAME_PROFILE" >&2
  exit 5
fi

# ---- CRIU claim fast path ---------------------------------------------------
# A generic per-(profile, family) checkpoint holds a fully booted engine —
# with the d3d9 command-stream proxy loaded — blocked at the replay-open
# gate. Restoring it takes ~150 ms; the requested replay is installed at the
# gate's fixed path before the release marker, so any same-family replay
# plays after only its own map load. The claim runs as one process tree
# (own Xvfb, embedded pulse, wineserver) in a PID namespace; its
# environment is published to the relay through the per-seat claim file.
# Falls through to the classic cold launch when no image exists.
SEAT=$((${W3CS_RELAY_PORT:-8145} - 8144))
CRIU_IMAGES=${W3CS_CRIU_IMAGES:-/opt/war3-runtime/classic/w3cs-criu/images}
CRIU_WORKER=${W3CS_CRIU_WORKER:-/usr/local/sbin/war3-classic-criu}
CRIU_STATE_ROOT=${W3CS_CRIU_STATE_ROOT:-/run/w3cs-criu}
REPLAY_FAMILY=$(head -c 52 "$REPLAY_SOURCE" | tail -c 4 | rev)
case "$REPLAY_FAMILY" in
  W3XP) REPLAY_FAMILY=tft ;;
  WAR3) REPLAY_FAMILY=roc ;;
  *) REPLAY_FAMILY= ;;
esac
# A Reign of Chaos replay needs the engine in Reign of Chaos mode. The CRIU
# claim path has always passed -classic for the roc family; the cold launch
# did not, so a roc replay with no claim image opened a Frozen Throne engine
# and died in game.dll on a null dereference (+0xb0) before writing a line of
# War3Log. Every engine these lanes launch shares one war3.exe, and the roc
# claims prove it takes the switch. 1.31 keeps its own argument list.
if [[ "$REPLAY_FAMILY" == roc && "$GAME_PROFILE" != native-1311-d3d9 ]]; then
  GAME_ARGS=(-classic "${GAME_ARGS[@]}")
fi
CRIU_BASE="$CRIU_IMAGES/${GAME_PROFILE}__${REPLAY_FAMILY}/seat$SEAT"
if [[ -n "$REPLAY_FAMILY" && -f "$CRIU_BASE/war3.json" && \
      -f "$CRIU_BASE/generic/inventory.img" && -x "$CRIU_WORKER" && \
      "${W3CS_CRIU_DISABLE:-0}" != 1 ]] && \
    sudo -n W3_CRIU_STATE_ROOT="$CRIU_STATE_ROOT" \
      "$CRIU_WORKER" status "$SEAT" >/dev/null 2>&1; then
  CLAIM_FILE="/tmp/w3cs-claim${W3CS_LOG_SUFFIX:-}.json"
  CLAIM_DISPLAY=":$((19 + SEAT))"
  CLAIM_RUNTIME="/run/war3-audio/slot$SEAT"
  CRIU_VIEW="$(dirname "$CRIU_IMAGES")/$GAME_PROFILE/game"
  GATE_DIR="/run/war3-replay-gate/slot$SEAT"
  write_claim() {
    printf '{"ready":%s,"display":"%s","pulseServer":"unix:%s/pulse/native","pulseDevice":"game%s.monitor","rootPid":%s}\n' \
      "$1" "$CLAIM_DISPLAY" "$CLAIM_RUNTIME" "$SEAT" "${2:-0}" \
      > "$CLAIM_FILE.tmp" && mv -f "$CLAIM_FILE.tmp" "$CLAIM_FILE"
  }
  criu_cleanup() {
    sudo -n W3_CRIU_STATE_ROOT="$CRIU_STATE_ROOT" \
      "$CRIU_WORKER" stop "$SEAT" >/dev/null 2>&1 || true
    rm -f "$CLAIM_FILE" "$CLAIM_FILE.tmp"
    rm -f -- "$SEAT_STATE"
  }
  trap criu_cleanup EXIT INT TERM
  write_claim false
  # The replay's exact map (identity-verified) goes into this seat's prefix
  # and the shared profile view; the worker restore installs the replay
  # itself at the gate path.
  if [[ -f "$REPLAY_METADATA" ]]; then
    python3 /home/ubuntu/w3cs-lab/stage-replay-session.py \
      --metadata "$REPLAY_METADATA" --replay "$REPLAY_SOURCE" \
      --prefix "$PREFIX" --game-directory "$CRIU_VIEW" \
      --replay-directory "$REPLAY_DIRECTORY"
  fi
  rm -f "$GATE_DIR/ready" "$GATE_DIR/release"
  sudo -n \
    W3_CRIU_STATE_ROOT="$CRIU_STATE_ROOT" \
    W3_CRIU_IMAGE_DIR="$CRIU_BASE/generic" \
    W3_CLASSIC_CRIU_MANIFEST="$CRIU_BASE/war3.json" \
    W3_CLASSIC_SLOT_PREFIX="$PREFIX" \
    W3_CLASSIC_MASTER_REPLAY="$REPLAY_SOURCE" \
    W3_CLASSIC_AUDIO_MODE=embedded \
    W3_CLASSIC_GAME_EXE="${GAME_EXE##*/}" \
    "$CRIU_WORKER" restore "$SEAT"
  CLAIM_ROOT=$(cat "$CRIU_STATE_ROOT/slot$SEAT/root.pid")
  CLAIM_GAME=$(cat "$CRIU_STATE_ROOT/slot$SEAT/game.pid")
  write_claim true "$CLAIM_ROOT"
  python3 - "$SEAT_STATE" "$GAME_PROFILE" "$PREFIX" "$CRIU_VIEW" \
    "$REPLAY_DIRECTORY" "$CLAIM_DISPLAY" "$REPLAY_ID" <<'PY'
import json
import os
from pathlib import Path
import sys

target = Path(sys.argv[1])
temporary = target.with_suffix(target.suffix + ".tmp")
temporary.write_text(json.dumps({
    "engineProfile": sys.argv[2],
    "winePrefix": sys.argv[3],
    "gameDirectory": sys.argv[4],
    "replayDirectory": sys.argv[5],
    "display": sys.argv[6],
    "replayId": sys.argv[7],
    "criuClaim": True,
}, sort_keys=True), encoding="utf-8")
os.replace(temporary, target)
PY
  touch "$GATE_DIR/release"
  echo "criu claim restored profile=$GAME_PROFILE family=$REPLAY_FAMILY" \
    "seat=$SEAT root=$CLAIM_ROOT game=$CLAIM_GAME"
  # The namespace tree is reparented to host init, so poll instead of wait.
  while kill -0 "$CLAIM_GAME" 2>/dev/null; do
    sleep .5
  done
  exit 0
fi

export WINEPREFIX="$PREFIX"
if [[ "$GAME_PROFILE" != native-1311-d3d9 ]]; then
  export WINEARCH=win32
else
  unset WINEARCH
fi
export WINEDEBUG=${WINEDEBUG:--all}
export PATH="$(dirname "$WINE_BIN"):$PATH"
export DISPLAY="$DISPLAY_NUMBER"
mkdir -p "$PREFIX/drive_c/users/ubuntu/Temp"

# Give each viewer an isolated game tree. Symlinks share the immutable MPQs,
# maps, and page cache, while patch DLLs and session maps stay private. This is
# fast enough for watch startup and safe when seats use different patches.
SESSION_ROOT=$(mktemp -d /tmp/w3cs-game-session.XXXXXX)
GAME_DIR="$SESSION_ROOT/game"

cleanup() {
  "$WINESERVER_BIN" -k 2>/dev/null || true
  rm -f -- "$SEAT_STATE"
  if [[ "${SESSION_ROOT:-}" == /tmp/w3cs-game-session.* ]]; then
    rm -rf -- "$SESSION_ROOT"
  fi
}
trap cleanup EXIT INT TERM

mkdir -p "$GAME_DIR"
cp -as "$BASE_GAME_DIR/." "$GAME_DIR/"
if [[ "$GAME_PROFILE" == native-1311-d3d9 ]]; then
  # CASC requires a writable per-process shared-memory file beside its
  # immutable data archives. The archives remain symlinked and page-shared.
  touch "$GAME_DIR/Data/data/shmem"
fi

"$WINESERVER_BIN" -k 2>/dev/null || true
if [[ -f "$REPLAY_METADATA" ]]; then
  python3 /home/ubuntu/w3cs-lab/stage-replay-session.py \
    --metadata "$REPLAY_METADATA" --replay "$REPLAY_SOURCE" \
    --prefix "$PREFIX" --game-directory "$GAME_DIR" \
    --replay-directory "$REPLAY_DIRECTORY"
else
  mkdir -p "$PREFIX/drive_c/users/ubuntu/Documents/Warcraft III/$REPLAY_DIRECTORY"
  cp "$REPLAY_SOURCE" \
    "$PREFIX/drive_c/users/ubuntu/Documents/Warcraft III/$REPLAY_DIRECTORY/trace.w3g"
fi
cd "$GAME_DIR"
# Drive this off the overlay path the profile case set, not off a second copy
# of the profile list: the two lists silently disagreed when native-1120 was
# added, so that lane ran the base engine and reported a missing CD key.
if [[ -n "${ENGINE_OVERLAY:-}" ]]; then
  [[ -f "$ENGINE_OVERLAY/war3.exe" ]] || {
    echo "engine overlay is not installed: $GAME_PROFILE" >&2
    exit 6
  }
  # ReplayKit is a build-time materializer only. The watch path installs one
  # immutable engine overlay, then launches the selected replay directly.
  # Large MPQs and maps stay shared. No replay-specific CRIU image is needed.
  # Remove EVERY target before copying: the view entries are cp -as
  # symlinks into the shared install, and cp onto a symlink follows it —
  # each session was silently rewriting the shared war3rk Storm.dll and
  # War3Patch.mpq with this profile's versions (found 2026-08-21 after it
  # broke the ReplayKit lane and a CRIU image).
  rm -f "$GAME_DIR/Game.dll" "$GAME_DIR/war3.exe" "$GAME_DIR/game.dll" \
    "$GAME_DIR/Storm.dll" "$GAME_DIR/War3Patch.mpq"
  for file in Storm.dll War3Patch.mpq war3.exe game.dll; do
    cp "$ENGINE_OVERLAY/$file" "$GAME_DIR/$file"
  done
fi

PROXY_DIRECTORY="$GAME_DIR/$(dirname "$GAME_EXE")"
rm -f "$PROXY_DIRECTORY/d3d9.dll"
cp /home/ubuntu/w3cs-lab/d3d9.dll "$PROXY_DIRECTORY/d3d9.dll"

if [[ "$GAME_PROFILE" == replaykit || \
      "$GAME_PROFILE" == replaykit-124ab ]]; then
  # This is the cold engine-profile builder path. ReplayKit reads lowercase
  # game.dll. Linux can retain an old lowercase file beside uppercase
  # Game.dll, so remove both names before restoring the 1.28f base.
  BASE_DIR="$CLASSIC_ROOT/war3-1285/Warcraft III"
  # rm every target first: the view files are symlinks into the shared
  # war3rk install, and cp onto a symlink writes THROUGH it (this branch
  # was overwriting the shared Warcraft III.exe, Storm.dll, and
  # War3Patch.mpq on every ReplayKit session).
  rm -f "$GAME_DIR/Game.dll" "$GAME_DIR/game.dll" "$GAME_DIR/war3.exe" \
    "$GAME_DIR/Warcraft III.exe" "$GAME_DIR/Storm.dll" \
    "$GAME_DIR/War3Patch.mpq"
  rm -rf -- "$GAME_DIR/ReplayKitBackup"
  cp "$BASE_DIR/Warcraft III.exe" "$GAME_DIR/Warcraft III.exe"
  cp "$BASE_DIR/Game.dll" "$GAME_DIR/game.dll"
  cp "$BASE_DIR/Storm.dll" "$GAME_DIR/Storm.dll"
  cp "$BASE_DIR/War3Patch.mpq" "$GAME_DIR/War3Patch.mpq"
  # ReplayKit resolves Wine's stable C: replay path reliably. A Z: path below
  # the private symlink view can be rejected by its legacy path parser.
  REPLAYKIT_WINDOWS="$REPLAY_WINDOWS"
  replaykit_args=("$GAME_EXE" "$REPLAYKIT_WINDOWS" -names -colors -window)
  if [[ "$GAME_PROFILE" == replaykit ]]; then
    replaykit_args+=(-totalsilence)
  fi
  wine "${replaykit_args[@]}" &
  replaykit_pid=$!
  game_pid=
  for _ in $(seq 1 240); do
    if [[ "$GAME_PROFILE" == replaykit-124ab ]]; then
      while read -r window; do
        [[ -n "$window" ]] || continue
        xdotool key --window "$window" Return >/dev/null 2>&1 || true
      done < <(xdotool search --onlyvisible --name '^ReplayKit v1.28b$' \
        2>/dev/null || true)
    fi
    game_pid=$(pgrep -n -u "$(id -u)" -x war3.exe 2>/dev/null || true)
    [[ -n "$game_pid" ]] && break
    if ! kill -0 "$replaykit_pid" 2>/dev/null; then
      wait "$replaykit_pid" || true
      echo "ReplayKit exited before starting the game" >&2
      exit 7
    fi
    sleep .25
  done
  [[ -n "$game_pid" ]] || {
    echo "ReplayKit timed out before starting the game" >&2
    exit 8
  }
  while kill -0 "$game_pid" 2>/dev/null; do
    sleep .5
  done
  exit 0
fi
if [[ "$GAME_PROFILE" == native-1311-d3d9 ]]; then
  export WINEDLLOVERRIDES='d3d9=n,b;dxgi,d3d11=b'
fi
python3 - "$SEAT_STATE" "$GAME_PROFILE" "$PREFIX" "$GAME_DIR" \
  "$REPLAY_DIRECTORY" "$DISPLAY_NUMBER" "$REPLAY_ID" <<'PY'
import json
import os
from pathlib import Path
import sys

target = Path(sys.argv[1])
temporary = target.with_suffix(target.suffix + ".tmp")
temporary.write_text(json.dumps({
    "engineProfile": sys.argv[2],
    "winePrefix": sys.argv[3],
    "gameDirectory": sys.argv[4],
    "replayDirectory": sys.argv[5],
    "display": sys.argv[6],
    "replayId": sys.argv[7],
}, sort_keys=True), encoding="utf-8")
os.replace(temporary, target)
PY
# Also stamp the resolution into this profile's prefix immediately before
# launch: 1.28.5 reads it from the registry, and every lane then keeps the
# game, command stream, input mapping, and browser on one value end to end.
if [[ "$SOURCE_WIDTH" =~ ^[0-9]{3,4}$ && "$SOURCE_HEIGHT" =~ ^[0-9]{3,4}$ ]]; then
  "$WINE_BIN" reg add \
    'HKCU\Software\Blizzard Entertainment\Warcraft III\Video' \
    /v reswidth /t REG_DWORD /d "$SOURCE_WIDTH" /f >/dev/null 2>&1 || true
  "$WINE_BIN" reg add \
    'HKCU\Software\Blizzard Entertainment\Warcraft III\Video' \
    /v resheight /t REG_DWORD /d "$SOURCE_HEIGHT" /f >/dev/null 2>&1 || true
fi

# Old-engine resolution status (verified live): pre-1.28 engines size their
# window AND backbuffer from the desktop at device creation (1432x1046 on
# the 1920x1440 root), ignore reswidth/resheight in windowed mode, break on
# unknown -width/-height arguments, and keep their backbuffer through a
# window-manager resize. Forcing 1024x768 for them therefore needs a proxy-
# level override of the CreateDevice presentation parameters (or fullscreen
# plumbing); until then these lanes stream at their native windowed size.
"$WINE_BIN" "$GAME_EXE" "${GAME_ARGS[@]}" -loadfile "$REPLAY_WINDOWS" &
launcher_pid=$!
set +e
wait "$launcher_pid"
launcher_status=$?
set -e

# Some Classic binaries replace their initial Wine process during startup.
# Waiting only for the launcher makes the EXIT trap kill the relaunched game
# in the loading screen. Keep the seat wrapper alive while any WC3 process in
# this isolated process group remains. Allow a short gap between the launcher
# exit and the replacement process becoming visible.
session_pgid=$(ps -o pgid= -p $$ | tr -d ' ')
game_processes() {
  pgrep -g "$session_pgid" -f \
    'Warcraft III\.exe|(^|[/ ])war3\.exe([ $]|$)' 2>/dev/null || true
}
replacement=
for _ in $(seq 1 20); do
  replacement=$(game_processes)
  [[ -n "$replacement" ]] && break
  sleep .25
done
while [[ -n "$replacement" ]]; do
  sleep .5
  replacement=$(game_processes)
done
exit "$launcher_status"