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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
MODE="check"
TRACE_ID=""
RUN_ID=""
WORK_ITEM_ID=""

CONTROL_HOST="${CONTROL_HOST:-wooo@192.168.0.110}"
CONTROL_PUBLIC_KEY_PATH="${CONTROL_PUBLIC_KEY_PATH:-/home/wooo/.ssh/id_ed25519.pub}"
TARGET_HOST="${TARGET_HOST:-ollama-111-gpu}"
TARGET_ADDRESS="${TARGET_ADDRESS:-192.168.0.111}"
TARGET_USER="${TARGET_USER:-ooo}"
REMOTE_SCRIPT="${REMOTE_SCRIPT:-/Users/ooo/.local/bin/awoooi-host-boot-readback}"
LOCAL_SCRIPT="${ROOT_DIR}/scripts/reboot-recovery/macos-host-boot-readback.sh"

SSH_OPTS=(
  -o BatchMode=yes
  -o ConnectTimeout="${SSH_CONNECT_TIMEOUT_SECONDS:-8}"
  -o ConnectionAttempts=1
  -o NumberOfPasswordPrompts=0
)

usage() {
  cat <<'USAGE'
Usage:
  install-macos111-host-boot-readback.sh --check [controlled identity]
  install-macos111-host-boot-readback.sh --apply --trace-id ID --run-id ID --work-item-id ID

Installs a no-secret Darwin boot and performance readback and authorizes host
110's existing public key with a forced command. The key cannot open a shell,
request a PTY, or open forwarding.
USAGE
}

require_value() {
  local flag="$1" value="${2:-}"
  if [ -z "$value" ]; then
    echo "BLOCKER missing_value=$flag" >&2
    exit 64
  fi
}

safe_identity() {
  [[ "$1" =~ ^[A-Za-z0-9][A-Za-z0-9._:@+-]{0,127}$ ]]
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --check)
      MODE="check"
      ;;
    --apply)
      MODE="apply"
      ;;
    --trace-id)
      require_value "$1" "${2:-}"
      TRACE_ID="$2"
      shift
      ;;
    --run-id)
      require_value "$1" "${2:-}"
      RUN_ID="$2"
      shift
      ;;
    --work-item-id)
      require_value "$1" "${2:-}"
      WORK_ITEM_ID="$2"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "UNKNOWN_ARG=$1" >&2
      usage >&2
      exit 64
      ;;
  esac
  shift
done

identity_count=0
[ -n "$TRACE_ID" ] && identity_count=$((identity_count + 1))
[ -n "$RUN_ID" ] && identity_count=$((identity_count + 1))
[ -n "$WORK_ITEM_ID" ] && identity_count=$((identity_count + 1))
if [ "$identity_count" -ne 0 ] && [ "$identity_count" -ne 3 ]; then
  echo "BLOCKER incomplete_control_identity" >&2
  exit 64
fi
if [ "$identity_count" -eq 3 ]; then
  for identity in "$TRACE_ID" "$RUN_ID" "$WORK_ITEM_ID"; do
    if ! safe_identity "$identity"; then
      echo "BLOCKER unsafe_control_identity" >&2
      exit 64
    fi
  done
fi
if [ "$MODE" = "apply" ] && [ "$identity_count" -ne 3 ]; then
  echo "BLOCKER controlled_identity_required" >&2
  exit 64
fi

if [ ! -f "$LOCAL_SCRIPT" ]; then
  echo "BLOCKER macos_boot_readback_source_missing"
  exit 66
fi

control_path_check() {
  ssh "${SSH_OPTS[@]}" "$CONTROL_HOST" \
    "ssh -i /home/wooo/.ssh/id_ed25519 -o IdentitiesOnly=yes -o BatchMode=yes -o StrictHostKeyChecking=yes -o UserKnownHostsFile=/home/wooo/.ssh/known_hosts -o ConnectTimeout=6 -o ConnectionAttempts=1 -o NumberOfPasswordPrompts=0 '${TARGET_USER}@${TARGET_ADDRESS}' ignored-by-forced-command"
}

validate_readback() {
  local output="$1"
  grep -Eq '^boot_id=darwin_[0-9]+ uptime_seconds=[0-9]+ systemd_state=darwin_ssh startup_enabled=(enabled|unknown) startup_active=(active|inactive)$' <<<"$output" \
    && grep -Eq '^perf_schema=agent99_perf_readback_v1 os=darwin cores=[1-9][0-9]* load1=[0-9]+([.][0-9]+)? mem_available_percent=[0-9]+([.][0-9]+)? disk_used_percent=[0-9]+([.][0-9]+)?$' <<<"$output"
}

check() {
  local direct_output control_output
  echo "AWOOOI_MACOS111_BOOT_READBACK=1"
  echo "MODE=check"
  echo "TRACE_ID=${TRACE_ID:-none}"
  echo "RUN_ID=${RUN_ID:-none}"
  echo "WORK_ITEM_ID=${WORK_ITEM_ID:-none}"
  if ! direct_output="$(ssh "${SSH_OPTS[@]}" "$TARGET_HOST" "'$REMOTE_SCRIPT'")"; then
    echo "MACOS111_READBACK_SCRIPT_READY=0"
    echo "BLOCKER macos111_readback_script_unavailable"
    return 75
  fi
  echo "$direct_output"
  if ! validate_readback "$direct_output"; then
    echo "MACOS111_READBACK_SCRIPT_READY=0"
    echo "BLOCKER macos111_readback_schema_incomplete"
    return 75
  fi
  echo "MACOS111_READBACK_SCRIPT_READY=1"

  if ! control_output="$(control_path_check)"; then
    echo "MACOS111_CONTROL_PATH_READY=0"
    echo "BLOCKER macos111_restricted_key_unavailable"
    return 75
  fi
  echo "$control_output"
  if ! validate_readback "$control_output"; then
    echo "MACOS111_CONTROL_PATH_READY=0"
    echo "BLOCKER macos111_control_readback_schema_incomplete"
    return 75
  fi
  echo "MACOS111_CONTROL_PATH_READY=1"
}

rollback_apply() {
  ssh "${SSH_OPTS[@]}" "$TARGET_HOST" \
    "TRACE_ID='$TRACE_ID' RUN_ID='$RUN_ID' WORK_ITEM_ID='$WORK_ITEM_ID' REMOTE_SCRIPT='$REMOTE_SCRIPT' bash -s" <<'REMOTE_ROLLBACK'
set -euo pipefail

auth_file="$HOME/.ssh/authorized_keys"
backup_dir="$HOME/.ssh/awoooi-backups/$RUN_ID"
rollback_verified=1

if [ -f "$backup_dir/authorized_keys" ]; then
  cp -p "$backup_dir/authorized_keys" "$auth_file"
  cmp -s "$backup_dir/authorized_keys" "$auth_file" || rollback_verified=0
elif [ -f "$backup_dir/authorized_keys.missing" ]; then
  rm -f "$auth_file"
  [ ! -e "$auth_file" ] || rollback_verified=0
else
  rollback_verified=0
fi

if [ -f "$backup_dir/macos-host-boot-readback.sh" ]; then
  cp -p "$backup_dir/macos-host-boot-readback.sh" "$REMOTE_SCRIPT"
  cmp -s "$backup_dir/macos-host-boot-readback.sh" "$REMOTE_SCRIPT" || rollback_verified=0
elif [ -f "$backup_dir/macos-host-boot-readback.sh.missing" ]; then
  rm -f "$REMOTE_SCRIPT"
  [ ! -e "$REMOTE_SCRIPT" ] || rollback_verified=0
else
  rollback_verified=0
fi

receipt="$backup_dir/rollback-receipt.txt"
receipt_tmp="${receipt}.tmp.$$"
[ ! -e "$receipt" ] || {
  echo "BLOCKER rollback_receipt_already_exists"
  exit 65
}
terminal="rollback_unverified"
[ "$rollback_verified" -eq 1 ] && terminal="rolled_back_verified"
printf 'schema_version=macos111_readback_rollback_receipt_v1 trace_id=%s run_id=%s work_item_id=%s terminal=%s rollback_verified=%s secret_value_read=0 private_key_value_read=0\n' \
  "$TRACE_ID" "$RUN_ID" "$WORK_ITEM_ID" "$terminal" "$rollback_verified" >"$receipt_tmp"
chmod 0600 "$receipt_tmp"
mv "$receipt_tmp" "$receipt"
echo "ROLLBACK_RECEIPT=$receipt"
echo "ROLLBACK_VERIFIED=$rollback_verified"
[ "$rollback_verified" -eq 1 ]
REMOTE_ROLLBACK
}

write_verified_receipt() {
  ssh "${SSH_OPTS[@]}" "$TARGET_HOST" \
    "TRACE_ID='$TRACE_ID' RUN_ID='$RUN_ID' WORK_ITEM_ID='$WORK_ITEM_ID' REMOTE_SCRIPT='$REMOTE_SCRIPT' bash -s" <<'REMOTE_RECEIPT'
set -euo pipefail

backup_dir="$HOME/.ssh/awoooi-backups/$RUN_ID"
receipt="$backup_dir/verifier-receipt.txt"
receipt_tmp="${receipt}.tmp.$$"
[ -d "$backup_dir" ] || {
  echo "BLOCKER apply_backup_missing"
  exit 65
}
[ ! -e "$receipt" ] || {
  echo "BLOCKER verifier_receipt_already_exists"
  exit 65
}
script_sha256="$(shasum -a 256 "$REMOTE_SCRIPT" | awk '{print $1}')"
printf 'schema_version=macos111_readback_apply_receipt_v1 trace_id=%s run_id=%s work_item_id=%s terminal=verified_ready script_sha256=%s rollback_available=1 secret_value_read=0 private_key_value_read=0\n' \
  "$TRACE_ID" "$RUN_ID" "$WORK_ITEM_ID" "$script_sha256" >"$receipt_tmp"
chmod 0600 "$receipt_tmp"
mv "$receipt_tmp" "$receipt"
echo "CONTROLLED_APPLY_RECEIPT=$receipt"
echo "CONTROLLED_APPLY_TERMINAL=verified_ready"
REMOTE_RECEIPT
}

if [ "$MODE" = "check" ]; then
  check
  exit $?
fi

tmp_dir="$(mktemp -d "${TMPDIR:-/tmp}/macos111-readback.XXXXXX")"
trap 'rm -rf "$tmp_dir"' EXIT

scp -q "${SSH_OPTS[@]}" \
  "${CONTROL_HOST}:${CONTROL_PUBLIC_KEY_PATH}" "$tmp_dir/control.pub"
ssh-keygen -lf "$tmp_dir/control.pub" >/dev/null

remote_stage="/Users/ooo/.cache/${RUN_ID}"
if ! preflight_output="$(ssh "${SSH_OPTS[@]}" "$TARGET_HOST" \
  "if [ -e '$remote_stage' ] || [ -e \"\$HOME/.ssh/awoooi-backups/$RUN_ID\" ]; then echo CONTROLLED_RUN_EXISTS; exit 65; fi; echo CONTROLLED_RUN_READY" 2>&1)"; then
  if grep -q '^CONTROLLED_RUN_EXISTS$' <<<"$preflight_output"; then
    echo "BLOCKER controlled_run_already_exists"
    exit 65
  fi
  echo "BLOCKER macos111_apply_preflight_transport_failed"
  exit 76
fi
if [ "$preflight_output" != "CONTROLLED_RUN_READY" ]; then
  echo "BLOCKER macos111_apply_preflight_unverified"
  exit 76
fi
if ! ssh "${SSH_OPTS[@]}" "$TARGET_HOST" "install -d -m 0700 '$remote_stage'"; then
  echo "BLOCKER macos111_stage_create_failed"
  exit 76
fi
if ! scp -q "${SSH_OPTS[@]}" "$LOCAL_SCRIPT" "$tmp_dir/control.pub" \
  "${TARGET_HOST}:${remote_stage}/"; then
  ssh "${SSH_OPTS[@]}" "$TARGET_HOST" "rm -rf '$remote_stage'" || true
  echo "BLOCKER macos111_stage_transfer_failed"
  echo "NO_WRITE_TERMINAL=staging_failed_cleanup_attempted"
  exit 76
fi

echo "AWOOOI_MACOS111_BOOT_READBACK=1"
echo "MODE=apply"
echo "TRACE_ID=$TRACE_ID"
echo "RUN_ID=$RUN_ID"
echo "WORK_ITEM_ID=$WORK_ITEM_ID"
echo "LOCAL_SCRIPT_SHA256=$(shasum -a 256 "$LOCAL_SCRIPT" | awk '{print $1}')"
echo "CONTROL_KEY_FINGERPRINT=$(ssh-keygen -lf "$tmp_dir/control.pub" | awk '{print $2}')"

if ! ssh "${SSH_OPTS[@]}" "$TARGET_HOST" \
  "TRACE_ID='$TRACE_ID' RUN_ID='$RUN_ID' WORK_ITEM_ID='$WORK_ITEM_ID' REMOTE_STAGE='$remote_stage' REMOTE_SCRIPT='$REMOTE_SCRIPT' bash -s" <<'REMOTE'
set -euo pipefail

auth_dir="$HOME/.ssh"
auth_file="$auth_dir/authorized_keys"
backup_dir="$auth_dir/awoooi-backups/$RUN_ID"
[ ! -e "$backup_dir" ] || {
  echo "BLOCKER controlled_run_already_exists"
  exit 65
}
install -d -m 0700 "$auth_dir" "$backup_dir" "$(dirname "$REMOTE_SCRIPT")"
if [ -f "$auth_file" ]; then
  cp -p "$auth_file" "$backup_dir/authorized_keys"
else
  : >"$backup_dir/authorized_keys.missing"
  : >"$auth_file"
fi
if [ -f "$REMOTE_SCRIPT" ]; then
  cp -p "$REMOTE_SCRIPT" "$backup_dir/macos-host-boot-readback.sh"
else
  : >"$backup_dir/macos-host-boot-readback.sh.missing"
fi

install -m 0755 "$REMOTE_STAGE/macos-host-boot-readback.sh" "$REMOTE_SCRIPT"
read -r key_type key_blob _ <"$REMOTE_STAGE/control.pub"
case "$key_type" in
  ssh-ed25519|ssh-rsa|ecdsa-sha2-*) ;;
  *)
    echo "BLOCKER unsupported_control_public_key_type"
    exit 65
    ;;
esac
[[ "$key_blob" =~ ^[A-Za-z0-9+/=]+$ ]] || {
  echo "BLOCKER invalid_control_public_key_blob"
  exit 65
}

auth_tmp="${auth_file}.tmp.$$"
awk -v blob="$key_blob" 'index($0, blob) == 0' "$auth_file" >"$auth_tmp"
printf 'restrict,command="%s" %s %s awoooi-host-probe-110\n' \
  "$REMOTE_SCRIPT" "$key_type" "$key_blob" >>"$auth_tmp"
chmod 0600 "$auth_tmp"
mv "$auth_tmp" "$auth_file"
chmod 0600 "$auth_file"
rm -rf "$REMOTE_STAGE"

echo "AUTHORIZED_KEYS_BACKUP=$backup_dir"
echo "READBACK_SCRIPT_SHA256=$(shasum -a 256 "$REMOTE_SCRIPT" | awk '{print $1}')"
echo "RESTRICTED_AUTHORIZED_KEY_INSTALLED=1"
REMOTE
then
  echo "BLOCKER macos111_apply_transport_failed"
  rollback_apply || true
  exit 76
fi

if ! check; then
  echo "BLOCKER macos111_post_verifier_failed"
  rollback_apply || exit 77
  exit 75
fi
if ! write_verified_receipt; then
  echo "BLOCKER macos111_durable_receipt_failed"
  rollback_apply || exit 77
  exit 76
fi