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

die() {
  echo "error: $*" >&2
  exit 2
}

usage() {
  cat <<'USAGE'
Usage:
  scripts/rerun/deploy-tipoc-binary.sh [options]

Deploys one local TiPOC binary to every node in an external-run inventory after
checking that no node is currently running `tipoc run`.

Options can be supplied as flags or matching environment variables.

Options:
  --source PATH           TIPOC_DEPLOY_SOURCE, clean source worktree
  --local-binary PATH     TIPOC_DEPLOY_LOCAL_BINARY
  --remote-binary PATH    TIPOC_DEPLOY_REMOTE_BINARY
  --node-file PATH        TIPOC_DEPLOY_NODE_FILE
  --ssh-key PATH          TIPOC_DEPLOY_SSH_KEY
  --ssh-user USER         TIPOC_DEPLOY_SSH_USER, default ec2-user
  --expected-arch ARCH    TIPOC_DEPLOY_EXPECTED_ARCH, default aarch64
  --bastion TARGET        TIPOC_DEPLOY_BASTION, required for proxy nodes
  --help

Node inventory format:
  node_id ssh_host ssh_mode private_ip expected_slots expected_sql_pool

ssh_mode must be direct or proxy. The private_ip and expected fields are parsed
for compatibility with external run launchers and are not used by deployment.
USAGE
}

shell_quote() {
  local value="$1"
  printf "'%s'" "${value//\'/\'\"\'\"\'}"
}

proxy_command_quote() {
  local value="$1"
  if [[ "$value" =~ ^[A-Za-z0-9_./:@%+=,-]+$ ]]; then
    printf '%s' "$value"
    return
  fi
  value="${value//\\/\\\\}"
  value="${value//\"/\\\"}"
  value="${value//\$/\\$}"
  value="${value//\`/\\\`}"
  printf '"%s"' "$value"
}

redact_stream() {
  sed -E \
    -e 's#(root:)[^@[:space:]]+(@tcp)#\1<redacted>\2#g' \
    -e 's#([?&](access-key|secret-access-key|session-token|X-Amz-[A-Za-z0-9_-]+)=)[^&[:space:]"'"'"']+#\1<redacted>#g' \
    -e 's#((password|passwd|pwd|secret|token|access_key|secret_key|session_token)=)[^&[:space:]"'"'"']+#\1<redacted>#Ig' \
    -e 's#(TIDB_DSN=).*#\1<redacted>#g'
}

require_file() {
  local path="$1" label="$2"
  [[ -n "$path" ]] || die "$label is required"
  [[ -s "$path" ]] || die "$label is missing or empty: $path"
}

require_path_value() {
  local value="$1" label="$2"
  [[ -n "$value" ]] || die "$label is required"
}

sha256_file() {
  local path="$1"
  if command -v sha256sum >/dev/null 2>&1; then
    sha256sum "$path" | awk '{print $1}'
  elif command -v shasum >/dev/null 2>&1; then
    shasum -a 256 "$path" | awk '{print $1}'
  else
    die "sha256sum or shasum is required"
  fi
}

TIPOC_DEPLOY_LOCAL_BINARY="${TIPOC_DEPLOY_LOCAL_BINARY:-}"
TIPOC_DEPLOY_SOURCE="${TIPOC_DEPLOY_SOURCE:-}"
TIPOC_DEPLOY_REMOTE_BINARY="${TIPOC_DEPLOY_REMOTE_BINARY:-}"
TIPOC_DEPLOY_NODE_FILE="${TIPOC_DEPLOY_NODE_FILE:-}"
TIPOC_DEPLOY_SSH_KEY="${TIPOC_DEPLOY_SSH_KEY:-}"
TIPOC_DEPLOY_SSH_USER="${TIPOC_DEPLOY_SSH_USER:-ec2-user}"
TIPOC_DEPLOY_EXPECTED_ARCH="${TIPOC_DEPLOY_EXPECTED_ARCH:-aarch64}"
TIPOC_DEPLOY_BASTION="${TIPOC_DEPLOY_BASTION:-}"

while [[ "$#" -gt 0 ]]; do
  case "$1" in
    --source)
      [[ "$#" -ge 2 ]] || die "--source requires a value"
      TIPOC_DEPLOY_SOURCE="$2"
      shift 2
      ;;
    --local-binary)
      [[ "$#" -ge 2 ]] || die "--local-binary requires a value"
      TIPOC_DEPLOY_LOCAL_BINARY="$2"
      shift 2
      ;;
    --remote-binary)
      [[ "$#" -ge 2 ]] || die "--remote-binary requires a value"
      TIPOC_DEPLOY_REMOTE_BINARY="$2"
      shift 2
      ;;
    --node-file)
      [[ "$#" -ge 2 ]] || die "--node-file requires a value"
      TIPOC_DEPLOY_NODE_FILE="$2"
      shift 2
      ;;
    --ssh-key)
      [[ "$#" -ge 2 ]] || die "--ssh-key requires a value"
      TIPOC_DEPLOY_SSH_KEY="$2"
      shift 2
      ;;
    --ssh-user)
      [[ "$#" -ge 2 ]] || die "--ssh-user requires a value"
      TIPOC_DEPLOY_SSH_USER="$2"
      shift 2
      ;;
    --expected-arch)
      [[ "$#" -ge 2 ]] || die "--expected-arch requires a value"
      TIPOC_DEPLOY_EXPECTED_ARCH="$2"
      shift 2
      ;;
    --bastion)
      [[ "$#" -ge 2 ]] || die "--bastion requires a value"
      TIPOC_DEPLOY_BASTION="$2"
      shift 2
      ;;
    --help)
      usage
      exit 0
      ;;
    *)
      die "unknown option: $1"
      ;;
  esac
done

require_path_value "$TIPOC_DEPLOY_SOURCE" "source worktree"
require_file "$TIPOC_DEPLOY_LOCAL_BINARY" "local TiPOC binary"
require_file "$TIPOC_DEPLOY_NODE_FILE" "node inventory file"
require_file "$TIPOC_DEPLOY_SSH_KEY" "SSH private key"
require_path_value "$TIPOC_DEPLOY_REMOTE_BINARY" "remote TiPOC binary path"
require_path_value "$TIPOC_DEPLOY_SSH_USER" "ssh user"
require_path_value "$TIPOC_DEPLOY_EXPECTED_ARCH" "expected architecture"
[[ "$TIPOC_DEPLOY_REMOTE_BINARY" == /* ]] || die "remote TiPOC binary path must be absolute"

command -v git >/dev/null 2>&1 || die "git is required"
command -v go >/dev/null 2>&1 || die "go is required"
SOURCE_ROOT="$(git -C "$TIPOC_DEPLOY_SOURCE" rev-parse --show-toplevel 2>/dev/null)" ||
  die "source is not a Git worktree: $TIPOC_DEPLOY_SOURCE"
EXPECTED_REVISION="$(git -C "$SOURCE_ROOT" rev-parse HEAD)"
[[ "$EXPECTED_REVISION" =~ ^[0-9a-f]{40}$ ]] || die "source HEAD is not a full Git SHA"
[[ -z "$(git -C "$SOURCE_ROOT" status --porcelain=v1)" ]] ||
  die "source worktree must be clean: $SOURCE_ROOT"

BUILD_INFO="$(go version -m "$TIPOC_DEPLOY_LOCAL_BINARY")" ||
  die "cannot inspect local binary build info"
EMBEDDED_VCS="$(printf '%s\n' "$BUILD_INFO" | sed -En 's/^[[:space:]]*build[[:space:]]+vcs=//p')"
EMBEDDED_REVISION="$(printf '%s\n' "$BUILD_INFO" | sed -En 's/^[[:space:]]*build[[:space:]]+vcs\.revision=//p')"
EMBEDDED_MODIFIED="$(printf '%s\n' "$BUILD_INFO" | sed -En 's/^[[:space:]]*build[[:space:]]+vcs\.modified=//p')"
[[ "$EMBEDDED_VCS" == "git" ]] || die "local binary vcs is not git"
[[ "$EMBEDDED_REVISION" == "$EXPECTED_REVISION" ]] ||
  die "local binary vcs.revision does not match source HEAD"
[[ "$EMBEDDED_MODIFIED" == "false" ]] || die "local binary vcs.modified is not false"

LOCAL_SHA256="$(sha256_file "$TIPOC_DEPLOY_LOCAL_BINARY")"
SSH_BASE=(-i "$TIPOC_DEPLOY_SSH_KEY" -o IdentitiesOnly=yes -o StrictHostKeyChecking=no -o BatchMode=yes)

ssh_target_for() {
  local host="$1"
  printf '%s@%s\n' "$TIPOC_DEPLOY_SSH_USER" "$host"
}

bastion_target() {
  if [[ "$TIPOC_DEPLOY_BASTION" == *@* ]]; then
    printf '%s\n' "$TIPOC_DEPLOY_BASTION"
  else
    printf '%s@%s\n' "$TIPOC_DEPLOY_SSH_USER" "$TIPOC_DEPLOY_BASTION"
  fi
}

build_node_ssh_args() {
  local mode="$1"
  NODE_SSH_ARGS=("${SSH_BASE[@]}")
  if [[ "$mode" == "proxy" ]]; then
    [[ -n "$TIPOC_DEPLOY_BASTION" ]] || die "TIPOC_DEPLOY_BASTION or --bastion is required for proxy nodes"
    local proxy_cmd="ssh"
    local arg
    for arg in "${SSH_BASE[@]}"; do
      proxy_cmd+=" $(proxy_command_quote "$arg")"
    done
    proxy_cmd+=" -W %h:%p $(proxy_command_quote "$(bastion_target)")"
    NODE_SSH_ARGS+=("-o" "ProxyCommand=$proxy_cmd")
  elif [[ "$mode" != "direct" ]]; then
    die "invalid ssh_mode: $mode"
  fi
}

ssh_node() {
  local host="$1" mode="$2"
  shift 2
  local target rc
  build_node_ssh_args "$mode"
  target="$(ssh_target_for "$host")"
  ssh "${NODE_SSH_ARGS[@]}" "$target" "$@" 2>&1 | redact_stream
  rc="${PIPESTATUS[0]}"
  return "$rc"
}

scp_to_node() {
  local host="$1" mode="$2" local_path="$3" remote_path="$4"
  local target rc
  build_node_ssh_args "$mode"
  target="$(ssh_target_for "$host")"
  scp -q "${NODE_SSH_ARGS[@]}" "$local_path" "$target:$remote_path" 2>&1 | redact_stream
  rc="${PIPESTATUS[0]}"
  return "$rc"
}

remote_preflight() {
  local node="$1" host="$2" mode="$3"
  echo "node=${node} status=preflight"
  ssh_node "$host" "$mode" \
    "bash -s -- $(shell_quote "$TIPOC_DEPLOY_REMOTE_BINARY") $(shell_quote "$TIPOC_DEPLOY_EXPECTED_ARCH")" <<'SH'
set -euo pipefail
remote_binary="$1"
expected_arch="$2"

arch="$(uname -m)"
printf 'node_arch=%s\n' "$arch"
if [[ "$arch" != "$expected_arch" ]]; then
  printf 'architecture_mismatch expected=%s actual=%s\n' "$expected_arch" "$arch" >&2
  exit 71
fi

if ps -eo pid=,args= | awk '$0 ~ /[t]ipoc[[:space:]]+run/ {print; found=1} END {exit found ? 0 : 1}'; then
  echo "active_tipoc_run=found" >&2
  exit 72
fi

mkdir -p "$(dirname "$remote_binary")"
if [[ -e "$remote_binary" && ! -f "$remote_binary" ]]; then
  printf 'remote path exists but is not a file: %s\n' "$remote_binary" >&2
  exit 73
fi
echo "active_tipoc_run=none"
SH
}

remote_activate_binary() {
  local node="$1" host="$2" mode="$3" remote_tmp="$4"
  echo "node=${node} status=activating"
  ssh_node "$host" "$mode" \
    "bash -s -- $(shell_quote "$remote_tmp") $(shell_quote "$TIPOC_DEPLOY_REMOTE_BINARY") $(shell_quote "$LOCAL_SHA256")" <<'SH'
set -euo pipefail
remote_tmp="$1"
remote_binary="$2"
expected_sha="$3"

sha256_of() {
  local path="$1"
  if command -v sha256sum >/dev/null 2>&1; then
    sha256sum "$path" | awk '{print $1}'
  elif command -v shasum >/dev/null 2>&1; then
    shasum -a 256 "$path" | awk '{print $1}'
  else
    echo "sha256sum or shasum is required on remote node" >&2
    exit 74
  fi
}

if ps -eo pid=,args= | awk '$0 ~ /[t]ipoc[[:space:]]+run/ {print; found=1} END {exit found ? 0 : 1}'; then
  rm -f "$remote_tmp"
  echo "active_tipoc_run=found_before_replace" >&2
  exit 72
fi

chmod 0755 "$remote_tmp"
mv "$remote_tmp" "$remote_binary"
actual_sha="$(sha256_of "$remote_binary")"
printf 'remote_sha256=%s\n' "$actual_sha"
if [[ "$actual_sha" != "$expected_sha" ]]; then
  printf 'sha256_mismatch expected=%s actual=%s\n' "$expected_sha" "$actual_sha" >&2
  exit 75
fi
if command -v file >/dev/null 2>&1; then
  file "$remote_binary" || true
fi
echo "deploy_status=ok"
SH
}

NODE_IDS=()
NODE_HOSTS=()
NODE_MODES=()

while read -r node host mode private_ip expected_slots expected_sql_pool extra; do
  [[ -z "${node:-}" || "${node:0:1}" == "#" ]] && continue
  [[ -z "${extra:-}" ]] || die "node inventory record has too many fields: $node"
  [[ -n "${host:-}" && -n "${mode:-}" && -n "${private_ip:-}" && -n "${expected_slots:-}" && -n "${expected_sql_pool:-}" ]] ||
    die "node inventory record has too few fields: $node"
  [[ "$mode" == "direct" || "$mode" == "proxy" ]] || die "node $node has invalid ssh_mode: $mode"
  NODE_IDS+=("$node")
  NODE_HOSTS+=("$host")
  NODE_MODES+=("$mode")
done < "$TIPOC_DEPLOY_NODE_FILE"

[[ "${#NODE_IDS[@]}" -gt 0 ]] || die "node inventory did not contain any nodes"

printf 'local_binary=%s\n' "$TIPOC_DEPLOY_LOCAL_BINARY"
printf 'source_revision=%s\n' "$EXPECTED_REVISION"
printf 'binary_vcs=%s\n' "$EMBEDDED_VCS"
printf 'binary_vcs_modified=%s\n' "$EMBEDDED_MODIFIED"
printf 'remote_binary=%s\n' "$TIPOC_DEPLOY_REMOTE_BINARY"
printf 'expected_arch=%s\n' "$TIPOC_DEPLOY_EXPECTED_ARCH"
printf 'local_sha256=%s\n' "$LOCAL_SHA256"
printf 'node_count=%s\n' "${#NODE_IDS[@]}"

for idx in "${!NODE_IDS[@]}"; do
  remote_preflight "${NODE_IDS[$idx]}" "${NODE_HOSTS[$idx]}" "${NODE_MODES[$idx]}"
done

for idx in "${!NODE_IDS[@]}"; do
  node="${NODE_IDS[$idx]}"
  host="${NODE_HOSTS[$idx]}"
  mode="${NODE_MODES[$idx]}"
  remote_tmp="${TIPOC_DEPLOY_REMOTE_BINARY}.tmp.${node}.$$"
  echo "node=${node} status=uploading"
  scp_to_node "$host" "$mode" "$TIPOC_DEPLOY_LOCAL_BINARY" "$remote_tmp"
  remote_activate_binary "$node" "$host" "$mode" "$remote_tmp"
done

echo "deploy_status=all_nodes_ok"