#!/usr/bin/env bash
# Exact-artifact admission for Stream.
#
# The previous whole-network experiment deployed image e1926f3, which was not
# the artifact its receipt covered. Nothing detected that. This binds one
# receipt to one image digest and gives deployment a way to refuse anything
# else.
#
#   scripts/admit build            build the exact linux/amd64 image, no publish
#   scripts/admit run              clean tree -> suites -> image -> receipt
#   scripts/admit publish [sha]    push the admitted image and pin the receipt
#                                  to the registry digest (default: HEAD)
#   scripts/admit verify <digest>  exit non-zero unless <digest> is admitted
#
# Design rules, learned from the failure this replaces:
#   - the receipt records the artifact's own digest, never a tag; tags move
#   - a skipped suite is written into the receipt as "skipped", never omitted,
#     so a partial run cannot read as a full one
#   - any suite failure aborts before a receipt exists
#   - the tree must be clean, so the receipt's revisions describe the source
#     that actually produced the binary
set -euo pipefail

cd "$(dirname "${BASH_SOURCE[0]}")/.."
ROOT="$PWD"
RECEIPTS="${STREAM_RECEIPTS_DIR:-$ROOT/receipts}"
IMAGE_REPO="${STREAM_IMAGE_REPO:-atcr.io/zat.dev/stream}"
# The default docker config routes atcr.io through a credential helper that
# errors; the isolated osxkeychain config is the one that authenticates.
# HOME may be unset in a minimal deployment environment, where `verify` still
# has to work -- so this must not be load-bearing under `set -u`.
if [ -z "${DOCKER_CONFIG:-}" ] && [ -n "${HOME:-}" ] && [ -d "$HOME/.config/docker-atcr" ]; then
  export DOCKER_CONFIG="$HOME/.config/docker-atcr"
fi

die() { printf 'admit: %s\n' "$*" >&2; exit 1; }
# stderr, not stdout: build_image returns the image id on stdout, and progress
# lines printed there end up captured inside the receipt's recorded identity.
note() { printf '\033[1m==> %s\033[0m\n' "$*" >&2; }

# Receipts are outputs, not source. A run writes its own receipt into the repo,
# so counting it as drift would make every completed run look dirty.
tree_drift() { git status --porcelain -- ':!receipts'; }

require_clean_tree() {
  [ -z "$(tree_drift)" ] ||
    die "worktree is dirty; the receipt must describe the source that built the artifact
$(tree_drift)"
}

# Every pinned input that can change behavior. A receipt that does not name
# these cannot be checked later.
revisions_json() {
  python3 - "$ROOT" <<'PY'
import json, re, subprocess, sys, pathlib
root = pathlib.Path(sys.argv[1])
zon = (root / "build.zig.zon").read_text()
deps = dict(re.findall(r'\.(\w[\w-]*) = \.\{\s*\.url = "([^"]+)"', zon))
def git(*a):
    return subprocess.check_output(["git", *a], cwd=root, text=True).strip()
upstream = pathlib.Path.home() / "github.com/bluesky-social/jetstream"
try:
    upstream_rev = subprocess.check_output(
        ["git", "rev-parse", "HEAD"], cwd=upstream, text=True).strip()
    upstream_dirty = bool(subprocess.check_output(
        ["git", "status", "--porcelain", "--untracked-files=no"], cwd=upstream, text=True).strip())
except Exception:
    upstream_rev, upstream_dirty = None, None
print(json.dumps({
    "stream": git("rev-parse", "HEAD"),
    "stream_describe": git("describe", "--always", "--dirty"),
    "zig": subprocess.check_output(["zig", "version"], text=True).strip(),
    "dependencies": deps,
    "upstream_jetstream": upstream_rev,
    "upstream_jetstream_dirty": upstream_dirty,
    "dashboard_sha256": __import__("hashlib").sha256(
        (root / "deploy/grafana/jetstream-upstream.json").read_bytes()).hexdigest(),
}, indent=2, sort_keys=True))
PY
}

# name|command. Negative suites come first: they are the ones that decide
# whether a green positive run means anything.
# The two simulator-backed oracles run FIRST. The pinned simulator accumulates
# world state for as long as it is up, and a full admission run is long enough
# to grow it by orders of magnitude; running them late meant they faced a world
# that made bootstrap outlast the harness's fixed serving timeout. Everything
# after them is simulator-independent.
SUITES=(
  "lifecycle-oracle|just oracle"
  "powerloss-oracle|just powerloss-oracle"
  "unit-debug|zig build test"
  "unit-releasesafe|zig build test -Doptimize=ReleaseSafe"
  "differential-oracle|just differential-oracle"
  "listener-contract|just listener-contract"
  "shutdown-contract|just shutdown-contract"
  "logging-contract|just logging-contract"
  "environment-contract|just environment-contract"
  "dashboard-test|just dashboard-test"
  "process-metrics-contract|just process-metrics-contract"
  "http-metrics-contract|just http-metrics-contract"
  "archive-contract|just archive-contract"
  "plan-config-contract|just plan-config-contract"
  "cursor-lookback-contract|just cursor-lookback-contract"
  "compaction-config-contract|just compaction-config-contract"
  "retry-config-contract|just retry-config-contract"
  "subscribe-config-contract|just subscribe-config-contract"
  "subscribe-read-batch-contract|just subscribe-read-batch-contract"
  "status-contract|just status-contract"
)

run_suites() {
  local out="$1" log_dir="$2" first=1
  mkdir -p "$log_dir"
  printf '{\n' >"$out"
  for entry in "${SUITES[@]}"; do
    local name="${entry%%|*}" cmd="${entry#*|}" status started elapsed
    # STREAM_ADMIT_SKIP is an escape hatch for an environment that genuinely
    # cannot run a suite. It is recorded, not hidden.
    if [[ ",${STREAM_ADMIT_SKIP:-}," == *",$name,"* ]]; then
      status="skipped"
      note "SKIP $name (STREAM_ADMIT_SKIP)"
    else
      note "$name"
      started=$SECONDS
      if (cd "$ROOT" && eval "$cmd") >"$log_dir/$name.log" 2>&1; then
        status="pass"
      else
        status="fail"
      fi
      elapsed=$((SECONDS - started))
      if [ "$status" = fail ]; then
        tail -30 "$log_dir/$name.log" >&2
        die "$name failed after ${elapsed}s; no receipt written"
      fi
      printf '  %s (%ss)\n' "$status" "$elapsed"
    fi
    [ $first -eq 1 ] || printf ',\n' >>"$out"
    first=0
    printf '  "%s": "%s"' "$name" "$status" >>"$out"
  done
  printf '\n}\n' >>"$out"
}

# Build the deploy artifact. STREAM_BUILD_HOST points at a native linux/amd64
# docker daemon (ssh://root@host); emulating amd64 on an arm64 workstation
# turns an 8-minute compile into 40-plus and produces a cross-emulated
# approximation of the thing that actually ships. When a remote is used the
# image is transferred back so the push uses local registry credentials and the
# build host never receives them.
build_image() {
  local tag="$1"
  if [ -n "${STREAM_BUILD_HOST:-}" ]; then
    # Deliberately not DOCKER_HOST=ssh://. That keeps one SSH-carried Docker API
    # stream open for the whole compile, and a multi-minute build drops it
    # ("broken pipe"), losing the run. Ship the context once, build in its own
    # session with keepalives, then stream the image back. The push still runs
    # locally, so the build host never receives registry credentials.
    local host key
    host="${STREAM_BUILD_HOST#ssh://}"
    key="${STREAM_SSH_KEY:-$HOME/.ssh/waow_ed25519}"
    local ssh_opts=(-i "$key" -o IdentitiesOnly=yes -o ServerAliveInterval=15
                    -o ServerAliveCountMax=10 -o ConnectTimeout=10)
    note "shipping build context to $host"
    # .git is part of the context: the Dockerfile copies it so the binary can
    # report the commit it was built from.
    tar -C "$ROOT" -czf - --exclude=.zig-cache --exclude=zig-out --exclude=receipts . |
      ssh "${ssh_opts[@]}" "$host" 'rm -rf /opt/streambuild && mkdir -p /opt/streambuild && tar -C /opt/streambuild -xzf -'
    note "building linux/amd64 image $tag natively on $host"
    # A long compile can drop its own output stream even with keepalives. That
    # says nothing about whether the build finished, so ask the remote daemon
    # instead of trusting the exit status of a broken pipe.
    ssh "${ssh_opts[@]}" "$host" "cd /opt/streambuild && docker build --platform linux/amd64 -t '$tag' ." >&2 || true
    ssh "${ssh_opts[@]}" "$host" "docker image inspect '$tag' >/dev/null" ||
      die "remote build did not produce $tag"
    note "transferring image to the local daemon"
    ssh "${ssh_opts[@]}" "$host" "docker save '$tag'" | docker load >&2
    docker image inspect --format '{{.Id}}' "$tag"
    return
  fi
  note "building linux/amd64 image $tag"
  docker build --platform linux/amd64 -t "$tag" "$ROOT" >&2
  # The local image ID is the artifact's identity before it is pushed. The
  # registry digest replaces it at publish time.
  docker image inspect --format '{{.Id}}' "$tag"
}

cmd_build() {
  require_clean_tree
  local sha; sha=$(git rev-parse --short HEAD)
  build_image "$IMAGE_REPO:$sha"
}

cmd_run() {
  require_clean_tree
  local sha stamp receipt log_dir suites_json revisions image_id tag
  sha=$(git rev-parse --short HEAD)
  stamp=$(date -u +%Y%m%dT%H%M%SZ)
  log_dir="$RECEIPTS/logs-$sha-$stamp"
  suites_json=$(mktemp)
  # The trap fires after this function returns, when the local is already out
  # of scope; without the default it aborts under `set -u`.
  trap 'rm -f "${suites_json:-}"' EXIT

  run_suites "$suites_json" "$log_dir"

  tag="$IMAGE_REPO:$sha"
  image_id=$(build_image "$tag")
  revisions=$(revisions_json)

  # Re-check: a suite run can take long enough for someone to touch the tree.
  require_clean_tree

  mkdir -p "$RECEIPTS"
  receipt="$RECEIPTS/$sha.json"
  # Revisions travel through a file, never interpolated into the program text:
  # JSON's `false`/`null` are not Python literals, which turned a completed run
  # into a NameError after every suite had already passed.
  local revisions_file
  revisions_file=$(mktemp)
  printf '%s\n' "$revisions" >"$revisions_file"
  python3 - "$receipt" "$image_id" "$tag" "$stamp" "$suites_json" "$revisions_file" <<'PY'
import json, sys
receipt, image_id, tag, stamp, suites_path, revisions_path = sys.argv[1:7]
json.dump({
    "schema": "stream-admission/1",
    "created_utc": stamp,
    "image": {"tag": tag, "local_id": image_id, "registry_digest": None},
    "revisions": json.load(open(revisions_path)),
    "suites": json.load(open(suites_path)),
}, open(receipt, "w"), indent=2, sort_keys=True)
open(receipt, "a").write("\n")
PY
  rm -f "$revisions_file"
  note "receipt $receipt"
  note "image not published; run 'scripts/admit publish' to pin its registry digest"
}

cmd_publish() {
  # Deliberately no clean-tree check: the receipt already records the revisions
  # its artifact was built from, so source state now is irrelevant. The real
  # guard is that the local image still matches the identity the receipt covers.
  local sha receipt tag digest
  sha="${1:-$(git rev-parse --short HEAD)}"
  receipt="$RECEIPTS/$sha.json"
  [ -f "$receipt" ] || die "no receipt for $sha; run 'scripts/admit run' first"
  tag=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["image"]["tag"])' "$receipt")

  # The image must still be the one the receipt covers.
  local recorded current
  recorded=$(python3 -c 'import json,sys;print(json.load(open(sys.argv[1]))["image"]["local_id"])' "$receipt")
  current=$(docker image inspect --format '{{.Id}}' "$tag")
  [ "$recorded" = "$current" ] ||
    die "local image $tag is $current but the receipt covers $recorded; rerun 'admit run'"

  note "pushing $tag"
  ATCR_AUTO_AUTH=1 docker push "$tag" >&2
  digest=$(docker image inspect --format '{{index .RepoDigests 0}}' "$tag" | cut -d@ -f2)
  [ -n "$digest" ] || die "could not read the pushed digest"

  python3 - "$receipt" "$digest" <<'PY'
import json, sys
receipt, digest = sys.argv[1], sys.argv[2]
data = json.load(open(receipt))
if data["image"]["registry_digest"] not in (None, digest):
    raise SystemExit(f"admit: receipt already pinned to {data['image']['registry_digest']}")
data["image"]["registry_digest"] = digest
json.dump(data, open(receipt, "w"), indent=2, sort_keys=True)
open(receipt, "a").write("\n")
PY
  note "admitted $digest"
}

# The deployment-side gate. Anything not carrying an admitted digest is
# refused, which is exactly what was missing when e1926f3 shipped.
cmd_verify() {
  local digest="${1:-}"
  [ -n "$digest" ] || die "usage: scripts/admit verify <sha256:...>"
  python3 - "$RECEIPTS" "$digest" <<'PY'
import json, pathlib, sys
receipts, digest = pathlib.Path(sys.argv[1]), sys.argv[2]
for path in sorted(receipts.glob("*.json")):
    data = json.load(open(path))
    if data.get("image", {}).get("registry_digest") != digest:
        continue
    suites = data.get("suites", {})
    bad = {k: v for k, v in suites.items() if v != "pass"}
    if bad:
        print(f"REJECT {digest}: {path.name} has non-passing suites: {bad}")
        raise SystemExit(1)
    print(f"ADMITTED {digest} ({path.name}, stream {data['revisions']['stream'][:12]})")
    raise SystemExit(0)
print(f"REJECT {digest}: no receipt admits this digest")
raise SystemExit(1)
PY
}

case "${1:-}" in
  build) shift; cmd_build "$@" ;;
  run) shift; cmd_run "$@" ;;
  publish) shift; cmd_publish "$@" ;;
  verify) shift; cmd_verify "$@" ;;
  *) die "usage: scripts/admit {build|run|publish [sha]|verify <digest>}" ;;
esac