#!/usr/bin/env bash
# proxmox-setup - Interactive Proxmox post-install configuration for Hetzner
#
# Configures a fresh Proxmox host with private networking, ZFS storage,
# WireGuard tunnel, firewall lockdown, and API tokens for incusos-proxmox.
#
# Runs from your workstation over SSH. Requires an SSH host alias configured
# in ~/.ssh/config (e.g. "hetzner-lab").
#
# Usage: proxmox-setup --host <ssh-alias> [OPTIONS]
# Run 'proxmox-setup --help' for full usage information.

set -euo pipefail

readonly VERSION="0.1.0"
readonly SCRIPT_NAME="proxmox-setup"
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

SSH_HOST=""
DRY_RUN=false
VERBOSE=false
QUIET=false
YES=false
SUBNET="10.10.0.0/24"
WG_SUBNET="10.10.99.0/24"
WG_PORT="51820"
SKIP_REPOS=false
SKIP_WIREGUARD=false
SKIP_FIREWALL=false
SKIP_STORAGE=false

# Derived (computed in detect step)
SUBNET_BASE=""          # e.g. 10.10.0
SUBNET_MASK=""          # e.g. 24
SUBNET_GW=""            # e.g. 10.10.0.1
WG_SUBNET_BASE=""       # e.g. 10.10.99
WG_SUBNET_MASK=""       # e.g. 24
WG_SERVER_IP=""         # e.g. 10.10.99.1
WG_CLIENT_IP=""         # e.g. 10.10.99.2
PUBLIC_IP=""
PVE_VERSION=""
HOSTNAME_DETECTED=""
KERNEL_VERSION=""

# ---------------------------------------------------------------------------
# Color and formatting
# ---------------------------------------------------------------------------

setup_colors() {
    if [[ -t 1 ]] && [[ "${NO_COLOR:-}" != "1" ]] && [[ "${TERM:-}" != "dumb" ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED='' GREEN='' YELLOW='' BLUE='' CYAN='' BOLD='' DIM='' RESET=''
    fi
}

info()    { [[ "$QUIET" == true ]] && return; echo -e "${BLUE}[info]${RESET}    $*"; }
success() { [[ "$QUIET" == true ]] && return; echo -e "${GREEN}[ok]${RESET}      $*"; }
warn()    { echo -e "${YELLOW}[warn]${RESET}    $*" >&2; }
error()   { echo -e "${RED}[error]${RESET}   $*" >&2; }
step()    { [[ "$QUIET" == true ]] && return; echo -e "${CYAN}[step]${RESET}    ${BOLD}$*${RESET}"; }
detail()  { [[ "$VERBOSE" != true ]] && return; echo -e "${DIM}          $*${RESET}"; }

dry_run_guard() {
    if [[ "$DRY_RUN" == true ]]; then
        info "(dry-run) $*"
        return 1
    fi
    return 0
}

# ---------------------------------------------------------------------------
# SSH helpers
# ---------------------------------------------------------------------------

# Run a command on the remote host via SSH
remote() {
    ssh -o ConnectTimeout=10 -o BatchMode=yes "$SSH_HOST" "$@"
}

# Run a command and capture output
remote_output() {
    ssh -o ConnectTimeout=10 -o BatchMode=yes "$SSH_HOST" "$@" 2>/dev/null
}

# Write content to a remote file via SSH
remote_write() {
    local dest="$1"
    local content="$2"
    ssh -o ConnectTimeout=10 -o BatchMode=yes "$SSH_HOST" "cat > '$dest'" <<< "$content"
}

# Check if remote file exists
remote_file_exists() {
    remote "test -f '$1'" 2>/dev/null
}

# ---------------------------------------------------------------------------
# Interactive helpers
# ---------------------------------------------------------------------------

confirm() {
    local prompt="$1"
    if [[ "$YES" == true ]]; then
        return 0
    fi
    echo -en "${BOLD}${prompt}${RESET} [y/N] "
    read -r answer
    [[ "$answer" =~ ^[Yy]$ ]]
}

prompt_value() {
    local prompt="$1"
    local default="$2"
    local result
    if [[ "$YES" == true ]]; then
        echo "$default"
        return
    fi
    echo -en "${BOLD}${prompt}${RESET} [${default}]: "
    read -r result
    echo "${result:-$default}"
}

# ---------------------------------------------------------------------------
# Help
# ---------------------------------------------------------------------------

usage() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} v${VERSION} - Interactive Proxmox post-install configuration for Hetzner

${BOLD}USAGE${RESET}
    ${SCRIPT_NAME} --host <ssh-alias> [OPTIONS]

${BOLD}DESCRIPTION${RESET}
    Configures a fresh Proxmox host over SSH with:
    - No-subscription repositories and system update
    - Private bridge (vmbr1) with NAT masquerading
    - ZFS storage pool on extra disks
    - WireGuard tunnel for secure remote access
    - nftables firewall lockdown
    - API token for incusos-proxmox automation

    Each step prompts for confirmation. Use --yes to skip prompts.

${BOLD}OPTIONS${RESET}
    ${BOLD}--host ALIAS${RESET}          SSH host alias (required, e.g. hetzner-lab)
    ${BOLD}--subnet CIDR${RESET}         Private bridge subnet (default: ${SUBNET})
    ${BOLD}--wg-subnet CIDR${RESET}      WireGuard subnet (default: ${WG_SUBNET})
    ${BOLD}--wg-port PORT${RESET}        WireGuard listen port (default: ${WG_PORT})
    ${BOLD}--skip-repos${RESET}          Skip repository changes
    ${BOLD}--skip-storage${RESET}        Skip storage pool creation
    ${BOLD}--skip-wireguard${RESET}      Skip WireGuard setup
    ${BOLD}--skip-firewall${RESET}       Skip firewall lockdown
    ${BOLD}-n, --dry-run${RESET}         Preview actions without executing
    ${BOLD}-y, --yes${RESET}             Skip confirmations (use defaults)
    ${BOLD}-v, --verbose${RESET}         Show detailed output
    ${BOLD}-q, --quiet${RESET}           Suppress informational output
    ${BOLD}-h, --help${RESET}            Show this help message
    ${BOLD}-V, --version${RESET}         Show version

${BOLD}EXAMPLES${RESET}
    # Preview everything
    ${SCRIPT_NAME} --host hetzner-lab --dry-run

    # Full setup with defaults
    ${SCRIPT_NAME} --host hetzner-lab --yes

    # Custom subnets, skip firewall
    ${SCRIPT_NAME} --host hetzner-lab --subnet 172.16.0.0/24 --skip-firewall

    # Just WireGuard and firewall (repos/storage already done)
    ${SCRIPT_NAME} --host hetzner-lab --skip-repos --skip-storage

${BOLD}PREREQUISITES${RESET}
    - SSH access to the Proxmox host (key-based auth recommended)
    - Add an entry to ~/.ssh/config:
        Host hetzner-lab
            HostName <public-ip>
            User root
            IdentityFile ~/.ssh/id_ed25519

${BOLD}SEE ALSO${RESET}
    hetzner/hetzner-setup.md    Full manual guide
    incusos/targets/hetzner/    Target config files
EOF
    exit 0
}

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --host)
                SSH_HOST="${2:?'--host requires an SSH alias'}"
                shift 2
                ;;
            --subnet)
                SUBNET="${2:?'--subnet requires a CIDR'}"
                shift 2
                ;;
            --wg-subnet)
                WG_SUBNET="${2:?'--wg-subnet requires a CIDR'}"
                shift 2
                ;;
            --wg-port)
                WG_PORT="${2:?'--wg-port requires a port number'}"
                shift 2
                ;;
            --skip-repos)
                SKIP_REPOS=true
                shift
                ;;
            --skip-storage)
                SKIP_STORAGE=true
                shift
                ;;
            --skip-wireguard)
                SKIP_WIREGUARD=true
                shift
                ;;
            --skip-firewall)
                SKIP_FIREWALL=true
                shift
                ;;
            -n|--dry-run)
                DRY_RUN=true
                shift
                ;;
            -y|--yes)
                YES=true
                shift
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -q|--quiet)
                QUIET=true
                shift
                ;;
            -h|--help)
                usage
                ;;
            -V|--version)
                echo "${SCRIPT_NAME} v${VERSION}"
                exit 0
                ;;
            -*)
                error "Unknown option: $1"
                echo "Run '${SCRIPT_NAME} --help' for usage information." >&2
                exit 1
                ;;
            *)
                error "Unexpected argument: $1"
                echo "Run '${SCRIPT_NAME} --help' for usage information." >&2
                exit 1
                ;;
        esac
    done

    if [[ -z "$SSH_HOST" ]]; then
        error "--host is required (SSH alias, e.g. hetzner-lab)"
        echo "Run '${SCRIPT_NAME} --help' for usage information." >&2
        exit 1
    fi
}

# ---------------------------------------------------------------------------
# Subnet helpers
# ---------------------------------------------------------------------------

parse_subnets() {
    # Parse private bridge subnet
    SUBNET_MASK="${SUBNET##*/}"
    local subnet_network="${SUBNET%%/*}"
    SUBNET_BASE="${subnet_network%.*}"
    SUBNET_GW="${SUBNET_BASE}.1"

    # Parse WireGuard subnet
    WG_SUBNET_MASK="${WG_SUBNET##*/}"
    local wg_network="${WG_SUBNET%%/*}"
    WG_SUBNET_BASE="${wg_network%.*}"
    WG_SERVER_IP="${WG_SUBNET_BASE}.1"
    WG_CLIENT_IP="${WG_SUBNET_BASE}.2"
}

# ---------------------------------------------------------------------------
# Step 1: Detect host information
# ---------------------------------------------------------------------------

step_detect() {
    step "Detecting host information"

    # Test SSH connectivity
    if ! remote "true" 2>/dev/null; then
        error "Cannot connect to '${SSH_HOST}' via SSH"
        error "Ensure the host is in ~/.ssh/config and key auth works:"
        echo "  ssh ${SSH_HOST} hostname" >&2
        exit 1
    fi
    success "SSH connection to ${SSH_HOST}"

    # Gather host info
    HOSTNAME_DETECTED=$(remote_output "hostname")
    KERNEL_VERSION=$(remote_output "uname -r")
    PVE_VERSION=$(remote_output "pvesh get /version --output-format json 2>/dev/null | python3 -c 'import sys,json; print(json.load(sys.stdin)[\"version\"])'" || echo "unknown")
    PUBLIC_IP=$(remote_output "hostname -I | awk '{print \$1}'" || echo "unknown")

    info "Hostname:     ${HOSTNAME_DETECTED}"
    info "Kernel:       ${KERNEL_VERSION}"
    info "PVE version:  ${PVE_VERSION}"
    info "Primary IP:   ${PUBLIC_IP}"

    # Block devices
    echo ""
    info "Block devices:"
    remote "lsblk -d -o NAME,SIZE,MODEL,ROTA,TYPE 2>/dev/null" | while IFS= read -r line; do
        echo "          $line"
    done

    # ZFS pools
    local zpool_output
    zpool_output=$(remote_output "zpool list -H 2>/dev/null" || true)
    if [[ -n "$zpool_output" ]]; then
        echo ""
        info "ZFS pools:"
        echo "$zpool_output" | while IFS= read -r line; do
            echo "          $line"
        done
    fi

    # Existing storage
    echo ""
    info "Proxmox storage:"
    remote "pvesm status 2>/dev/null" | while IFS= read -r line; do
        echo "          $line"
    done

    # Network interfaces
    echo ""
    info "Network bridges:"
    remote "ip -br addr show type bridge 2>/dev/null" | while IFS= read -r line; do
        echo "          $line"
    done

    # WireGuard status
    local wg_status
    wg_status=$(remote_output "wg show 2>/dev/null" || true)
    if [[ -n "$wg_status" ]]; then
        echo ""
        info "WireGuard active:"
        echo "$wg_status" | head -5 | while IFS= read -r line; do
            echo "          $line"
        done
    fi

    echo ""
}

# ---------------------------------------------------------------------------
# Step 2: Repositories
# ---------------------------------------------------------------------------

step_repos() {
    if [[ "$SKIP_REPOS" == true ]]; then
        info "Skipping repository configuration (--skip-repos)"
        return 0
    fi

    step "Repository configuration"

    # Check current state
    local has_enterprise
    has_enterprise=$(remote_output "grep -l '^deb.*enterprise' /etc/apt/sources.list.d/*.list 2>/dev/null" || true)
    local has_nosub
    has_nosub=$(remote_output "test -f /etc/apt/sources.list.d/pve-no-subscription.list && echo yes" || true)

    if [[ -z "$has_enterprise" ]] && [[ "$has_nosub" == "yes" ]]; then
        success "Already configured (no-subscription repos active)"
        return 0
    fi

    info "Will disable enterprise repos and enable no-subscription repos"

    if ! confirm "Apply repository changes?"; then
        info "Skipped"
        return 0
    fi

    if dry_run_guard "Disable enterprise repos, enable no-subscription"; then
        return 0
    fi

    # Disable enterprise repos
    remote "sed -i 's/^deb/# deb/' /etc/apt/sources.list.d/pve-enterprise.list 2>/dev/null || true"
    remote "test -f /etc/apt/sources.list.d/ceph.list && sed -i 's/^deb/# deb/' /etc/apt/sources.list.d/ceph.list || true"

    # Add no-subscription repo
    remote "echo 'deb http://download.proxmox.com/debian/pve bookworm pve-no-subscription' > /etc/apt/sources.list.d/pve-no-subscription.list"

    # Remove subscription nag
    remote "sed -Ezi.bak \"s/(Ext\\.Msg\\.show\\(\\{[^}]+title: gettext\\('No valid sub)/void\\(\\{ \\/\\/\\1/\" /usr/share/javascript/proxmox-widget-toolkit/proxmoxlib.js 2>/dev/null || true"
    remote "systemctl restart pveproxy 2>/dev/null || true"

    success "Repositories configured"
}

# ---------------------------------------------------------------------------
# Step 3: System update
# ---------------------------------------------------------------------------

step_update() {
    if [[ "$SKIP_REPOS" == true ]]; then
        info "Skipping system update (--skip-repos)"
        return 0
    fi

    step "System update"

    if ! confirm "Run apt update && apt dist-upgrade?"; then
        info "Skipped"
        return 0
    fi

    if dry_run_guard "apt update && apt dist-upgrade -y"; then
        return 0
    fi

    info "Updating package lists..."
    remote "apt update -qq"
    info "Upgrading packages (this may take a few minutes)..."
    remote "DEBIAN_FRONTEND=noninteractive apt dist-upgrade -y -qq"
    success "System updated"
}

# ---------------------------------------------------------------------------
# Step 4: Private bridge
# ---------------------------------------------------------------------------

step_bridge() {
    step "Private bridge (vmbr1)"

    # Check if vmbr1 already exists
    local existing_bridge
    existing_bridge=$(remote_output "ip addr show vmbr1 2>/dev/null" || true)
    if [[ -n "$existing_bridge" ]]; then
        success "vmbr1 already exists"
        remote "ip -br addr show vmbr1" | while IFS= read -r line; do
            echo "          $line"
        done
        return 0
    fi

    info "Will create vmbr1 with:"
    info "  Subnet:    ${SUBNET}"
    info "  Gateway:   ${SUBNET_GW}/${SUBNET_MASK}"
    info "  NAT:       MASQUERADE to vmbr0"

    local bridge_config
    bridge_config="
auto vmbr1
iface vmbr1 inet static
    address ${SUBNET_GW}/${SUBNET_MASK}
    bridge-ports none
    bridge-stp off
    bridge-fd 0
    post-up   echo 1 > /proc/sys/net/ipv4/ip_forward
    post-up   iptables -t nat -A POSTROUTING -s ${SUBNET} -o vmbr0 -j MASQUERADE
    post-down iptables -t nat -D POSTROUTING -s ${SUBNET} -o vmbr0 -j MASQUERADE"

    if [[ "$VERBOSE" == true ]]; then
        echo ""
        echo "$bridge_config"
        echo ""
    fi

    if ! confirm "Add vmbr1 to /etc/network/interfaces?"; then
        info "Skipped"
        return 0
    fi

    if dry_run_guard "Create vmbr1 (${SUBNET_GW}/${SUBNET_MASK})"; then
        return 0
    fi

    # Append to interfaces file
    remote "printf '%s\n' '' >> /etc/network/interfaces"
    remote_write "/tmp/vmbr1-config.tmp" "$bridge_config"
    remote "cat /tmp/vmbr1-config.tmp >> /etc/network/interfaces && rm -f /tmp/vmbr1-config.tmp"

    # Persist IP forwarding
    remote_write "/etc/sysctl.d/99-forward.conf" "net.ipv4.ip_forward = 1"
    remote "sysctl -p /etc/sysctl.d/99-forward.conf >/dev/null 2>&1"

    # Apply
    remote "ifreload -a 2>/dev/null || ifup vmbr1"

    success "vmbr1 created (${SUBNET_GW}/${SUBNET_MASK})"
}

# ---------------------------------------------------------------------------
# Step 5: Storage
# ---------------------------------------------------------------------------

step_storage() {
    if [[ "$SKIP_STORAGE" == true ]]; then
        info "Skipping storage setup (--skip-storage)"
        return 0
    fi

    step "ZFS storage pool"

    # Check if local-zfs already exists in PVE
    local existing_zfs
    existing_zfs=$(remote_output "pvesm status 2>/dev/null | grep -w local-zfs" || true)
    if [[ -n "$existing_zfs" ]]; then
        success "local-zfs already registered with Proxmox"
        echo "          $existing_zfs"
        return 0
    fi

    # Check if a zpool named local-zfs already exists
    local existing_pool
    existing_pool=$(remote_output "zpool list -H local-zfs 2>/dev/null" || true)
    if [[ -n "$existing_pool" ]]; then
        info "ZFS pool 'local-zfs' exists but not registered with Proxmox"
        if confirm "Register local-zfs with Proxmox?"; then
            if dry_run_guard "pvesm add zfspool local-zfs"; then
                return 0
            fi
            remote "pvesm add zfspool local-zfs -pool local-zfs"
            remote "pvesm set local-zfs -content images,rootdir"
            success "local-zfs registered"
        fi
        return 0
    fi

    # Find available disks (not in any zpool, not mounted)
    info "Scanning for available disks..."
    local available_disks
    available_disks=$(remote_output "
        used_disks=\$(zpool status 2>/dev/null | grep -oP '/dev/\\S+' | sed 's|/dev/||' || true)
        lsblk -dno NAME,SIZE,TYPE | while read name size type; do
            [[ \"\$type\" != \"disk\" ]] && continue
            # Skip if in a zpool
            echo \"\$used_disks\" | grep -qw \"\$name\" && continue
            # Skip if has mounted partitions
            lsblk -no MOUNTPOINT /dev/\$name 2>/dev/null | grep -q '/' && continue
            echo \"\$name \$size\"
        done
    " || true)

    if [[ -z "$available_disks" ]]; then
        info "No unused disks found -- skipping storage pool creation"
        info "If you have disks to use, create the pool manually:"
        info "  zpool create local-zfs mirror /dev/sdX /dev/sdY"
        info "  pvesm add zfspool local-zfs -pool local-zfs"
        return 0
    fi

    info "Available disks:"
    echo "$available_disks" | while IFS= read -r line; do
        echo "          /dev/$line"
    done

    local disk_count
    disk_count=$(echo "$available_disks" | wc -l)
    local disk_names
    disk_names=$(echo "$available_disks" | awk '{print "/dev/"$1}' | tr '\n' ' ')

    local pool_type
    if [[ "$disk_count" -eq 1 ]]; then
        pool_type="single"
        info "Recommendation: single disk (no redundancy)"
    elif [[ "$disk_count" -eq 2 ]]; then
        pool_type="mirror"
        info "Recommendation: mirror (2 disks, full redundancy)"
    else
        pool_type="raidz1"
        info "Recommendation: raidz1 (${disk_count} disks, 1-disk redundancy)"
    fi

    local pool_name
    pool_name=$(prompt_value "Pool name" "local-zfs")

    if ! confirm "Create ZFS pool '${pool_name}' (${pool_type}) with: ${disk_names}?"; then
        info "Skipped"
        return 0
    fi

    if dry_run_guard "zpool create ${pool_name} ${pool_type} ${disk_names}"; then
        return 0
    fi

    local zpool_cmd="zpool create ${pool_name}"
    case "$pool_type" in
        mirror) zpool_cmd="${zpool_cmd} mirror ${disk_names}" ;;
        raidz1) zpool_cmd="${zpool_cmd} raidz1 ${disk_names}" ;;
        single) zpool_cmd="${zpool_cmd} ${disk_names}" ;;
    esac

    remote "$zpool_cmd"
    remote "pvesm add zfspool ${pool_name} -pool ${pool_name}"
    remote "pvesm set ${pool_name} -content images,rootdir"

    success "ZFS pool '${pool_name}' created and registered"
}

# ---------------------------------------------------------------------------
# Step 6: WireGuard
# ---------------------------------------------------------------------------

step_wireguard() {
    if [[ "$SKIP_WIREGUARD" == true ]]; then
        info "Skipping WireGuard setup (--skip-wireguard)"
        return 0
    fi

    step "WireGuard tunnel"

    # Check if already configured
    local wg_active
    wg_active=$(remote_output "wg show wg0 2>/dev/null" || true)
    if [[ -n "$wg_active" ]]; then
        success "WireGuard wg0 already active"
        remote "wg show wg0" 2>/dev/null | head -8 | while IFS= read -r line; do
            echo "          $line"
        done
        return 0
    fi

    info "Will configure WireGuard:"
    info "  Server:    ${WG_SERVER_IP}/${WG_SUBNET_MASK} on port ${WG_PORT}"
    info "  Client:    ${WG_CLIENT_IP}/${WG_SUBNET_MASK}"
    info "  Routing:   wg0 <-> vmbr1 forwarding"

    if ! confirm "Set up WireGuard?"; then
        info "Skipped"
        return 0
    fi

    if dry_run_guard "Install and configure WireGuard (${WG_SERVER_IP}:${WG_PORT})"; then
        return 0
    fi

    # Install
    info "Installing wireguard..."
    remote "DEBIAN_FRONTEND=noninteractive apt install -y -qq wireguard >/dev/null 2>&1"

    # Generate server keys
    info "Generating server keypair..."
    remote "wg genkey | tee /etc/wireguard/server-private.key | wg pubkey > /etc/wireguard/server-public.key"
    remote "chmod 600 /etc/wireguard/server-private.key"

    local server_privkey
    server_privkey=$(remote_output "cat /etc/wireguard/server-private.key")
    local server_pubkey
    server_pubkey=$(remote_output "cat /etc/wireguard/server-public.key")

    # Generate client keys locally
    info "Generating client keypair..."
    local client_privkey client_pubkey
    client_privkey=$(wg genkey)
    client_pubkey=$(echo "$client_privkey" | wg pubkey)

    # Write server config
    local server_conf="[Interface]
PrivateKey = ${server_privkey}
Address = ${WG_SERVER_IP}/${WG_SUBNET_MASK}
ListenPort = ${WG_PORT}

# Allow forwarding between WireGuard clients and the private VM bridge
PostUp = iptables -A FORWARD -i wg0 -o vmbr1 -j ACCEPT; iptables -A FORWARD -i vmbr1 -o wg0 -j ACCEPT
PostDown = iptables -D FORWARD -i wg0 -o vmbr1 -j ACCEPT; iptables -D FORWARD -i vmbr1 -o wg0 -j ACCEPT

[Peer]
# Workstation
PublicKey = ${client_pubkey}
AllowedIPs = ${WG_CLIENT_IP}/32"

    remote_write "/etc/wireguard/wg0.conf" "$server_conf"
    remote "chmod 600 /etc/wireguard/wg0.conf"

    # Enable and start
    remote "systemctl enable --now wg-quick@wg0"

    success "WireGuard server configured"

    # Print client config
    echo ""
    info "Client config -- save to /etc/wireguard/hetzner-lab.conf or import into WireGuard app:"
    echo ""
    echo "  [Interface]"
    echo "  PrivateKey = ${client_privkey}"
    echo "  Address = ${WG_CLIENT_IP}/${WG_SUBNET_MASK}"
    echo ""
    echo "  [Peer]"
    echo "  PublicKey = ${server_pubkey}"
    echo "  Endpoint = ${PUBLIC_IP}:${WG_PORT}"
    echo "  AllowedIPs = 10.10.0.0/16"
    echo "  PersistentKeepalive = 25"
    echo ""
    warn "Save the client config now -- private key is not stored anywhere else"
}

# ---------------------------------------------------------------------------
# Step 7: Firewall
# ---------------------------------------------------------------------------

step_firewall() {
    if [[ "$SKIP_FIREWALL" == true ]]; then
        info "Skipping firewall setup (--skip-firewall)"
        return 0
    fi

    step "Firewall lockdown (nftables)"

    # Check if already configured
    local existing_rules
    existing_rules=$(remote_output "test -f /etc/nftables-hetzner.conf && echo yes" || true)
    if [[ "$existing_rules" == "yes" ]]; then
        success "Firewall rules already present (/etc/nftables-hetzner.conf)"
        return 0
    fi

    info "Will restrict public interface to SSH (22) and WireGuard (${WG_PORT}) only"
    info "Web UI (8006) and VM traffic only via WireGuard"
    warn "Make sure WireGuard is working before applying firewall rules!"

    if ! confirm "Apply firewall rules?"; then
        info "Skipped"
        return 0
    fi

    if dry_run_guard "Apply nftables rules (SSH + WG only on public interface)"; then
        return 0
    fi

    local nft_rules="#!/usr/sbin/nft -f

flush ruleset

table inet filter {
    chain input {
        type filter hook input priority 0; policy drop;

        # Loopback
        iif lo accept

        # Established/related connections
        ct state established,related accept

        # ICMP (ping)
        ip protocol icmp accept
        ip6 nexthdr icmpv6 accept

        # SSH on public interface
        iifname \"vmbr0\" tcp dport 22 accept

        # WireGuard on public interface
        iifname \"vmbr0\" udp dport ${WG_PORT} accept

        # Allow everything on private bridge and WireGuard
        iifname \"vmbr1\" accept
        iifname \"wg0\" accept

        # Log and drop everything else
        log prefix \"nft-drop: \" limit rate 5/minute counter drop
    }

    chain forward {
        type filter hook forward priority 0; policy drop;

        # Established/related
        ct state established,related accept

        # WireGuard <-> private bridge
        iifname \"wg0\" oifname \"vmbr1\" accept
        iifname \"vmbr1\" oifname \"wg0\" accept

        # Private bridge -> internet (NAT)
        iifname \"vmbr1\" oifname \"vmbr0\" accept
    }

    chain output {
        type filter hook output priority 0; policy accept;
    }
}"

    remote_write "/etc/nftables-hetzner.conf" "$nft_rules"
    remote "nft -f /etc/nftables-hetzner.conf"
    remote "cp /etc/nftables-hetzner.conf /etc/nftables.conf"
    remote "systemctl enable nftables"

    success "Firewall applied (SSH + WireGuard only on public interface)"
}

# ---------------------------------------------------------------------------
# Step 8: API token
# ---------------------------------------------------------------------------

step_api_token() {
    step "API token for incusos-proxmox"

    # Check if role already exists
    local existing_role
    existing_role=$(remote_output "pveum role list --output-format json 2>/dev/null | python3 -c 'import sys,json; [print(r[\"roleid\"]) for r in json.load(sys.stdin) if r[\"roleid\"]==\"IncusOSDeployer\"]'" || true)

    if [[ -n "$existing_role" ]]; then
        success "IncusOSDeployer role already exists"
    else
        if ! confirm "Create API role, pool, and token?"; then
            info "Skipped"
            return 0
        fi

        if dry_run_guard "Create IncusOSDeployer role + automation@pve!deploy token"; then
            return 0
        fi

        # Create role
        remote "pveum role add IncusOSDeployer -privs \
            'VM.Allocate VM.Config.CDROM VM.Config.CPU VM.Config.Disk VM.Config.HWType \
             VM.Config.Memory VM.Config.Network VM.Config.Options VM.PowerMgmt \
             VM.Monitor VM.Audit VM.Console \
             Datastore.AllocateSpace Datastore.Audit \
             Pool.Allocate Pool.Audit \
             SDN.Use Sys.Modify'"
        success "Created IncusOSDeployer role"
    fi

    # Create pool if needed
    local existing_pool
    existing_pool=$(remote_output "pveum pool list --output-format json 2>/dev/null | python3 -c 'import sys,json; [print(p[\"poolid\"]) for p in json.load(sys.stdin) if p[\"poolid\"]==\"IncusLab\"]'" || true)

    if [[ -z "$existing_pool" ]] && [[ "$DRY_RUN" != true ]]; then
        remote "pveum pool add IncusLab -comment 'IncusOS lab VMs'"
        success "Created IncusLab pool"
    fi

    # Create token if user doesn't exist yet
    local existing_token
    existing_token=$(remote_output "pveum user token list automation@pve --output-format json 2>/dev/null | python3 -c 'import sys,json; [print(t[\"tokenid\"]) for t in json.load(sys.stdin) if t[\"tokenid\"]==\"deploy\"]'" || true)

    if [[ -n "$existing_token" ]]; then
        success "automation@pve!deploy token already exists"
        info "If you need to regenerate the secret, delete and recreate the token"
        return 0
    fi

    if [[ "$DRY_RUN" == true ]]; then
        info "(dry-run) Would create automation@pve!deploy token"
        return 0
    fi

    # Create token
    local token_output
    token_output=$(remote "pveum user token add automation@pve deploy --privsep 0 --output-format json 2>/dev/null" || true)

    if [[ -z "$token_output" ]]; then
        # User might not exist yet
        remote "pveum user add automation@pve --comment 'IncusOS automation' 2>/dev/null || true"
        token_output=$(remote "pveum user token add automation@pve deploy --privsep 0 --output-format json")
    fi

    local token_secret
    token_secret=$(echo "$token_output" | python3 -c "import sys,json; print(json.load(sys.stdin)['value'])" 2>/dev/null || echo "$token_output")

    # Assign ACLs
    remote "pveum acl modify /pool/IncusLab -user automation@pve -role IncusOSDeployer"

    # Try to find the storage name for ACL assignment
    local storage_name
    storage_name=$(remote_output "pvesm status 2>/dev/null | awk '/zfspool/{print \$1}' | head -1" || true)
    if [[ -n "$storage_name" ]]; then
        remote "pveum acl modify /storage/${storage_name} -user automation@pve -role IncusOSDeployer"
    fi
    remote "pveum acl modify /storage/local -user automation@pve -role IncusOSDeployer"

    success "API token created"
    echo ""
    warn "Token secret (save now -- shown only once):"
    echo ""
    echo "  ${token_secret}"
    echo ""
}

# ---------------------------------------------------------------------------
# Step 9: Summary output
# ---------------------------------------------------------------------------

step_summary() {
    step "Summary"

    echo ""
    echo -e "${BOLD}proxmox.yaml for incusos/targets/hetzner/:${RESET}"
    echo ""
    echo "  host: ${SUBNET_GW}"
    echo "  method: api"
    echo "  api_token_id: automation@pve!deploy"
    echo "  node: ${HOSTNAME_DETECTED:-pve}"
    echo "  storage: local-zfs"
    echo "  iso_storage: local"
    echo "  bridge: vmbr1"
    echo "  gateway: ${SUBNET_GW}"
    echo "  dns: ${SUBNET_GW}"
    echo "  pool: IncusLab"
    echo "  ssh_user: root"

    echo ""
    echo -e "${BOLD}env file additions:${RESET}"
    echo ""
    echo "  # Hetzner"
    echo "  HETZNER_PROXMOX_TOKEN_SECRET=<token from step above>"
    echo "  HETZNER_PROXMOX_ROOT_PASSWORD=<your root password>"

    echo ""
    echo -e "${BOLD}Next steps:${RESET}"
    echo ""
    echo "  1. Save WireGuard client config and connect"
    echo "  2. Copy proxmox.yaml.example and fill in credentials:"
    echo "     cp incusos/targets/hetzner/proxmox.yaml.example incusos/targets/hetzner/proxmox.yaml"
    echo "  3. Test with dry run:"
    echo "     export PROXMOX_TOKEN_SECRET=\"\$HETZNER_PROXMOX_TOKEN_SECRET\""
    echo "     incusos-proxmox --dry-run --proxmox incusos/targets/hetzner/proxmox.yaml \\"
    echo "                     incusos/targets/hetzner/lab-cluster.yaml"
    echo ""
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

main() {
    setup_colors
    parse_args "$@"
    parse_subnets

    echo ""
    echo -e "${BOLD}${SCRIPT_NAME}${RESET} v${VERSION} -- Proxmox post-install configuration"
    if [[ "$DRY_RUN" == true ]]; then
        echo -e "${YELLOW}(dry-run mode -- no changes will be made)${RESET}"
    fi
    echo ""

    step_detect
    step_repos
    step_update
    step_bridge
    step_storage
    step_wireguard
    step_firewall
    step_api_token
    step_summary

    success "Setup complete"
}

main "$@"