#!/usr/bin/env bash
set -Eeuo pipefail
IFS=$'\n\t'

readonly PROGRAM="vps-agent"
readonly BINARY="/usr/local/bin/vps-agent"
readonly PREVIOUS="/usr/local/lib/vps-agent/vps-agent.previous"
readonly CONFIG_DIR="/etc/vps-agent"
readonly CREDENTIAL_DIR="${CONFIG_DIR}/credentials"
readonly STATE_DIR="/var/lib/vps-agent"
readonly UNIT="/etc/systemd/system/vps-agent.service"
readonly DEFAULT_RELEASES="https://github.com/vps-control/vps-control/releases"

VERSION="latest"
RELEASES_URL="${VPS_AGENT_RELEASES_URL:-$DEFAULT_RELEASES}"
PURGE=false
ASSUME_YES=false
TMP_DIR=""
DOWNLOADED_BINARY=""

usage() {
  printf '%s\n' "Usage: $0 <install|update|remove|rollback|status|restart|version> [options]" \
    "  --version VERSION     release tag, or latest" \
    "  --releases-url URL    releases root (default: VPS_AGENT_RELEASES_URL)" \
    "  --purge               remove configuration and credentials" \
    "  --yes                 confirm destructive purge"
}

die() { printf 'error: %s\n' "$*" >&2; exit 1; }
info() { printf '==> %s\n' "$*"; }

cleanup() {
  if [[ -n "$TMP_DIR" && -d "$TMP_DIR" ]]; then
    rm -rf -- "$TMP_DIR"
  fi
}
trap cleanup EXIT

require_root() { [[ ${EUID:-$(id -u)} -eq 0 ]] || die "run this action as root"; }

check_host() {
  [[ "$(uname -s)" == "Linux" ]] || die "only Linux is supported"
  command -v systemctl >/dev/null || die "systemd is required"
  [[ -d /run/systemd/system ]] || die "systemd is not running"
}

release_arch() {
  case "$(uname -m)" in
    x86_64|amd64) printf 'amd64\n' ;;
    aarch64|arm64) printf 'arm64\n' ;;
    *) die "unsupported architecture: $(uname -m)" ;;
  esac
}

download_release() {
  command -v curl >/dev/null || die "curl is required"
  command -v sha256sum >/dev/null || die "sha256sum is required"
  TMP_DIR="$(mktemp -d -t vps-agent.XXXXXXXX)"
  chmod 700 "$TMP_DIR"
  local artifact="vps-agent-linux-$(release_arch)"
  local release_url
  if [[ "$VERSION" == "latest" ]]; then
    release_url="${RELEASES_URL%/}/latest/download"
  else
    [[ "$VERSION" =~ ^v?[0-9]+\.[0-9]+\.[0-9]+([.-][A-Za-z0-9.-]+)?$ ]] || die "invalid release version"
    release_url="${RELEASES_URL%/}/download/${VERSION}"
  fi
  info "downloading ${artifact} (${VERSION})"
  curl --fail --silent --show-error --location --proto '=https' --tlsv1.2 \
    --output "$TMP_DIR/$artifact" "$release_url/$artifact"
  curl --fail --silent --show-error --location --proto '=https' --tlsv1.2 \
    --output "$TMP_DIR/SHA256SUMS" "$release_url/SHA256SUMS"

  local expected actual
  expected="$(awk -v name="$artifact" '$2 == name || $2 == "*" name { print $1; exit }' "$TMP_DIR/SHA256SUMS")"
  [[ "$expected" =~ ^[a-fA-F0-9]{64}$ ]] || die "artifact checksum is absent or malformed"
  actual="$(sha256sum "$TMP_DIR/$artifact" | awk '{print $1}')"
  [[ "$actual" == "$expected" ]] || die "checksum verification failed"

  if [[ -n "${VPS_AGENT_MINISIGN_PUBLIC_KEY:-}" ]]; then
    command -v minisign >/dev/null || die "minisign is required when a public key is configured"
    curl --fail --silent --show-error --location --proto '=https' --tlsv1.2 \
      --output "$TMP_DIR/SHA256SUMS.minisig" "$release_url/SHA256SUMS.minisig"
    minisign -Vm "$TMP_DIR/SHA256SUMS" -P "$VPS_AGENT_MINISIGN_PUBLIC_KEY"
  fi
  chmod 0755 "$TMP_DIR/$artifact"
  "$TMP_DIR/$artifact" --version >/dev/null || die "downloaded binary self-check failed"
  DOWNLOADED_BINARY="$TMP_DIR/$artifact"
}

install_unit() {
  if [[ -f "$(cd "$(dirname "$0")" && pwd)/../agent/packaging/vps-agent.service" ]]; then
    install -o root -g root -m 0644 \
      "$(cd "$(dirname "$0")" && pwd)/../agent/packaging/vps-agent.service" "$UNIT"
    return
  fi
  install -o root -g root -m 0644 /dev/null "$UNIT"
  tee "$UNIT" >/dev/null <<'EOF'
[Unit]
Description=VPS Control monitoring agent
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=vps-agent
Group=vps-agent
ExecStart=/usr/local/bin/vps-agent --config /etc/vps-agent/config.toml
Restart=on-failure
RestartSec=5s
TimeoutStopSec=15s
MemoryHigh=48M
MemoryMax=64M
CPUQuota=10%
TasksMax=64
NoNewPrivileges=true
PrivateTmp=true
PrivateDevices=true
ProtectHome=true
ProtectSystem=strict
ProtectKernelTunables=true
ProtectKernelModules=true
ProtectKernelLogs=true
ProtectControlGroups=true
ProtectClock=true
ProtectHostname=true
RestrictSUIDSGID=true
LockPersonality=true
RestrictRealtime=true
RestrictNamespaces=true
CapabilityBoundingSet=
AmbientCapabilities=
SystemCallArchitectures=native
UMask=0077
ProcSubset=all
ProtectProc=invisible
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6 AF_NETLINK
[Install]
WantedBy=multi-user.target
EOF
}

create_config() {
  [[ -f "$CONFIG_DIR/config.toml" ]] && return
  install -o root -g vps-agent -m 0640 /dev/null "$CONFIG_DIR/config.toml"
  tee "$CONFIG_DIR/config.toml" >/dev/null <<'EOF'
[server]
listen = "127.0.0.1:9105"
[auth]
mode = "token"
token_file = "/etc/vps-agent/credentials/control-plane.token"
[limits]
max_concurrent_requests = 32
request_timeout_seconds = 5
max_body_bytes = 65536
read_requests_per_second = 10.0
write_requests_per_second = 2.0
burst = 20.0
command_timeout_seconds = 4
[services]
allowed = []
[firewall]
management_enabled = false
rollback_timeout_seconds = 60
max_snapshots = 5
[logging]
level = "info"
[telemetry]
connections_enabled = true
process_info_enabled = true
[logs]
enabled = false
allowed_units = ["vps-agent.service"]
max_entries_per_request = 200
max_message_bytes = 16384
[tls]
enabled = false
certificate_file = ""
private_key_file = ""
EOF
}

create_token() {
  local token_file="$CREDENTIAL_DIR/control-plane.token"
  [[ -f "$token_file" ]] && return
  if command -v openssl >/dev/null; then
    umask 077
    openssl rand -hex 32 > "$token_file"
  else
    umask 077
    od -An -N32 -tx1 /dev/urandom | tr -d ' \n' > "$token_file"
    printf '\n' >> "$token_file"
  fi
  chown vps-agent:vps-agent "$token_file"
  chmod 0600 "$token_file"
  info "created agent token at $token_file; copy it through a secure channel"
}

wait_healthy() {
  local attempts=15
  while (( attempts > 0 )); do
    if systemctl is-active --quiet vps-agent.service; then
      return 0
    fi
    sleep 1
    attempts=$((attempts - 1))
  done
  systemctl status --no-pager vps-agent.service || true
  return 1
}

do_install() {
  require_root; check_host
  [[ ! -e "$BINARY" ]] || die "$BINARY already exists; use update"
  local downloaded
  download_release
  downloaded="$DOWNLOADED_BINARY"
  getent group vps-agent >/dev/null || groupadd --system vps-agent
  id -u vps-agent >/dev/null 2>&1 || useradd --system --gid vps-agent --home-dir /nonexistent --shell /usr/sbin/nologin vps-agent
  install -d -o root -g vps-agent -m 0750 "$CONFIG_DIR" "$CREDENTIAL_DIR"
  install -d -o vps-agent -g vps-agent -m 0700 "$STATE_DIR"
  install -d -o root -g root -m 0755 "$(dirname "$PREVIOUS")"
  install -o root -g root -m 0755 "$downloaded" "$BINARY"
  create_config
  create_token
  install_unit
  systemctl daemon-reload
  systemctl enable --now vps-agent.service
  wait_healthy || die "agent failed to become healthy"
  systemctl status --no-pager vps-agent.service
}

do_update() {
  require_root; check_host
  [[ -x "$BINARY" ]] || die "agent is not installed"
  local downloaded staged
  download_release
  downloaded="$DOWNLOADED_BINARY"
  staged="$(dirname "$BINARY")/.vps-agent.new"
  install -o root -g root -m 0755 "$downloaded" "$staged"
  install -o root -g root -m 0755 "$BINARY" "$PREVIOUS"
  mv -f -- "$staged" "$BINARY"
  systemctl restart vps-agent.service
  if ! wait_healthy; then
    info "health check failed; restoring previous binary"
    install -o root -g root -m 0755 "$PREVIOUS" "$BINARY"
    systemctl restart vps-agent.service
    wait_healthy || die "rollback also failed; inspect systemctl status"
    die "update rolled back"
  fi
  info "updated to $($BINARY --version)"
}

do_rollback() {
  require_root; check_host
  [[ -x "$PREVIOUS" ]] || die "no previous binary is available"
  local current="$(dirname "$BINARY")/.vps-agent.current"
  install -o root -g root -m 0755 "$BINARY" "$current"
  install -o root -g root -m 0755 "$PREVIOUS" "$BINARY"
  mv -f -- "$current" "$PREVIOUS"
  systemctl restart vps-agent.service
  wait_healthy || die "rolled-back version did not become healthy"
}

do_remove() {
  require_root; check_host
  if [[ "$PURGE" == true && "$ASSUME_YES" != true ]]; then
    [[ -t 0 ]] || die "--purge requires an interactive terminal or --yes"
    read -r -p "Permanently delete $CONFIG_DIR and credentials? [y/N] " answer
    [[ "$answer" == "y" || "$answer" == "Y" ]] || die "purge cancelled"
  fi
  systemctl disable --now vps-agent.service 2>/dev/null || true
  rm -f -- "$UNIT" "$BINARY" "$PREVIOUS"
  rm -rf -- "$STATE_DIR"
  systemctl daemon-reload
  if [[ "$PURGE" == true ]]; then
    rm -rf -- "$CONFIG_DIR"
    info "removed configuration and credentials; this cannot be recovered"
  else
    info "preserved $CONFIG_DIR (use remove --purge to delete it)"
  fi
  userdel vps-agent 2>/dev/null || true
  groupdel vps-agent 2>/dev/null || true
}

[[ $# -ge 1 ]] || { usage; exit 2; }
ACTION="$1"; shift
while [[ $# -gt 0 ]]; do
  case "$1" in
    --version) [[ $# -ge 2 ]] || die "--version needs a value"; VERSION="$2"; shift 2 ;;
    --releases-url) [[ $# -ge 2 ]] || die "--releases-url needs a value"; RELEASES_URL="$2"; shift 2 ;;
    --purge) PURGE=true; shift ;;
    --yes) ASSUME_YES=true; shift ;;
    -h|--help) usage; exit 0 ;;
    *) die "unknown option: $1" ;;
  esac
done

case "$ACTION" in
  install) do_install ;;
  update) do_update ;;
  remove) do_remove ;;
  rollback) do_rollback ;;
  status) systemctl status --no-pager vps-agent.service ;;
  restart) require_root; systemctl restart vps-agent.service; wait_healthy ;;
  version) [[ -x "$BINARY" ]] && "$BINARY" --version || die "agent is not installed" ;;
  *) usage; exit 2 ;;
esac