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

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

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
  install-macos111-host-boot-readback.sh --apply

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

while [ "$#" -gt 0 ]; do
  case "$1" in
    --check)
      MODE="check"
      ;;
    --apply)
      MODE="apply"
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "UNKNOWN_ARG=$1" >&2
      usage >&2
      exit 64
      ;;
  esac
  shift
done

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"
}

check() {
  local direct_output control_output
  echo "AWOOOI_MACOS111_BOOT_READBACK=1"
  echo "MODE=check"
  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"
  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"
  echo "MACOS111_CONTROL_PATH_READY=1"
}

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

RUN_ID="macos111-boot-readback-$(date +%Y%m%dT%H%M%S%z)-$$"
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}"
ssh "${SSH_OPTS[@]}" "$TARGET_HOST" "install -d -m 0700 '$remote_stage'"
scp -q "${SSH_OPTS[@]}" "$LOCAL_SCRIPT" "$tmp_dir/control.pub" \
  "${TARGET_HOST}:${remote_stage}/"

echo "AWOOOI_MACOS111_BOOT_READBACK=1"
echo "MODE=apply"
echo "RUN_ID=$RUN_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}')"

ssh "${SSH_OPTS[@]}" "$TARGET_HOST" \
  "RUN_ID='$RUN_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"
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

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

check