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

# ---------------- Defaults ----------------
TIMEOUT_MIN=2
HOST=""
ACTION=""
UPDATE_REMOTE=0

SCRIPT_PATH="/usr/local/bin/vpnswitch"
REMOTE_USER="mark"
SSH_OPTS=(-o BatchMode=yes -o ConnectTimeout=5)

HEADSCALE_URL="https://headscale.neland.dk"

# Lokal fil med nøgler (kun brugt ved upload/-u i controller-mode)
KEYS_FILE="${KEYS_FILE:-$HOME/keys.txt}"

usage() {
  cat <<EOF
Usage:
  vpnswitch on      [-h <host|number>] [-t minutes] [-u]
  vpnswitch confirm [-h <host|number>] [-u]
  vpnswitch off     [-h <host|number>] [-u]
  vpnswitch status  [-h <host|number>]
  vpnswitch upload  -h <host|number>

Modes:
  - Without -h: local mode (runs on this machine)
  - With -h:    controller mode (SSH to target host as ${REMOTE_USER})

Options:
  -h   Hostname or number (2 → aqualog2). Omit for local mode.
  -t   Rollback timeout in minutes (default: 2) [only for 'on']
  -u   Upload/update vpnswitch on target host before action (controller-mode only)

Keys:
  Upload embeds keys from $KEYS_FILE into remote script:
    HS_KEY=...
    TS_KEY=...
  Override with: KEYS_FILE=/path/to/keys.txt

Examples:
  # controller mode
  vpnswitch on -h 2 -t 2 -u
  vpnswitch confirm -h 2
  vpnswitch upload -h 2

  # local mode (run directly on the node)
  vpnswitch on -t 2
  vpnswitch confirm
EOF
  exit 1
}

die() { echo "ERROR: $*" >&2; exit 1; }

normalize_host() {
  local h="$1"
  if [[ "$h" =~ ^[0-9]+$ ]]; then
    echo "aqualog${h}"
  else
    echo "$h"
  fi
}

unit_name_for_host() {
  local h="$1"
  h="${h//./_}"
  h="${h//-/_}"
  echo "vpnswitch-rollback-${h}"
}

# ---------------- Parse action ----------------
[[ $# -ge 1 ]] || usage
ACTION="$1"
shift

case "$ACTION" in
  on|off|confirm|status|upload) ;;
  *) die "Invalid action: $ACTION" ;;
esac

# ---------------- Parse options ----------------
while getopts ":h:t:u" opt; do
  case "$opt" in
    h) HOST="$OPTARG" ;;
    t) TIMEOUT_MIN="$OPTARG" ;;
    u) UPDATE_REMOTE=1 ;;
    :) die "Option -$OPTARG requires an argument" ;;
    \?) die "Unknown option: -$OPTARG" ;;
  esac
done

if ! [[ "$TIMEOUT_MIN" =~ ^[0-9]+$ ]]; then
  die "Timeout (-t) must be an integer number of minutes"
fi

LOCAL_MODE=0
TARGET_HOST=""

if [[ -z "$HOST" ]]; then
  LOCAL_MODE=1
  TARGET_HOST="localhost"
else
  TARGET_HOST="$(normalize_host "$HOST")"
fi

# Unit name: unique per target host in controller mode; per local hostname in local mode
if [[ "$LOCAL_MODE" -eq 1 ]]; then
  UNIT_NAME="$(unit_name_for_host "$(hostname -s)")"
else
  UNIT_NAME="$(unit_name_for_host "$TARGET_HOST")"
fi

# ---------------- Execution primitives ----------------

run_script() {
  local host="$1"
  local script="$2"

  if [[ "$LOCAL_MODE" -eq 1 ]]; then
    bash -s <<EOF
set -euo pipefail
$script
EOF
  else
    ssh "${SSH_OPTS[@]}" "${REMOTE_USER}@${host}" "bash -s" <<EOF
set -euo pipefail
$script
EOF
  fi
}

scp_to_target() {
  local src="$1"
  local dst="$2"

  if [[ "$LOCAL_MODE" -eq 1 ]]; then
    # Local "upload" means install on local machine (requires sudo)
    cp -f "$src" "$dst"
  else
    scp "${SSH_OPTS[@]}" "$src" "${REMOTE_USER}@${TARGET_HOST}:$dst" >/dev/null
  fi
}

# ---------------- Keys handling (controller upload) ----------------
load_keys_local() {
  [[ -f "$KEYS_FILE" ]] || die "Keys file not found: $KEYS_FILE"

  # shellcheck disable=SC1090
  source "$KEYS_FILE"

  [[ -n "${HS_KEY:-}" ]] || die "HS_KEY missing in $KEYS_FILE"
  [[ -n "${TS_KEY:-}" ]] || die "TS_KEY missing in $KEYS_FILE"
}

build_script_with_keys() {
  load_keys_local

  local tmp hsq tsq
  tmp="$(mktemp /tmp/vpnswitch.embedkeys.XXXXXX)"

  # Shell-escape keys safely
  hsq="$(printf "%q" "$HS_KEY")"
  tsq="$(printf "%q" "$TS_KEY")"

  # Insert keys AFTER "set -euo pipefail" so:
  # - shebang stays on line 1
  # - keys are in the script body
  # - keys are exported for subshells (bash -s)
  awk -v HS="$hsq" -v TS="$tsq" '
    {
      print $0
      if (!inserted && $0 ~ /^set -euo pipefail[[:space:]]*$/) {
        print ""
        print "# --- BEGIN embedded keys (generated locally) ---"
        print "HS_KEY=" HS
        print "TS_KEY=" TS
        print "export HS_KEY TS_KEY"
        print "# --- END embedded keys ---"
        print ""
        inserted=1
      }
    }
    END {
      if (!inserted) {
        print "ERROR: Could not find line: set -euo pipefail" > "/dev/stderr"
        exit 2
      }
    }
  ' "$0" > "$tmp"

  chmod 600 "$tmp"
  echo "$tmp"
}

# ---------------- Upload/update ----------------
upload_to_target() {
  # upload requires a target host (controller) OR local install (local mode)
  if [[ "$LOCAL_MODE" -eq 1 ]]; then
    die "upload requires -h <host|number> (controller-mode) to push to another node"
  fi

  local remote_tmp="/tmp/vpnswitch.$$.tmp"
  local local_tmp
  local_tmp="$(build_script_with_keys)"

  echo "Uploading vpnswitch to ${TARGET_HOST} as ${REMOTE_USER} (sudo install)..."
  scp "${SSH_OPTS[@]}" "$local_tmp" "${REMOTE_USER}@${TARGET_HOST}:${remote_tmp}" >/dev/null

  run_script "$TARGET_HOST" "
    sudo install -m 0755 '${remote_tmp}' '${SCRIPT_PATH}'
    rm -f '${remote_tmp}'
  "

  rm -f "$local_tmp"
}

maybe_update_remote() {
  # -u is meaningful only in controller-mode and only for on/off/confirm
  if [[ "$UPDATE_REMOTE" -eq 1 ]]; then
    if [[ "$LOCAL_MODE" -eq 1 ]]; then
      echo "Note: -u ignored in local mode"
    else
      upload_to_target
    fi
  fi
}

cleanup_unit() {
  local unit="$1"

  run_script "$TARGET_HOST" "
    sudo systemctl stop '${unit}.timer' '${unit}.service' 2>/dev/null || true
    sudo systemctl disable '${unit}.timer' '${unit}.service' 2>/dev/null || true
    sudo systemctl reset-failed '${unit}.timer' '${unit}.service' 2>/dev/null || true

    # Remove transient units if they exist
    sudo rm -f '/run/systemd/transient/${unit}.timer' '/run/systemd/transient/${unit}.service' 2>/dev/null || true

    # Remove fragment unit files if they exist (should normally not exist, but clean anyway)
    sudo rm -f '/etc/systemd/system/${unit}.timer' '/etc/systemd/system/${unit}.service' 2>/dev/null || true

    sudo systemctl daemon-reload
  "
}

# ---------------- Actions ----------------
schedule_rollback() {
  local unit="$1"
  local timeout="$2"

  echo "Scheduling rollback to Tailscale in ${timeout} minutes (unit: ${unit})..."

  cleanup_unit "$unit"

  run_script "$TARGET_HOST" "
    sudo systemd-run --unit='${unit}' --on-active='${timeout}min' '${SCRIPT_PATH}' off
  "
}

cancel_rollback() {
  local unit="$1"
  echo "Cancelling rollback (unit: ${unit})..."
  cleanup_unit "$unit"
}

do_on() {
  maybe_update_remote

  schedule_rollback "$UNIT_NAME" "$TIMEOUT_MIN"

  echo "Switching to Headscale: ${HEADSCALE_URL}"
  run_script "$TARGET_HOST" "
    [[ -n \"\${HS_KEY:-}\" ]] || { echo 'HS_KEY missing (did you upload with embedded keys?)'; exit 1; }
    sudo tailscale logout || true
    echo tailscale up --login-server '${HEADSCALE_URL}' --authkey \"\${HS_KEY}\"
    sudo tailscale up --login-server '${HEADSCALE_URL}' --authkey \"\${HS_KEY}\"
    sudo tailscale status
  "

  echo "NOTE: Run 'vpnswitch confirm' to cancel rollback."
}

do_off() {
  maybe_update_remote

  cancel_rollback "$UNIT_NAME"

  echo "Switching to Tailscale (default login-server)"
  run_script "$TARGET_HOST" "
    [[ -n \"\${TS_KEY:-}\" ]] || { echo 'TS_KEY missing (did you upload with embedded keys?)'; exit 1; }
    sudo tailscale logout || true
    echo tailscale up --authkey \"\${TS_KEY}\"
    sudo tailscale up --authkey \"\${TS_KEY}\"
    sudo tailscale status
  "
}

do_confirm() {
  maybe_update_remote
  cancel_rollback "$UNIT_NAME"
  echo "Confirmed: rollback cancelled."
}

do_status() {
  echo "Status for ${TARGET_HOST}"
  run_script "$TARGET_HOST" "
    echo '--- tailscale status (summary) ---'
    sudo tailscale status 2>/dev/null | head -n 30 || true
    echo
    echo '--- current login-server (ControlURL) ---'
    sudo tailscale debug prefs 2>/dev/null | sed -n 's/.*\"ControlURL\": \"\\(.*\\)\".*/\\1/p' | head -n 1 || true
    echo
    echo '--- rollback timer ---'
    sudo systemctl list-timers --all | grep -F '${UNIT_NAME}.timer' || echo 'No rollback timer scheduled'
  "
}

# ---------------- Dispatch ----------------
case "$ACTION" in
  upload)
    # upload must have -h
    [[ -n "$HOST" ]] || die "upload requires -h <host|number>"
    upload_to_target
    ;;
  on)
    do_on
    ;;
  off)
    do_off
    ;;
  confirm)
    do_confirm
    ;;
  status)
    do_status
    ;;
esac
