#!/usr/bin/env bash

set -euo pipefail

# Proxmox cluster power helper
# - Start or shutdown a set of QEMU VMs and/or LXC containers by ID, or all.
# - Auth via API token or username/password (env vars or secret-tool).
#
# Requirements: curl, jq; optional: secret-tool (GNOME keyring)
#
# Environment variables (examples):
#   PVE_HOST=proxmox.example.com[:8006]
#   PVE_TOKEN_ID="user@pam!automation"     # when using API token
#   PVE_TOKEN_SECRET="xxxxxxxx-xxxx-xxxx"  # when using API token
#   PVE_USER="user"                        # when using password login
#   PVE_REALM="pam"                        # default pam
#   PVE_PASSWORD="..."                     # or provided via keyring
#   PVE_SCHEME="https"                    # default https
#   PVE_VERIFY_SSL="true|false"            # default true
#   PVE_NODE_FILTER=""                     # optional: restrict to node name
#
# Examples:
#   scripts/proxmox-power.sh --op shutdown --all
#   scripts/proxmox-power.sh --op start --ids 100 101 --only-qemu
#   PVE_TOKEN_ID=me@pam!ci PVE_TOKEN_SECRET=... scripts/proxmox-power.sh --op shutdown --all

SCHEME=${PVE_SCHEME:-https}
HOST=${PVE_HOST:-}
VERIFY_SSL=${PVE_VERIFY_SSL:-true}
INSECURE_FLAG=""

if [[ ${VERIFY_SSL} != "true" ]]; then
  INSECURE_FLAG="-k"
fi

usage() {
  cat <<EOF
Usage: $0 --op start|shutdown [--all | --ids <vmid> [<vmid> ...]] [options]

Options:
  --host HOST              Proxmox host (env PVE_HOST). Example: proxmox.example.com:8006
  --op OP                  Operation: start or shutdown
  --all                    Apply to all VMs/containers in the cluster (honors filters)
  --ids LIST               Space-separated list of VMIDs to operate on
  --only-qemu              Only operate on QEMU VMs
  --only-lxc               Only operate on LXC containers
  --include-stopped        Include stopped guests when op=shutdown (no-op otherwise)
  --force                  If shutdown times out, force stop
  --timeout SEC            Shutdown wait timeout (default 120)
  --concurrency N          Parallel operations (default 4)
  --node NODE              Restrict to a specific node name
  --dry-run                Show actions without executing
  --insecure               Do not verify SSL (same as PVE_VERIFY_SSL=false)
  -h, --help               Show this help

Auth (choose one):
  API Token: env PVE_TOKEN_ID and PVE_TOKEN_SECRET
  Password:  env PVE_USER, PVE_PASSWORD (or from keyring), optional PVE_REALM (default pam)

Keyring:
  If PVE_PASSWORD is empty and 'secret-tool' is available, the script tries:
    secret-tool lookup service proxmox user "+$PVE_USER+" realm "+${PVE_REALM:-pam}+"
  If PVE_TOKEN_SECRET is empty, it tries:
    secret-tool lookup service proxmox token_id "+$PVE_TOKEN_ID+"
EOF
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || { echo "Error: required command '$1' not found" >&2; exit 1; }
}

get_keyring() {
  local value=""
  if command -v secret-tool >/dev/null 2>&1; then
    value=$(secret-tool lookup "$@" || true)
  fi
  printf '%s' "$value"
}

# Globals set by auth_init
AUTH_HEADER=""
COOKIE_HEADER=""
CSRF_HEADER=""

auth_init() {
  local base_url="$SCHEME://$HOST/api2/json"

  if [[ -n "${PVE_TOKEN_ID:-}" && -z "${PVE_TOKEN_SECRET:-}" ]]; then
    PVE_TOKEN_SECRET=$(get_keyring service proxmox token_id "${PVE_TOKEN_ID}") || true
  fi

  if [[ -n "${PVE_TOKEN_ID:-}" && -n "${PVE_TOKEN_SECRET:-}" ]]; then
    AUTH_HEADER=("-H" "Authorization: PVEAPIToken=${PVE_TOKEN_ID}=${PVE_TOKEN_SECRET}")
    return 0
  fi

  local user="${PVE_USER:-}"
  local realm="${PVE_REALM:-pam}"
  local password="${PVE_PASSWORD:-}"

  if [[ -z "$user" ]]; then
    echo "Error: set PVE_TOKEN_ID/PVE_TOKEN_SECRET or PVE_USER[/PVE_PASSWORD]" >&2
    exit 2
  fi

  if [[ -z "$password" ]]; then
    password=$(get_keyring service proxmox user "$user" realm "$realm") || true
  fi

  if [[ -z "$password" ]]; then
    echo "Error: password not provided and not found in keyring for user '$user' realm '$realm'" >&2
    exit 2
  fi

  # Login to get ticket and CSRF token
  local resp
  resp=$(curl -sS $INSECURE_FLAG -X POST \("${AUTH_HEADER[*]}"\) \
    -d "username=${user}@${realm}" \
    -d "password=${password}" \
    "$base_url/access/ticket")

  local ticket csrf
  ticket=$(echo "$resp" | jq -r '.data.ticket // empty')
  csrf=$(echo "$resp" | jq -r '.data.CSRFPreventionToken // empty')
  if [[ -z "$ticket" || -z "$csrf" ]]; then
    echo "Error: failed to obtain auth ticket (check credentials)" >&2
    echo "$resp" | jq -r '.' >&2 || true
    exit 3
  fi
  COOKIE_HEADER=("-H" "Cookie: PVEAuthCookie=${ticket}")
  CSRF_HEADER=("-H" "CSRFPreventionToken: ${csrf}")
}

api_get() {
  local path="$1"; shift
  local url="$SCHEME://$HOST/api2/json$path"
  curl -sS $INSECURE_FLAG "${AUTH_HEADER[@]}" "${COOKIE_HEADER[@]}" -X GET "$url" "$@"
}

api_post() {
  local path="$1"; shift
  local url="$SCHEME://$HOST/api2/json$path"
  curl -sS $INSECURE_FLAG "${AUTH_HEADER[@]}" "${COOKIE_HEADER[@]}" "${CSRF_HEADER[@]}" -X POST "$url" "$@"
}

# Parse CLI
OP=""
DO_ALL=false
IDS=()
ONLY_QEMU=false
ONLY_LXC=false
INCLUDE_STOPPED=false
FORCE=false
TIMEOUT=120
CONCURRENCY=4
NODE_FILTER="${PVE_NODE_FILTER:-}"
DRY_RUN=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    --op) OP="$2"; shift 2;;
    --all) DO_ALL=true; shift;;
    --ids) shift; while [[ $# -gt 0 && ! "$1" =~ ^-- ]]; do IDS+=("$1"); shift; done ;;
    --only-qemu) ONLY_QEMU=true; shift;;
    --only-lxc) ONLY_LXC=true; shift;;
    --include-stopped) INCLUDE_STOPPED=true; shift;;
    --force) FORCE=true; shift;;
    --timeout) TIMEOUT="$2"; shift 2;;
    --concurrency) CONCURRENCY="$2"; shift 2;;
    --node) NODE_FILTER="$2"; shift 2;;
    --host) HOST="$2"; shift 2;;
    --dry-run) DRY_RUN=true; shift;;
    --insecure) VERIFY_SSL=false; INSECURE_FLAG="-k"; shift;;
    -h|--help) usage; exit 0;;
    *) echo "Unknown argument: $1" >&2; usage; exit 2;;
  esac
done

require_cmd curl
require_cmd jq

if [[ -z "$HOST" ]]; then
  echo "Error: --host or PVE_HOST is required" >&2
  usage
  exit 2
fi

case "$OP" in
  start|shutdown) :;;
  *) echo "Error: --op must be 'start' or 'shutdown'" >&2; usage; exit 2;;
esac

if ! $DO_ALL && [[ ${#IDS[@]} -eq 0 ]]; then
  echo "Error: specify --all or a list of --ids" >&2
  exit 2
fi

if $ONLY_QEMU && $ONLY_LXC; then
  echo "Error: cannot use --only-qemu and --only-lxc together" >&2
  exit 2
fi

auth_init

# Collect targets
resources=$(api_get "/cluster/resources?type=vm")

filter_jq='[.data[] | {type, vmid: (.vmid|tostring), status, node}]'
items=$(echo "$resources" | jq "$filter_jq")

if [[ -n "$NODE_FILTER" ]]; then
  items=$(echo "$items" | jq --arg node "$NODE_FILTER" '[.[] | select(.node==$node)]')
fi

if $ONLY_QEMU; then
  items=$(echo "$items" | jq '[.[] | select(.type=="qemu")]')
elif $ONLY_LXC; then
  items=$(echo "$items" | jq '[.[] | select(.type=="lxc")]')
fi

select_ids=()
if $DO_ALL; then
  mapfile -t select_ids < <(echo "$items" | jq -r '.[].vmid')
else
  select_ids=("${IDS[@]}")
fi

if [[ ${#select_ids[@]} -eq 0 ]]; then
  echo "No matching guests found." >&2
  exit 0
fi

# Build an associative map of vmid -> node,type,status
declare -A VM_NODE VM_TYPE VM_STATUS
while IFS=$'\t' read -r vid node type status; do
  VM_NODE[$vid]="$node"
  VM_TYPE[$vid]="$type"
  VM_STATUS[$vid]="$status"
done < <(
  echo "$items" | jq -r '.[] | "\(.vmid)\t\(.node)\t\(.type)\t\(.status)"'
)

work_list=()
for vid in "${select_ids[@]}"; do
  if [[ -z "${VM_NODE[$vid]:-}" ]]; then
    echo "Skip vmid=$vid (not found by filters)" >&2
    continue
  fi
  # Idempotence: skip if already desired state
  st="${VM_STATUS[$vid]}"
  case "$OP" in
    start)
      if [[ "$st" == "running" ]]; then
        echo "Already running: $vid (${VM_TYPE[$vid]} on ${VM_NODE[$vid]})"
        continue
      fi
      ;;
    shutdown)
      if [[ "$st" != "running" && $INCLUDE_STOPPED == false ]]; then
        echo "Already stopped: $vid (${VM_TYPE[$vid]} on ${VM_NODE[$vid]})"
        continue
      fi
      ;;
  esac
  work_list+=("$vid")
done

if [[ ${#work_list[@]} -eq 0 ]]; then
  echo "Nothing to do."
  exit 0
fi

run_action() {
  local vid="$1"
  local node="${VM_NODE[$vid]}"
  local type="${VM_TYPE[$vid]}"
  local path_base="/nodes/${node}/${type}/${vid}/status"

  echo "[$OP] ${type}:${vid} on node ${node}"
  if $DRY_RUN; then
    return 0
  fi

  case "$OP" in
    start)
      api_post "${path_base}/start" >/dev/null
      ;;
    shutdown)
      # Try graceful shutdown
      api_post "${path_base}/shutdown" -d "timeout=${TIMEOUT}" >/dev/null || true
      # Optionally force stop if still running after timeout
      # We poll once after timeout window to check status
      sleep 2
      local st_json
      st_json=$(api_get "/nodes/${node}/${type}/${vid}/status/current")
      local cur
      cur=$(echo "$st_json" | jq -r '.data.status // .data.status.current // empty')
      if [[ "$cur" == "running" && $FORCE == true ]]; then
        echo "Forcing stop: ${type}:${vid}"
        api_post "${path_base}/stop" >/dev/null || true
      fi
      ;;
  esac
}

# Parallelize with xargs -P
export -f run_action api_post api_get
export SCHEME HOST INSECURE_FLAG AUTH_HEADER COOKIE_HEADER CSRF_HEADER TIMEOUT FORCE DRY_RUN
declare -p VM_NODE VM_TYPE VM_STATUS >/dev/null 2>&1 || true

printf '%s\n' "${work_list[@]}" | xargs -I{} -P "$CONCURRENCY" bash -c 'run_action "$@"' _ {}

echo "Done: $OP ${#work_list[@]} item(s)."