#!/bin/bash

# DartNode Node Deployment Script
# Supports Proxmox VE 8.3+ (Linux) and macOS Apple Silicon hosts
# One-liner: curl -sSL https://pkg.dev.snaju.com/dn/dn-deploy.sh | bash
# With token: curl -sSL https://pkg.dev.snaju.com/dn/dn-deploy.sh | bash -s -- --token YOUR_TOKEN

# ============================================
# macOS Detection — delegate to dartnode.sh
# ============================================
if [ "$(uname -s)" = "Darwin" ]; then
    echo "[DartNode] macOS detected — delegating to macOS installer..."
    # Filter out "--" separator and pass remaining args to dartnode.sh
    MACOS_ARGS=()
    for arg in "$@"; do
        [ "$arg" != "--" ] && MACOS_ARGS+=("$arg")
    done
    exec bash <(curl -sSL https://pkg.dev.snaju.com/dn/dartnode.sh) install "${MACOS_ARGS[@]}"
    exit $?
fi

set -e

# Ensure sbin paths are in PATH (needed for mdadm, pvcreate, etc.)
export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:$PATH"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
MAGENTA='\033[0;35m'
NC='\033[0m' # No Color
BOLD='\033[1m'
DIM='\033[2m'

# Configuration
INTERFACES_FILE="/etc/network/interfaces"
MDADM_CONF="/etc/mdadm/mdadm.conf"
API_BASE="https://api.dartnode.net"

# Registration token (can be set via --token flag or DARTNODE_TOKEN env var)
REGISTRATION_TOKEN="${DARTNODE_TOKEN:-}"

# State tracking
NETWORK_CONFIGURED=false
RAID_CONFIGURED=false
LVM_CONFIGURED=false
STORAGE_CONFIGURED=false
DAEMON_INSTALLED=false
REGISTERED_WITH_API=false
RAID_DEVICE=""

# Restore mode: node lost its boot drive but customer data volumes survived.
# All storage steps become strictly non-destructive (assemble/activate/re-add
# only — never wipefs/mdadm --create/pvcreate/vgcreate).
RESTORE_MODE=false
DISCOVERED_VG=""

function header() {
    clear
    echo -e "${CYAN}"
    echo "    ___           _       __          _      "
    echo "   /   \\__ _ _ __| |_  /\\ \\ \\___   __| | ___ "
    echo "  / /\\ / _\` | '__| __|/  \\/ / _ \\ / _\` |/ _ \\"
    echo " / /_// (_| | |  | |_/ /\\  / (_) | (_| |  __/"
    echo "/___,' \\__,_|_|   \\__\\_\\ \\/ \\___/ \\__,_|\\___|"
    echo -e "${NC}"
    echo -e "${BOLD}DartNode Proxmox Node Deployment Script${NC}"
    echo "Copyright (c) 2024-2026, Snaju Inc. All Rights Reserved."
    echo ""
    echo "=============================================="
    echo ""
}

function log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

function log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

function log_warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

function log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

function log_step() {
    echo -e "${MAGENTA}[STEP]${NC} $1"
}

function check_root() {
    if [ "$EUID" -ne 0 ]; then
        log_error "Please run as root"
        exit 1
    fi
}

function check_proxmox() {
    if [ ! -f /etc/pve/datacenter.cfg ]; then
        log_error "This doesn't appear to be a Proxmox node."
        exit 1
    fi

    PVE_VERSION=$(pveversion | grep -oP 'pve-manager/\K[0-9]+\.[0-9]+')
    log_info "Detected Proxmox VE version: $PVE_VERSION"
}

function fix_proxmox_repos() {
    log_step "Configuring Proxmox repositories..."

    local REPOS_CHANGED=false

    # Disable enterprise PVE repo if it exists
    if [ -f /etc/apt/sources.list.d/pve-enterprise.list ]; then
        if ! grep -q "^#" /etc/apt/sources.list.d/pve-enterprise.list 2>/dev/null || grep -q "^deb" /etc/apt/sources.list.d/pve-enterprise.list 2>/dev/null; then
            log_info "Disabling Proxmox VE enterprise repository..."
            sed -i 's/^deb/#deb/' /etc/apt/sources.list.d/pve-enterprise.list
            REPOS_CHANGED=true
        fi
    fi

    # Disable enterprise Ceph repo if it exists
    if [ -f /etc/apt/sources.list.d/ceph.list ]; then
        if grep -q "enterprise.proxmox.com" /etc/apt/sources.list.d/ceph.list 2>/dev/null; then
            log_info "Disabling Proxmox Ceph enterprise repository..."
            sed -i 's/^deb/#deb/' /etc/apt/sources.list.d/ceph.list
            REPOS_CHANGED=true
        fi
    fi

    # Add no-subscription PVE repo if not present
    local PVE_NO_SUB="/etc/apt/sources.list.d/pve-no-subscription.list"
    if [ ! -f "$PVE_NO_SUB" ] || ! grep -q "pve-no-subscription" "$PVE_NO_SUB" 2>/dev/null; then
        log_info "Adding Proxmox VE no-subscription repository..."
        echo "deb http://download.proxmox.com/debian/pve bookworm pve-no-subscription" > "$PVE_NO_SUB"
        REPOS_CHANGED=true
    fi

    # Add no-subscription Ceph repo if not present (for Proxmox 8+)
    local CEPH_NO_SUB="/etc/apt/sources.list.d/ceph-no-subscription.list"
    if [ ! -f "$CEPH_NO_SUB" ] || ! grep -q "ceph-quincy" "$CEPH_NO_SUB" 2>/dev/null; then
        log_info "Adding Ceph no-subscription repository..."
        echo "deb http://download.proxmox.com/debian/ceph-quincy bookworm no-subscription" > "$CEPH_NO_SUB"
        REPOS_CHANGED=true
    fi

    if [ "$REPOS_CHANGED" = true ]; then
        log_info "Updating package lists..."
        apt-get update -qq
        log_success "Repositories configured!"
    else
        log_success "Repositories already configured."
    fi
}

function confirm() {
    local prompt="$1"
    local default="${2:-y}"

    if [[ "$default" == "y" ]]; then
        prompt="${CYAN}▶${NC} $prompt ${DIM}[Y/n]${NC}: "
    else
        prompt="${CYAN}▶${NC} $prompt ${DIM}[y/N]${NC}: "
    fi

    echo ""
    echo -ne "$prompt"
    read response < /dev/tty
    response=${response:-$default}

    [[ "$response" =~ ^[Yy]$ ]]
}

# ============================================
# JSON Parsing Helpers
# ============================================

function json_get() {
    local json="$1"
    local key="$2"

    if command -v python3 &> /dev/null; then
        echo "$json" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('data',d).get('$key',''))" 2>/dev/null
    else
        echo "$json" | grep -o "\"$key\":\"[^\"]*\"" | cut -d'"' -f4
    fi
}

function json_get_nested() {
    local json="$1"
    local parent="$2"
    local key="$3"

    if command -v python3 &> /dev/null; then
        echo "$json" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('data',d).get('$parent',{}).get('$key',''))" 2>/dev/null
    else
        echo "$json" | grep -o "\"$key\":[^,}]*" | head -1 | sed 's/.*://' | tr -d '"'
    fi
}

function json_get_bool() {
    local json="$1"
    local key="$2"

    if command -v python3 &> /dev/null; then
        echo "$json" | python3 -c "import sys,json; d=json.load(sys.stdin); print('true' if d.get('data',d).get('$key') else 'false')" 2>/dev/null
    else
        echo "$json" | grep -q "\"$key\":true" && echo "true" || echo "false"
    fi
}

# ============================================
# Token Validation Functions
# ============================================

function validate_token() {
    local token="$1"

    log_step "Validating registration token..."

    local response
    response=$(curl -sS "${API_BASE}/node/status?token=${token}" 2>/dev/null)

    if [ -z "$response" ]; then
        log_error "Unable to connect to DartNode API"
        return 1
    fi

    if echo "$response" | grep -q '"success":true'; then
        local is_valid is_used node_name
        is_valid=$(json_get_bool "$response" "isValid")
        is_used=$(json_get_bool "$response" "isUsed")
        node_name=$(json_get "$response" "nodeName")

        if [ "$is_valid" = "true" ]; then
            log_success "Token is valid for node: ${node_name}"
            TOKEN_NODE_NAME="$node_name"
            return 0
        elif [ "$is_used" = "true" ]; then
            log_warn "Token has already been used"
            return 2
        else
            log_error "Token is expired"
            return 1
        fi
    else
        local error_msg
        error_msg=$(json_get "$response" "error")
        log_error "Invalid token: ${error_msg:-Unknown error}"
        return 1
    fi
}

function fetch_token_config() {
    local token="$1"

    log_step "Fetching configuration from DartNode..."

    local response
    response=$(curl -sS "${API_BASE}/node/config?token=${token}" 2>/dev/null)

    if [ -z "$response" ]; then
        log_error "Unable to fetch configuration"
        return 1
    fi

    if echo "$response" | grep -q '"success":true'; then
        # Parse and store config values
        TOKEN_NODE_NAME=$(json_get "$response" "nodeName")
        TOKEN_NODE_DB_ID=$(json_get "$response" "nodeDbId")
        TOKEN_NODE_TYPE=$(json_get "$response" "nodeType")
        TOKEN_STORAGE_NAME=$(json_get "$response" "storageName")
        TOKEN_STORAGE_TYPE=$(json_get "$response" "storageType")
        TOKEN_MODE=$(json_get "$response" "mode")

        # NFS shares the fleet actually mounts, as "name|server|export|path|content|options"
        # lines. Empty when the API couldn't reach a reference node (or no python3).
        TOKEN_NFS_STORAGES=""
        if command -v python3 &> /dev/null; then
            TOKEN_NFS_STORAGES=$(echo "$response" | python3 -c "
import sys, json
d = json.load(sys.stdin)
d = d.get('data', d)
for s in (d.get('nfsStorages') or []):
    print('|'.join([str(s.get(k) or '') for k in ('storage', 'server', 'export', 'path', 'content', 'options')]))
" 2>/dev/null)
        fi
        TOKEN_REDIS_HOST=$(json_get_nested "$response" "redis" "host")
        TOKEN_REDIS_PORT=$(json_get_nested "$response" "redis" "port")
        TOKEN_REDIS_PASSWORD=$(json_get_nested "$response" "redis" "password")
        TOKEN_REDIS_DB=$(json_get_nested "$response" "redis" "db")
        TOKEN_AUTH_SECRET=$(json_get "$response" "authSecret")

        # JSON null comes back as "None" (python3) or "null" — treat as unset
        case "$TOKEN_STORAGE_NAME" in null|None) TOKEN_STORAGE_NAME="" ;; esac
        case "$TOKEN_STORAGE_TYPE" in null|None) TOKEN_STORAGE_TYPE="" ;; esac
        case "$TOKEN_NODE_DB_ID" in null|None) TOKEN_NODE_DB_ID="" ;; esac

        # Set defaults for required fields
        TOKEN_REDIS_HOST=${TOKEN_REDIS_HOST:-23.165.104.28}
        TOKEN_REDIS_PORT=${TOKEN_REDIS_PORT:-6379}
        TOKEN_REDIS_DB=${TOKEN_REDIS_DB:-0}
        TOKEN_REDIS_PASSWORD=${TOKEN_REDIS_PASSWORD:-}
        TOKEN_MODE=${TOKEN_MODE:-pve}
        TOKEN_AUTH_SECRET=${TOKEN_AUTH_SECRET:-dartnode-vnc-secret-2024}

        # Debug: Show what was parsed
        if [ -n "$DEBUG" ]; then
            echo "[DEBUG] Parsed config:"
            echo "  NODE_NAME: $TOKEN_NODE_NAME"
            echo "  REDIS_HOST: $TOKEN_REDIS_HOST"
            echo "  REDIS_PORT: $TOKEN_REDIS_PORT"
            echo "  REDIS_PASSWORD: ${TOKEN_REDIS_PASSWORD:0:8}..."
        fi

        log_success "Configuration fetched for node: ${TOKEN_NODE_NAME}"
        return 0
    else
        local error_msg
        error_msg=$(json_get "$response" "error")
        log_error "Failed to fetch configuration: ${error_msg:-Unknown error}"
        return 1
    fi
}

# ============================================
# Network Configuration
# ============================================

function list_network_interfaces() {
    echo ""
    log_info "Available network interfaces:"
    echo ""

    ip -o link show | awk -F': ' '{print $2}' | grep -vE '^(lo|veth|tap|fwbr|fwln|fwpr)' | while read iface; do
        STATE=$(cat /sys/class/net/$iface/operstate 2>/dev/null || echo "unknown")
        SPEED=$(cat /sys/class/net/$iface/speed 2>/dev/null || echo "?")
        MAC=$(cat /sys/class/net/$iface/address 2>/dev/null || echo "unknown")

        if [ "$STATE" == "up" ]; then
            STATE_COLOR="${GREEN}UP${NC}"
        else
            STATE_COLOR="${RED}DOWN${NC}"
        fi

        printf "  %-20s  State: %-12b  Speed: %s Mb/s  MAC: %s\n" "$iface" "$STATE_COLOR" "$SPEED" "$MAC"
    done
    echo ""
}

function prompt_choice() {
    local prompt="$1"
    local default="$2"

    echo "" > /dev/tty
    echo -ne "${CYAN}▶${NC} $prompt ${DIM}[$default]${NC}: " > /dev/tty
    read response < /dev/tty
    echo "${response:-$default}"
}

function setup_network_bridges() {
    echo ""
    echo -e "${BOLD}Network Configuration Mode:${NC}"
    echo ""
    echo -e "  ${CYAN}1)${NC} Standard - Separate physical ports for WAN and LAN"
    echo -e "  ${CYAN}2)${NC} VLAN - Single trunk port with VLAN tagging"
    echo -e "  ${CYAN}3)${NC} Skip network configuration"

    NET_MODE=$(prompt_choice "Select mode" "1")

    case $NET_MODE in
        1)
            setup_network_standard
            ;;
        2)
            setup_network_vlan
            ;;
        3)
            log_info "Network configuration skipped."
            return 0
            ;;
        *)
            log_error "Invalid selection."
            return 1
            ;;
    esac
}

function setup_network_standard() {
    echo ""
    log_info "Standard network configuration (separate ports)"
    log_warn "Common SuperMicro interfaces: enp1s0f0 (WAN), enp1s0f1 (LAN)"

    list_network_interfaces

    # Check existing configuration
    EXISTING_VMBR0=$(grep -A5 "^iface vmbr0" $INTERFACES_FILE 2>/dev/null | grep bridge-ports | awk '{print $2}' || echo "")
    EXISTING_VMBR0_IP=$(grep -A5 "^iface vmbr0" $INTERFACES_FILE 2>/dev/null | grep "address" | awk '{print $2}' || echo "")
    EXISTING_VMBR1=$(grep -A5 "^iface vmbr1" $INTERFACES_FILE 2>/dev/null | grep bridge-ports | awk '{print $2}' || echo "")
    EXISTING_VMBR1_IP=$(grep -A5 "^iface vmbr1" $INTERFACES_FILE 2>/dev/null | grep "address" | awk '{print $2}' || echo "")

    SKIP_WAN=false
    SKIP_LAN=false

    # ===== WAN (vmbr0) Configuration =====
    echo ""
    echo -e "${BOLD}━━━ WAN Bridge (vmbr0) ━━━${NC}"

    if [ -n "$EXISTING_VMBR0" ]; then
        echo ""
        log_success "vmbr0 already configured:"
        echo -e "  ${DIM}Interface:${NC} $EXISTING_VMBR0"
        echo -e "  ${DIM}IP:${NC} ${EXISTING_VMBR0_IP:-DHCP}"

        if ! confirm "Reconfigure vmbr0 (WAN)?"; then
            SKIP_WAN=true
            WAN_IFACE="$EXISTING_VMBR0"
            log_info "Keeping existing vmbr0 configuration."
        fi
    fi

    if [ "$SKIP_WAN" = false ]; then
        WAN_IFACE=$(prompt_choice "Enter WAN interface for vmbr0" "${EXISTING_VMBR0:-enp1s0f0}")

        if [ ! -d "/sys/class/net/$WAN_IFACE" ]; then
            log_error "Interface $WAN_IFACE does not exist!"
            if ! confirm "Continue anyway?" "n"; then
                return 1
            fi
        fi

        echo ""
        echo -e "${BOLD}WAN IP Configuration:${NC}"
        echo -e "  ${CYAN}1)${NC} Static IP"
        echo -e "  ${CYAN}2)${NC} DHCP"
        WAN_IP_MODE=$(prompt_choice "Select" "1")

        if [ "$WAN_IP_MODE" == "1" ]; then
            WAN_IP=$(prompt_choice "Enter WAN IP address (e.g., 38.134.40.96/24)" "$EXISTING_VMBR0_IP")
            WAN_GW=$(prompt_choice "Enter WAN gateway (e.g., 38.134.40.1)" "")
        fi
    fi

    # ===== LAN (vmbr1) Configuration =====
    echo ""
    echo -e "${BOLD}━━━ LAN Bridge (vmbr1) ━━━${NC}"

    if [ -n "$EXISTING_VMBR1" ]; then
        echo ""
        log_success "vmbr1 already configured:"
        echo -e "  ${DIM}Interface:${NC} $EXISTING_VMBR1"
        echo -e "  ${DIM}IP:${NC} ${EXISTING_VMBR1_IP:-DHCP}"

        if ! confirm "Reconfigure vmbr1 (LAN)?"; then
            SKIP_LAN=true
            LAN_IFACE="$EXISTING_VMBR1"
            log_info "Keeping existing vmbr1 configuration."
        fi
    fi

    if [ "$SKIP_LAN" = false ]; then
        LAN_IFACE=$(prompt_choice "Enter LAN interface for vmbr1" "${EXISTING_VMBR1:-enp1s0f1}")

        if [ ! -d "/sys/class/net/$LAN_IFACE" ]; then
            log_error "Interface $LAN_IFACE does not exist!"
            if ! confirm "Continue anyway?" "n"; then
                return 1
            fi
        fi

        echo ""
        echo -e "${BOLD}LAN IP Configuration:${NC}"
        echo -e "  ${CYAN}1)${NC} Static IP"
        echo -e "  ${CYAN}2)${NC} DHCP"
        LAN_IP_MODE=$(prompt_choice "Select" "1")

        if [ "$LAN_IP_MODE" == "1" ]; then
            LAN_IP=$(prompt_choice "Enter LAN IP address (e.g., 10.1.1.7/16)" "$EXISTING_VMBR1_IP")
        fi

        # Ask about static route for internal network
        echo ""
        if confirm "Add static route for 10.0.0.0/8 network (internal/LAN traffic)?"; then
            ADD_LAN_ROUTE="y"
            # Try to guess the gateway from the LAN IP
            if [ -n "$LAN_IP" ]; then
                DEFAULT_LAN_GW=$(echo "$LAN_IP" | sed 's/\.[0-9]*\/.*/.1/')
            else
                DEFAULT_LAN_GW="10.1.0.1"
            fi
            LAN_ROUTE_GW=$(prompt_choice "Enter LAN gateway for 10.0.0.0/8 route" "$DEFAULT_LAN_GW")
        else
            ADD_LAN_ROUTE="n"
        fi
    else
        ADD_LAN_ROUTE="n"
    fi

    # ===== Summary and Apply =====
    if [ "$SKIP_WAN" = true ] && [ "$SKIP_LAN" = true ]; then
        log_info "Both bridges already configured. No changes needed."
        NETWORK_CONFIGURED=true
        return 0
    fi

    # If skipping WAN, load existing WAN config
    if [ "$SKIP_WAN" = true ]; then
        WAN_IFACE="$EXISTING_VMBR0"
        if [ -n "$EXISTING_VMBR0_IP" ]; then
            WAN_IP_MODE="1"
            WAN_IP="$EXISTING_VMBR0_IP"
            WAN_GW=$(grep -A10 "^iface vmbr0" $INTERFACES_FILE 2>/dev/null | grep "gateway" | awk '{print $2}' || echo "")
        else
            WAN_IP_MODE="2"
        fi
    fi

    # If skipping LAN, load existing LAN config
    if [ "$SKIP_LAN" = true ]; then
        LAN_IFACE="$EXISTING_VMBR1"
        if [ -n "$EXISTING_VMBR1_IP" ]; then
            LAN_IP_MODE="1"
            LAN_IP="$EXISTING_VMBR1_IP"
        else
            LAN_IP_MODE="2"
        fi
        # Check for existing route
        if grep -q "10.0.0.0/8" $INTERFACES_FILE 2>/dev/null; then
            ADD_LAN_ROUTE="y"
            LAN_ROUTE_GW=$(grep -A20 "^iface vmbr1" $INTERFACES_FILE 2>/dev/null | grep "10.0.0.0/8" | grep -oE '[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "")
        fi
    fi

    echo ""
    log_info "Configuration summary:"
    echo "  vmbr0 (WAN): $WAN_IFACE"
    if [ "$WAN_IP_MODE" == "1" ]; then
        echo "    IP: $WAN_IP, Gateway: $WAN_GW"
    else
        echo "    IP: DHCP"
    fi
    if [ "$SKIP_WAN" = true ]; then
        echo -e "    ${DIM}(keeping existing)${NC}"
    fi
    echo "  vmbr1 (LAN): $LAN_IFACE"
    if [ "$LAN_IP_MODE" == "1" ]; then
        echo "    IP: $LAN_IP"
    else
        echo "    IP: DHCP"
    fi
    if [[ "$ADD_LAN_ROUTE" =~ ^[Yy]$ ]]; then
        echo "    Route: 10.0.0.0/8 via $LAN_ROUTE_GW"
    fi
    if [ "$SKIP_LAN" = true ]; then
        echo -e "    ${DIM}(keeping existing)${NC}"
    fi
    echo ""

    if ! confirm "Apply this network configuration?"; then
        log_warn "Network configuration skipped."
        return 0
    fi

    cp $INTERFACES_FILE "${INTERFACES_FILE}.bak.$(date +%Y%m%d%H%M%S)"
    log_info "Backed up current network configuration."

    generate_standard_interfaces "$WAN_IFACE" "$WAN_IP_MODE" "$WAN_IP" "$WAN_GW" "$LAN_IFACE" "$LAN_IP_MODE" "$LAN_IP" "$ADD_LAN_ROUTE" "$LAN_ROUTE_GW"

    log_success "Network bridges configured!"
    log_warn "Network changes will take effect after reboot or 'ifreload -a'"

    # Apply network changes now
    if confirm "Apply network changes now (ifreload -a)?"; then
        log_step "Applying network configuration..."
        ifreload -a 2>/dev/null || ifdown -a && ifup -a 2>/dev/null || true
        sleep 2

        # Run connectivity checks
        check_network_connectivity
    fi

    NETWORK_CONFIGURED=true
}

function check_network_connectivity() {
    echo ""
    log_step "Checking network connectivity..."
    echo ""

    # Check WAN connectivity (1.1.1.1)
    echo -n "  WAN (1.1.1.1): "
    if ping -c 2 -W 3 1.1.1.1 &>/dev/null; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${RED}FAILED${NC}"
        log_warn "WAN connectivity check failed. Check vmbr0 configuration."
    fi

    # Check LAN connectivity (10.1.0.5)
    echo -n "  LAN (10.1.0.5): "
    if ping -c 2 -W 3 10.1.0.5 &>/dev/null; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${RED}FAILED${NC}"
        log_warn "LAN connectivity check failed. Check vmbr1 configuration and 10.0.0.0/8 route."
    fi

    echo ""
}

function generate_standard_interfaces() {
    local WAN_IFACE="$1"
    local WAN_IP_MODE="$2"
    local WAN_IP="$3"
    local WAN_GW="$4"
    local LAN_IFACE="$5"
    local LAN_IP_MODE="$6"
    local LAN_IP="$7"
    local ADD_ROUTE="$8"
    local ROUTE_GW="$9"

    cat > $INTERFACES_FILE <<EOF
auto lo
iface lo inet loopback

# -----------------------------------------------------------------------
# PHYSICAL INTERFACES
# -----------------------------------------------------------------------

iface $WAN_IFACE inet manual

iface $LAN_IFACE inet manual

# -----------------------------------------------------------------------
# BRIDGES
# -----------------------------------------------------------------------

# WAN Bridge (vmbr0)
auto vmbr0
EOF

    if [ "$WAN_IP_MODE" == "1" ]; then
        cat >> $INTERFACES_FILE <<EOF
iface vmbr0 inet static
        address $WAN_IP
        gateway $WAN_GW
        bridge-ports $WAN_IFACE
        bridge-stp off
        bridge-fd 0
EOF
    else
        cat >> $INTERFACES_FILE <<EOF
iface vmbr0 inet dhcp
        bridge-ports $WAN_IFACE
        bridge-stp off
        bridge-fd 0
EOF
    fi

    cat >> $INTERFACES_FILE <<EOF

# LAN Bridge (vmbr1)
auto vmbr1
EOF

    if [ "$LAN_IP_MODE" == "1" ]; then
        cat >> $INTERFACES_FILE <<EOF
iface vmbr1 inet static
        address $LAN_IP
        bridge-ports $LAN_IFACE
        bridge-stp off
        bridge-fd 0
EOF
    else
        cat >> $INTERFACES_FILE <<EOF
iface vmbr1 inet dhcp
        bridge-ports $LAN_IFACE
        bridge-stp off
        bridge-fd 0
EOF
    fi

    # Add static route for internal network if requested
    if [[ "$ADD_ROUTE" =~ ^[Yy]$ ]] && [ -n "$ROUTE_GW" ]; then
        cat >> $INTERFACES_FILE <<EOF
        up ip route add 10.0.0.0/8 via $ROUTE_GW || true
EOF
    fi

    cat >> $INTERFACES_FILE <<EOF

source /etc/network/interfaces.d/*
EOF
}

function setup_network_vlan() {
    echo ""
    log_info "VLAN network configuration (single trunk port)"
    log_warn "This configures VLAN sub-interfaces on a single physical port"

    list_network_interfaces

    # Check existing configuration
    EXISTING_VMBR0=$(grep -A5 "^iface vmbr0" $INTERFACES_FILE 2>/dev/null | grep bridge-ports | awk '{print $2}' || echo "")
    EXISTING_VMBR0_IP=$(grep -A5 "^iface vmbr0" $INTERFACES_FILE 2>/dev/null | grep "address" | awk '{print $2}' || echo "")
    EXISTING_VMBR1=$(grep -A5 "^iface vmbr1" $INTERFACES_FILE 2>/dev/null | grep bridge-ports | awk '{print $2}' || echo "")
    EXISTING_VMBR1_IP=$(grep -A5 "^iface vmbr1" $INTERFACES_FILE 2>/dev/null | grep "address" | awk '{print $2}' || echo "")

    # Try to detect trunk interface from existing VLAN config (e.g., "enp1s0f0.20" -> "enp1s0f0")
    EXISTING_TRUNK=""
    if [ -n "$EXISTING_VMBR0" ] && [[ "$EXISTING_VMBR0" == *.* ]]; then
        EXISTING_TRUNK="${EXISTING_VMBR0%.*}"
    fi

    SKIP_WAN=false
    SKIP_LAN=false

    TRUNK_IFACE=$(prompt_choice "Enter trunk interface (physical port)" "${EXISTING_TRUNK:-}")

    if [ -z "$TRUNK_IFACE" ]; then
        log_error "Trunk interface is required!"
        return 1
    fi

    if [ ! -d "/sys/class/net/$TRUNK_IFACE" ]; then
        log_error "Interface $TRUNK_IFACE does not exist!"
        if ! confirm "Continue anyway?" "n"; then
            return 1
        fi
    fi

    echo ""
    log_info "VLAN Configuration"
    echo -e "  ${DIM}Common VLANs:${NC}"
    echo -e "  ${DIM}  VLAN 10 = LAN/Private network${NC}"
    echo -e "  ${DIM}  VLAN 20 = WAN/Public network${NC}"

    # ===== WAN (vmbr0) Configuration =====
    echo ""
    echo -e "${BOLD}━━━ WAN Bridge (vmbr0) ━━━${NC}"

    if [ -n "$EXISTING_VMBR0" ]; then
        echo ""
        log_success "vmbr0 already configured:"
        echo -e "  ${DIM}Bridge Port:${NC} $EXISTING_VMBR0"
        echo -e "  ${DIM}IP:${NC} ${EXISTING_VMBR0_IP:-DHCP}"

        if ! confirm "Reconfigure vmbr0 (WAN)?"; then
            SKIP_WAN=true
            # Extract VLAN from existing config
            if [[ "$EXISTING_VMBR0" == *.* ]]; then
                WAN_VLAN="${EXISTING_VMBR0##*.}"
            fi
            log_info "Keeping existing vmbr0 configuration."
        fi
    fi

    if [ "$SKIP_WAN" = false ]; then
        WAN_VLAN=$(prompt_choice "Enter WAN VLAN ID" "20")

        echo ""
        echo -e "${BOLD}WAN (VLAN $WAN_VLAN) IP Configuration:${NC}"
        echo -e "  ${CYAN}1)${NC} Static IP"
        echo -e "  ${CYAN}2)${NC} DHCP"
        WAN_IP_MODE=$(prompt_choice "Select" "1")

        if [ "$WAN_IP_MODE" == "1" ]; then
            WAN_IP=$(prompt_choice "Enter WAN IP address (e.g., 108.165.121.14/24)" "$EXISTING_VMBR0_IP")
            WAN_GW=$(prompt_choice "Enter WAN gateway (e.g., 108.165.121.1)" "")
        fi
    fi

    # ===== LAN (vmbr1) Configuration =====
    echo ""
    echo -e "${BOLD}━━━ LAN Bridge (vmbr1) ━━━${NC}"

    if [ -n "$EXISTING_VMBR1" ]; then
        echo ""
        log_success "vmbr1 already configured:"
        echo -e "  ${DIM}Bridge Port:${NC} $EXISTING_VMBR1"
        echo -e "  ${DIM}IP:${NC} ${EXISTING_VMBR1_IP:-DHCP}"

        if ! confirm "Reconfigure vmbr1 (LAN)?"; then
            SKIP_LAN=true
            # Extract VLAN from existing config
            if [[ "$EXISTING_VMBR1" == *.* ]]; then
                LAN_VLAN="${EXISTING_VMBR1##*.}"
            fi
            log_info "Keeping existing vmbr1 configuration."
        fi
    fi

    if [ "$SKIP_LAN" = false ]; then
        LAN_VLAN=$(prompt_choice "Enter LAN VLAN ID" "10")

        echo ""
        echo -e "${BOLD}LAN (VLAN $LAN_VLAN) IP Configuration:${NC}"
        echo -e "  ${CYAN}1)${NC} Static IP"
        echo -e "  ${CYAN}2)${NC} DHCP"
        LAN_IP_MODE=$(prompt_choice "Select" "2")

        if [ "$LAN_IP_MODE" == "1" ]; then
            LAN_IP=$(prompt_choice "Enter LAN IP address (e.g., 10.4.0.10/24)" "$EXISTING_VMBR1_IP")
        fi

        if confirm "Add route for 10.0.0.0/8 network?"; then
            ADD_LAN_ROUTE="y"
            LAN_ROUTE_GW=$(prompt_choice "Enter LAN gateway for 10.0.0.0/8 route" "10.4.0.1")
        else
            ADD_LAN_ROUTE="n"
        fi
    else
        ADD_LAN_ROUTE="n"
    fi

    # ===== Summary and Apply =====
    if [ "$SKIP_WAN" = true ] && [ "$SKIP_LAN" = true ]; then
        log_info "Both bridges already configured. No changes needed."
        NETWORK_CONFIGURED=true
        return 0
    fi

    # If skipping WAN, load existing WAN config
    if [ "$SKIP_WAN" = true ]; then
        if [ -n "$EXISTING_VMBR0_IP" ]; then
            WAN_IP_MODE="1"
            WAN_IP="$EXISTING_VMBR0_IP"
            WAN_GW=$(grep -A10 "^iface vmbr0" $INTERFACES_FILE 2>/dev/null | grep "gateway" | awk '{print $2}' || echo "")
        else
            WAN_IP_MODE="2"
        fi
    fi

    # If skipping LAN, load existing LAN config
    if [ "$SKIP_LAN" = true ]; then
        if [ -n "$EXISTING_VMBR1_IP" ]; then
            LAN_IP_MODE="1"
            LAN_IP="$EXISTING_VMBR1_IP"
        else
            LAN_IP_MODE="2"
        fi
        # Check for existing route
        if grep -q "10.0.0.0/8" $INTERFACES_FILE 2>/dev/null; then
            ADD_LAN_ROUTE="y"
            LAN_ROUTE_GW=$(grep -A15 "^iface vmbr1" $INTERFACES_FILE 2>/dev/null | grep "10.0.0.0/8" | grep -oP 'via \K[\d.]+' || echo "10.4.0.1")
        fi
    fi

    echo ""
    log_info "Configuration summary:"
    echo "  Trunk Interface: $TRUNK_IFACE"
    echo ""
    echo "  vmbr0 (WAN - VLAN $WAN_VLAN): ${TRUNK_IFACE}.${WAN_VLAN}"
    if [ "$WAN_IP_MODE" == "1" ]; then
        echo "    IP: $WAN_IP, Gateway: $WAN_GW"
    else
        echo "    IP: DHCP"
    fi
    if [ "$SKIP_WAN" = true ]; then
        echo -e "    ${DIM}(keeping existing)${NC}"
    fi
    echo ""
    echo "  vmbr1 (LAN - VLAN $LAN_VLAN): ${TRUNK_IFACE}.${LAN_VLAN}"
    if [ "$LAN_IP_MODE" == "1" ]; then
        echo "    IP: $LAN_IP"
    else
        echo "    IP: DHCP"
    fi
    if [[ "$ADD_LAN_ROUTE" =~ ^[Yy]$ ]]; then
        echo "    Route: 10.0.0.0/8 via $LAN_ROUTE_GW"
    fi
    if [ "$SKIP_LAN" = true ]; then
        echo -e "    ${DIM}(keeping existing)${NC}"
    fi
    echo ""

    if ! confirm "Apply this network configuration?"; then
        log_warn "Network configuration skipped."
        return 0
    fi

    cp $INTERFACES_FILE "${INTERFACES_FILE}.bak.$(date +%Y%m%d%H%M%S)"
    log_info "Backed up current network configuration."

    generate_vlan_interfaces "$TRUNK_IFACE" "$WAN_VLAN" "$WAN_IP_MODE" "$WAN_IP" "$WAN_GW" "$LAN_VLAN" "$LAN_IP_MODE" "$LAN_IP" "$ADD_LAN_ROUTE" "$LAN_ROUTE_GW"

    log_success "VLAN network bridges configured!"
    log_warn "Network changes will take effect after reboot or 'ifreload -a'"

    # Apply network changes now
    if confirm "Apply network changes now (ifreload -a)?"; then
        log_step "Applying network configuration..."
        ifreload -a 2>/dev/null || ifdown -a && ifup -a 2>/dev/null || true
        sleep 2

        # Run connectivity checks
        check_network_connectivity
    fi

    NETWORK_CONFIGURED=true
}

function generate_vlan_interfaces() {
    local TRUNK_IFACE="$1"
    local WAN_VLAN="$2"
    local WAN_IP_MODE="$3"
    local WAN_IP="$4"
    local WAN_GW="$5"
    local LAN_VLAN="$6"
    local LAN_IP_MODE="$7"
    local LAN_IP="$8"
    local ADD_LAN_ROUTE="$9"
    local LAN_ROUTE_GW="${10}"

    cat > $INTERFACES_FILE <<EOF
auto lo
iface lo inet loopback

# -----------------------------------------------------------------------
# PHYSICAL TRUNK INTERFACE
# -----------------------------------------------------------------------

iface $TRUNK_IFACE inet manual

# -----------------------------------------------------------------------
# VLAN SUB-INTERFACES
# -----------------------------------------------------------------------

iface ${TRUNK_IFACE}.${WAN_VLAN} inet manual

iface ${TRUNK_IFACE}.${LAN_VLAN} inet manual

# -----------------------------------------------------------------------
# BRIDGES
# -----------------------------------------------------------------------

# WAN Bridge (VLAN $WAN_VLAN)
auto vmbr0
EOF

    if [ "$WAN_IP_MODE" == "1" ]; then
        cat >> $INTERFACES_FILE <<EOF
iface vmbr0 inet static
        bridge-ports ${TRUNK_IFACE}.${WAN_VLAN}
        bridge-stp off
        bridge-fd 0
        bridge-vlan-aware yes
        address $WAN_IP
        gateway $WAN_GW
EOF
    else
        cat >> $INTERFACES_FILE <<EOF
iface vmbr0 inet dhcp
        bridge-ports ${TRUNK_IFACE}.${WAN_VLAN}
        bridge-stp off
        bridge-fd 0
        bridge-vlan-aware yes
EOF
    fi

    cat >> $INTERFACES_FILE <<EOF

# LAN Bridge (VLAN $LAN_VLAN)
auto vmbr1
EOF

    if [ "$LAN_IP_MODE" == "1" ]; then
        cat >> $INTERFACES_FILE <<EOF
iface vmbr1 inet static
        bridge-ports ${TRUNK_IFACE}.${LAN_VLAN}
        bridge-stp off
        bridge-fd 0
        bridge-vlan-aware yes
        address $LAN_IP
EOF
    else
        cat >> $INTERFACES_FILE <<EOF
iface vmbr1 inet dhcp
        bridge-ports ${TRUNK_IFACE}.${LAN_VLAN}
        bridge-stp off
        bridge-fd 0
        bridge-vlan-aware yes
EOF
    fi

    if [[ "$ADD_LAN_ROUTE" =~ ^[Yy]$ ]]; then
        cat >> $INTERFACES_FILE <<EOF
        up ip route add 10.0.0.0/8 via $LAN_ROUTE_GW
EOF
    fi

    cat >> $INTERFACES_FILE <<EOF

source /etc/network/interfaces.d/*
EOF
}

# ============================================
# RAID Configuration
# ============================================

function list_available_disks() {
    echo ""
    log_info "Available disks for RAID:"
    echo ""

    lsblk -d -o NAME,SIZE,TYPE,MODEL,SERIAL | grep -E "^(NAME|sd|nvme)" | while read line; do
        echo "  $line"
    done

    echo ""
    log_info "Disks currently in use (partitioned/mounted):"
    lsblk -o NAME,SIZE,TYPE,MOUNTPOINT | grep -E "(sd|nvme)" | head -20
    echo ""
}

function setup_raid_array() {
    echo ""
    echo -e "${BOLD}=== RAID Array Configuration ===${NC}"
    echo ""

    if ! command -v mdadm &> /dev/null; then
        log_info "Installing mdadm..."
        apt-get install -y mdadm

        # Verify installation
        if ! command -v mdadm &> /dev/null; then
            log_error "Failed to install mdadm. Please install it manually: apt-get install mdadm"
            return 1
        fi
        log_success "mdadm installed successfully"
    fi

    EXISTING_ARRAYS=$(cat /proc/mdstat | grep "^md" | awk '{print "/dev/"$1}' || echo "")
    if [ -n "$EXISTING_ARRAYS" ]; then
        log_warn "Existing RAID arrays detected:"
        cat /proc/mdstat
        echo ""

        if ! confirm "Continue with RAID setup? (Existing arrays will not be modified)"; then
            return 0
        fi
    fi

    list_available_disks

    echo -e "${BOLD}RAID Level Options:${NC}"
    echo ""
    echo -e "  ${CYAN}1)${NC} RAID-1  ${DIM}(Mirroring, 2 disks, 50% capacity)${NC}"
    echo -e "  ${CYAN}2)${NC} RAID-10 ${DIM}(Striped mirrors, 4+ disks, 50% capacity)${NC}"
    echo -e "  ${CYAN}3)${NC} RAID-5  ${DIM}(Striping with parity, 3+ disks, ~66-87% capacity)${NC}"
    echo -e "  ${CYAN}4)${NC} RAID-6  ${DIM}(Striping with double parity, 4+ disks)${NC}"
    echo -e "  ${CYAN}5)${NC} Skip RAID setup"

    RAID_CHOICE=$(prompt_choice "Select RAID level" "1")

    case $RAID_CHOICE in
        1)
            RAID_LEVEL="1"
            MIN_DISKS=2
            ;;
        2)
            RAID_LEVEL="10"
            MIN_DISKS=4
            ;;
        3)
            RAID_LEVEL="5"
            MIN_DISKS=3
            ;;
        4)
            RAID_LEVEL="6"
            MIN_DISKS=4
            ;;
        5)
            log_info "RAID setup skipped."
            return 0
            ;;
        *)
            log_error "Invalid selection."
            return 1
            ;;
    esac

    log_info "Selected RAID level: $RAID_LEVEL (minimum $MIN_DISKS disks required)"
    echo ""
    echo -e "${DIM}Enter the disks to use for the RAID array.${NC}"
    echo -e "${DIM}Example: /dev/sda /dev/sdb /dev/sdc /dev/sdd${NC}"
    echo -e "${DIM}         or: sda sdb sdc sdd${NC}"
    DISK_INPUT=$(prompt_choice "Disks" "")

    DISKS=()
    for disk in $DISK_INPUT; do
        if [[ ! "$disk" == /dev/* ]]; then
            disk="/dev/$disk"
        fi
        DISKS+=("$disk")
    done

    DISK_COUNT=${#DISKS[@]}

    if [ $DISK_COUNT -lt $MIN_DISKS ]; then
        log_error "RAID-$RAID_LEVEL requires at least $MIN_DISKS disks. You provided $DISK_COUNT."
        return 1
    fi

    for disk in "${DISKS[@]}"; do
        if [ ! -b "$disk" ]; then
            log_error "Disk $disk does not exist!"
            return 1
        fi
    done

    echo ""
    log_warn "The following disks will be used:"
    for disk in "${DISKS[@]}"; do
        lsblk -d -o NAME,SIZE,MODEL "$disk"
    done
    echo ""

    log_error "WARNING: ALL DATA ON THESE DISKS WILL BE DESTROYED!"
    echo ""

    if ! confirm "Are you absolutely sure you want to proceed?" "n"; then
        log_info "RAID setup cancelled."
        return 0
    fi

    CONFIRM_DESTROY=$(prompt_choice "Type 'DESTROY' to confirm data destruction" "")
    if [ "$CONFIRM_DESTROY" != "DESTROY" ]; then
        log_info "RAID setup cancelled."
        return 0
    fi

    EXISTING_MD=$(ls /dev/md* 2>/dev/null | grep -oP 'md\d+' | sort -V | tail -1 || echo "")
    if [ -n "$EXISTING_MD" ]; then
        NEXT_MD_NUM=$((${EXISTING_MD#md} + 1))
    else
        NEXT_MD_NUM=0
    fi
    MD_DEVICE="/dev/md${NEXT_MD_NUM}"

    MD_DEVICE=$(prompt_choice "RAID device name" "$MD_DEVICE")

    echo ""
    log_info "Creating RAID-$RAID_LEVEL array on $MD_DEVICE..."

    for disk in "${DISKS[@]}"; do
        log_info "Wiping $disk..."
        wipefs -a "$disk" 2>/dev/null || true
        dd if=/dev/zero of="$disk" bs=1M count=100 2>/dev/null || true
    done

    DISK_STRING="${DISKS[*]}"

    log_info "Running: mdadm --create $MD_DEVICE --level=$RAID_LEVEL --raid-devices=$DISK_COUNT $DISK_STRING"

    mdadm --create "$MD_DEVICE" \
        --level="$RAID_LEVEL" \
        --raid-devices="$DISK_COUNT" \
        --metadata=1.2 \
        ${DISKS[@]}

    log_info "Waiting for array to initialize..."
    sleep 2

    cat /proc/mdstat

    mkdir -p /etc/mdadm
    mdadm --detail --scan >> $MDADM_CONF
    update-initramfs -u

    log_success "RAID array $MD_DEVICE created successfully!"

    RAID_DEVICE="$MD_DEVICE"
    RAID_CONFIGURED=true
}

# ============================================
# LVM Configuration
# ============================================

function setup_lvm() {
    echo ""
    echo -e "${BOLD}=== LVM Configuration ===${NC}"
    echo ""

    if ! command -v pvcreate &> /dev/null; then
        log_info "Installing lvm2..."
        apt-get update && apt-get install -y lvm2
    fi

    log_info "Current LVM configuration:"
    echo ""
    echo "Physical Volumes:"
    pvs 2>/dev/null || echo "  (none)"
    echo ""
    echo "Volume Groups:"
    vgs 2>/dev/null || echo "  (none)"
    echo ""

    if [ -n "$RAID_DEVICE" ]; then
        log_info "Using newly created RAID device: $RAID_DEVICE"
        PV_DEVICE="$RAID_DEVICE"
    else
        echo -e "${BOLD}Available block devices:${NC}"
        lsblk -d -o NAME,SIZE,TYPE | grep -E "(disk|raid)"
        echo ""

        PV_DEVICE=$(prompt_choice "Enter device for Physical Volume (e.g., /dev/md0 or /dev/sda)" "")

        if [ -z "$PV_DEVICE" ]; then
            log_warn "LVM setup skipped."
            return 0
        fi

        if [[ ! "$PV_DEVICE" == /dev/* ]]; then
            PV_DEVICE="/dev/$PV_DEVICE"
        fi
    fi

    if [ ! -b "$PV_DEVICE" ]; then
        log_error "Device $PV_DEVICE does not exist!"
        return 1
    fi

    if pvs "$PV_DEVICE" &>/dev/null; then
        log_warn "$PV_DEVICE is already a Physical Volume."
        EXISTING_VG=$(pvs --noheadings -o vg_name "$PV_DEVICE" | tr -d ' ')
        if [ -n "$EXISTING_VG" ]; then
            log_info "It belongs to Volume Group: $EXISTING_VG"

            if [ "$EXISTING_VG" == "vg0" ]; then
                log_success "Already configured as expected (vg0). LVM setup complete!"
                LVM_CONFIGURED=true
                return 0
            fi
        fi

        if ! confirm "Continue and add to vg0?"; then
            return 0
        fi
    else
        log_info "Creating Physical Volume on $PV_DEVICE..."

        # Wait for udev to settle and device to be fully ready
        log_info "Waiting for device to settle..."
        udevadm settle --timeout=30 2>/dev/null || true

        # If this is a RAID device, check sync status and warn user
        if [[ "$PV_DEVICE" == /dev/md* ]]; then
            SYNC_STATUS=$(cat /proc/mdstat 2>/dev/null | grep -A1 "$(basename $PV_DEVICE)" | grep -oP '\d+%' || echo "")
            if [ -n "$SYNC_STATUS" ]; then
                log_warn "RAID array is still syncing ($SYNC_STATUS). LVM operations may be slow."
                log_info "You can continue, but operations will be faster after sync completes."
            fi
        fi

        pvcreate "$PV_DEVICE"
        log_success "Physical Volume created."
    fi

    if vgs vg0 &>/dev/null; then
        log_info "Volume Group 'vg0' already exists."
        log_info "Extending vg0 with $PV_DEVICE..."
        vgextend vg0 "$PV_DEVICE" 2>/dev/null || log_warn "Device may already be in vg0"
    else
        log_info "Creating Volume Group 'vg0' with $PV_DEVICE..."

        # Ensure udev has settled before vgcreate
        udevadm settle --timeout=30 2>/dev/null || true

        vgcreate vg0 "$PV_DEVICE"
    fi

    log_success "LVM configuration complete!"
    echo ""
    log_info "Final LVM status:"
    echo ""
    echo "Physical Volumes:"
    pvs
    echo ""
    echo "Volume Groups:"
    vgs
    echo ""

    LVM_CONFIGURED=true
}

# ============================================
# Proxmox Storage Configuration
# ============================================

function setup_proxmox_storage() {
    echo ""
    echo -e "${BOLD}=== Proxmox Storage Configuration ===${NC}"
    echo ""

    log_info "Configuring Proxmox storage..."

    log_step "Updating 'local' storage to support images and ISO..."
    pvesm set local --content iso,vztmpl,snippets,backup,images 2>/dev/null || \
        log_warn "Could not update local storage (may already be configured)"

    log_step "Updating 'local-lvm' storage to support images..."
    pvesm set local-lvm --content images,rootdir 2>/dev/null || \
        log_warn "Could not update local-lvm storage (may already be configured)"

    echo ""
    if vgs vg0 &>/dev/null; then
        log_info "Setting up LVM storage on vg0..."

        if pvesm status | grep -q "VM_LVM"; then
            log_info "VM_LVM storage already exists."
            pvesm set VM_LVM --content images,rootdir 2>/dev/null || true
        else
            LVM_STORAGE_NAME=$(prompt_choice "Enter LVM storage name" "VM_LVM")

            log_step "Creating '$LVM_STORAGE_NAME' LVM storage on vg0..."

            # Use standard LVM (not LVM-thin)
            log_step "Adding LVM storage to Proxmox..."
            pvesm add lvm "$LVM_STORAGE_NAME" \
                --vgname vg0 \
                --content images,rootdir \
                --nodes $(hostname) 2>/dev/null || {
                log_error "Failed to create LVM storage"
            }

            if pvesm status | grep -q "$LVM_STORAGE_NAME"; then
                log_success "LVM storage '$LVM_STORAGE_NAME' created!"

                # Store for API registration
                STORAGE_NAME="$LVM_STORAGE_NAME"
            fi
        fi
    else
        log_warn "vg0 not found. Skipping VM_LVM storage setup."
        log_info "Run the LVM configuration step first to create vg0."
    fi

    log_success "Storage configuration complete!"

    echo ""
    log_info "Current storage configuration:"
    pvesm status
    echo ""

    STORAGE_CONFIGURED=true
}

# ============================================
# Restore Mode (non-destructive storage recovery)
# ============================================
# Used when the boot drive (OS + /etc/pve) was lost but the customer data
# disks (RAID members / LVM PVs holding vm-*-disk-* volumes) are intact.
# These functions only ASSEMBLE and ACTIVATE what already exists.

function assemble_raid_array() {
    echo ""
    echo -e "${BOLD}=== RAID Recovery (assemble existing array) ===${NC}"
    echo ""

    if ! command -v mdadm &> /dev/null; then
        log_info "Installing mdadm..."
        apt-get install -y mdadm

        if ! command -v mdadm &> /dev/null; then
            log_error "Failed to install mdadm. Please install it manually: apt-get install mdadm"
            return 1
        fi
    fi

    log_info "Scanning for existing RAID arrays (no disks will be modified)..."

    # Assemble anything with intact superblocks. Non-zero exit just means
    # "nothing new assembled" (already running or no arrays present).
    mdadm --assemble --scan 2>/dev/null || true
    udevadm settle --timeout=30 2>/dev/null || true

    ACTIVE_ARRAYS=$(grep "^md" /proc/mdstat 2>/dev/null | awk '{print "/dev/"$1}' || echo "")

    if [ -z "$ACTIVE_ARRAYS" ]; then
        log_warn "No RAID arrays found. If this node used RAID, the members were not detected."
        log_warn "Check 'mdadm --examine /dev/sdX' on the data disks before proceeding."
        log_info "If this node used plain disks (no RAID), this is expected — continuing."
        return 0
    fi

    log_success "Active RAID arrays:"
    cat /proc/mdstat
    echo ""

    # Warn loudly about degraded arrays but keep going — data is still readable.
    for md in $ACTIVE_ARRAYS; do
        DEGRADED=$(mdadm --detail "$md" 2>/dev/null | grep -cE "degraded|faulty|removed" || true)
        if [ "$DEGRADED" -gt 0 ]; then
            log_warn "$md appears DEGRADED — restore can continue, but replace the failed member ASAP:"
            mdadm --detail "$md" | grep -E "State :|faulty|removed" || true
        fi
    done

    # Persist assembly so the array survives the next reboot of the new OS.
    mkdir -p /etc/mdadm
    touch "$MDADM_CONF"
    mdadm --detail --scan 2>/dev/null | while read -r line; do
        grep -qF "$line" "$MDADM_CONF" 2>/dev/null || echo "$line" >> "$MDADM_CONF"
    done
    update-initramfs -u

    RAID_DEVICE=$(echo "$ACTIVE_ARRAYS" | head -1)
    RAID_CONFIGURED=true
    log_success "RAID array(s) assembled and persisted to $MDADM_CONF"
}

function import_existing_lvm() {
    echo ""
    echo -e "${BOLD}=== LVM Recovery (activate existing volume group) ===${NC}"
    echo ""

    if ! command -v pvscan &> /dev/null; then
        log_info "Installing lvm2..."
        apt-get update && apt-get install -y lvm2
    fi

    log_info "Rescanning LVM metadata (no volumes will be created or wiped)..."
    pvscan --cache 2>/dev/null || true
    vgscan --mknodes 2>/dev/null || true

    echo ""
    echo "Physical Volumes:"
    pvs 2>/dev/null || echo "  (none)"
    echo ""
    echo "Volume Groups:"
    vgs 2>/dev/null || echo "  (none)"
    echo ""

    VG_LIST=$(vgs --noheadings -o vg_name 2>/dev/null | tr -d ' ' || echo "")

    if [ -z "$VG_LIST" ]; then
        log_error "No volume groups found!"
        log_error "Either the RAID array is not assembled yet, or the data disks are not attached."
        log_error "Do NOT run the normal 'lvm' setup — pvcreate would destroy surviving customer data."
        return 1
    fi

    if echo "$VG_LIST" | grep -qx "vg0"; then
        DISCOVERED_VG="vg0"
    elif [ "$(echo "$VG_LIST" | wc -l)" -eq 1 ]; then
        DISCOVERED_VG="$VG_LIST"
        log_warn "Volume group is named '$DISCOVERED_VG' (expected 'vg0')."
    else
        log_info "Multiple volume groups found:"
        echo "$VG_LIST"
        DISCOVERED_VG=$(prompt_choice "Which VG holds the VM disks?" "vg0")
    fi

    log_info "Activating volume group '$DISCOVERED_VG'..."
    vgchange -ay "$DISCOVERED_VG"
    udevadm settle --timeout=30 2>/dev/null || true

    VM_DISK_COUNT=$(lvs --noheadings -o lv_name "$DISCOVERED_VG" 2>/dev/null | grep -c "vm-.*-disk-" || true)

    echo ""
    if [ "$VM_DISK_COUNT" -gt 0 ]; then
        log_success "Found ${VM_DISK_COUNT} VM disk volume(s) in '$DISCOVERED_VG':"
        lvs -o lv_name,lv_size,lv_attr "$DISCOVERED_VG" | grep -E "LV|vm-" || true
    else
        log_warn "Volume group '$DISCOVERED_VG' is active but contains NO vm-*-disk-* volumes."
        log_warn "Verify you activated the right VG before restoring VMs from the admin panel."
    fi

    LVM_CONFIGURED=true
    log_success "LVM recovery complete — '$DISCOVERED_VG' is active."
}

function restore_proxmox_storage() {
    echo ""
    echo -e "${BOLD}=== Proxmox Storage Recovery (re-register existing VG) ===${NC}"
    echo ""

    if [ -z "$DISCOVERED_VG" ]; then
        log_error "No volume group discovered — run the LVM recovery step first."
        return 1
    fi

    log_step "Updating 'local' storage to support images, snippets and ISO..."
    pvesm set local --content iso,vztmpl,snippets,backup,images 2>/dev/null || \
        log_warn "Could not update local storage (may already be configured)"

    # The storage ID must be EXACTLY what the DartNode DB has for this node —
    # every service's disk config references "<storageName>:vm-<vmid>-disk-N".
    DEFAULT_STORAGE_NAME="${TOKEN_STORAGE_NAME:-VM_LVM}"
    if [ -n "$TOKEN_STORAGE_NAME" ]; then
        log_info "DartNode DB expects storage name: ${BOLD}${TOKEN_STORAGE_NAME}${NC}"
    else
        log_warn "DartNode DB has no existing storage name for this node (new node?)."
    fi

    RESTORE_STORAGE_NAME=$(prompt_choice "Proxmox storage ID to map onto '$DISCOVERED_VG'" "$DEFAULT_STORAGE_NAME")

    if pvesm status 2>/dev/null | awk '{print $1}' | grep -qx "$RESTORE_STORAGE_NAME"; then
        log_info "Storage '$RESTORE_STORAGE_NAME' already registered."
        pvesm set "$RESTORE_STORAGE_NAME" --content images,rootdir 2>/dev/null || true
    else
        log_step "Registering LVM storage '$RESTORE_STORAGE_NAME' on VG '$DISCOVERED_VG'..."
        pvesm add lvm "$RESTORE_STORAGE_NAME" \
            --vgname "$DISCOVERED_VG" \
            --content images,rootdir \
            --nodes "$(hostname)" || {
            log_error "Failed to register LVM storage"
            return 1
        }
    fi

    if pvesm status | grep -q "$RESTORE_STORAGE_NAME"; then
        log_success "Storage '$RESTORE_STORAGE_NAME' is registered."
        STORAGE_NAME="$RESTORE_STORAGE_NAME"
    fi

    echo ""
    log_info "Existing VM volumes visible to Proxmox:"
    pvesm list "$RESTORE_STORAGE_NAME" 2>/dev/null | head -40 || true
    echo ""

    STORAGE_CONFIGURED=true
    log_success "Storage recovery complete!"
}

function setup_nfs_storage() {
    echo ""
    echo -e "${BOLD}=== NFS Storage Configuration ===${NC}"
    echo ""

    log_info "Configure NFS shares for cloud-init images, backups, and VM storage."
    echo ""

    if ! confirm "Configure NFS storage?"; then
        log_info "NFS configuration skipped."
        return 0
    fi

    # Check existing NFS storage
    EXISTING_NFS=$(pvesm status --enabled 2>/dev/null | grep "nfs" | awk '{print $1}' || echo "")
    if [ -n "$EXISTING_NFS" ]; then
        log_info "Existing NFS storage:"
        echo "$EXISTING_NFS" | while read storage; do
            echo "  - $storage"
        done
        echo ""
    fi

    # Preferred path: mirror the NFS shares the rest of the fleet mounts
    # (pulled live from an online node via /node/config). No guessed defaults.
    if [ -n "$TOKEN_NFS_STORAGES" ]; then
        log_info "NFS shares in use across the DartNode fleet:"
        while IFS='|' read -r name server export path content options; do
            [ -z "$name" ] && continue
            echo "  - ${name}  ->  ${server}:${export}"
        done <<< "$TOKEN_NFS_STORAGES"
        echo ""

        if confirm "Add these NFS shares (existing ones are skipped)?"; then
            while IFS='|' read -r name server export path content options; do
                [ -z "$name" ] && continue

                if pvesm status 2>/dev/null | awk '{print $1}' | grep -qx "$name"; then
                    log_info "${name} storage already exists. Skipping..."
                    continue
                fi

                log_step "Adding '${name}' (${server}:${export})..."
                PVESM_CMD=(pvesm add nfs "$name" --server "$server" --export "$export" --content "$content")
                [ -n "$path" ] && PVESM_CMD+=(--path "$path")
                [ -n "$options" ] && PVESM_CMD+=(--options "$options")

                PVESM_OUT=$("${PVESM_CMD[@]}" 2>&1) && \
                    log_success "${name} storage added!" || {
                    log_warn "Failed to add ${name} storage:"
                    echo "$PVESM_OUT" | sed 's/^/    /'
                }
            done <<< "$TOKEN_NFS_STORAGES"

            log_success "NFS storage configuration complete!"
            echo ""
            log_info "Current NFS storage:"
            pvesm status --enabled 2>/dev/null | grep -E "(Name|nfs)" || echo "  (none)"
            return 0
        fi

        log_info "Falling back to manual NFS configuration..."
        echo ""
    else
        log_warn "Could not fetch the fleet's NFS shares from the API — using manual configuration."
        echo ""
    fi

    # Cloud-init images NFS share
    echo ""
    log_step "1/3: Cloud-init Images NFS Share"
    if pvesm status | grep -q "cloud-init-img"; then
        log_info "cloud-init-img storage already exists. Skipping..."
    else
        CI_NFS_SERVER=$(prompt_choice "NFS Server IP for cloud-init images" "10.1.0.5")
        CI_NFS_EXPORT=$(prompt_choice "NFS Export path" "/volume1/CloudInitImages")

        PVESM_OUT=$(pvesm add nfs cloud-init-img \
            --server "$CI_NFS_SERVER" \
            --export "$CI_NFS_EXPORT" \
            --path /mnt/pve/cloud-init-img \
            --content snippets,iso,images 2>&1) && \
            log_success "cloud-init-img storage added!" || {
            log_warn "Failed to add cloud-init-img storage:"
            echo "$PVESM_OUT" | sed 's/^/    /'
        }
    fi

    # Backups NFS share
    echo ""
    log_step "2/3: Backups NFS Share"
    if pvesm status | grep -q "dn-backups"; then
        log_info "dn-backups storage already exists. Skipping..."
    else
        BK_NFS_SERVER=$(prompt_choice "NFS Server IP for backups" "10.1.0.5")
        BK_NFS_EXPORT=$(prompt_choice "NFS Export path" "/volume1/PXEBackups")

        PVESM_OUT=$(pvesm add nfs dn-backups \
            --server "$BK_NFS_SERVER" \
            --export "$BK_NFS_EXPORT" \
            --path /mnt/pve/dn-backups \
            --content images,backup,iso 2>&1) && \
            log_success "dn-backups storage added!" || {
            log_warn "Failed to add dn-backups storage:"
            echo "$PVESM_OUT" | sed 's/^/    /'
        }
    fi

    # Additional VM storage (store-1)
    echo ""
    log_step "3/3: Additional VM Storage NFS Share (optional)"
    if pvesm status | grep -q "store-1"; then
        log_info "store-1 storage already exists. Skipping..."
    else
        if confirm "Add store-1 NFS share for additional VM storage?"; then
            ST_NFS_SERVER=$(prompt_choice "NFS Server IP for store-1" "10.1.0.214")
            ST_NFS_EXPORT=$(prompt_choice "NFS Export path" "/mnt/HOU-1-S1-P1/store-1-vms")

            PVESM_OUT=$(pvesm add nfs store-1 \
                --server "$ST_NFS_SERVER" \
                --export "$ST_NFS_EXPORT" \
                --path /mnt/pve/store-1 \
                --content images 2>&1) && \
                log_success "store-1 storage added!" || {
                log_warn "Failed to add store-1 storage:"
                echo "$PVESM_OUT" | sed 's/^/    /'
            }
        fi
    fi

    log_success "NFS storage configuration complete!"
    echo ""
    log_info "Current NFS storage:"
    pvesm status --enabled 2>/dev/null | grep -E "(Name|nfs)" || echo "  (none)"
}

# ============================================
# DartNode Daemon Installation
# ============================================

function install_dartnode_daemon() {
    echo ""
    echo -e "${BOLD}=== DartNode Stats Daemon ===${NC}"
    echo ""

    if command -v dartnode &> /dev/null; then
        log_info "DartNode daemon is already installed."

        CURRENT_VERSION=$(dartnode version 2>/dev/null | head -1 || echo "unknown")
        log_info "Current version: $CURRENT_VERSION"

        if confirm "Update to latest version?"; then
            log_info "Updating DartNode daemon..."
            dartnode update
        fi

        if confirm "Reconfigure daemon settings?"; then
            if [ -n "$REGISTRATION_TOKEN" ]; then
                # Use token-based configuration
                configure_daemon_with_token
            else
                dartnode config
            fi
        fi

        if ! systemctl is-active --quiet dn-daemon; then
            log_info "Starting DartNode daemon..."
            dartnode start
        fi

        DAEMON_INSTALLED=true
        return 0
    fi

    if ! confirm "Install DartNode PVE Stats Daemon?"; then
        log_info "Daemon installation skipped."
        return 0
    fi

    log_info "Downloading and installing DartNode daemon..."
    echo ""

    # Download and run the installer WITHOUT the token flag
    # Registration with the API is handled in Step 7 once all node data is collected
    curl -sSL https://pkg.dev.snaju.com/dn/dartnode.sh | bash -s install

    if ! command -v dartnode &> /dev/null; then
        log_error "DartNode daemon installation failed!"
        return 1
    fi

    log_success "DartNode daemon installed!"

    # Configure the daemon with token settings
    if [ -n "$REGISTRATION_TOKEN" ]; then
        log_step "Configuring daemon with Redis and node settings..."
        configure_daemon_with_token

        # Start the daemon
        log_step "Starting DartNode daemon service..."
        systemctl enable dn-daemon 2>/dev/null || true
        systemctl restart dn-daemon 2>/dev/null || dartnode start 2>/dev/null || true

        if systemctl is-active --quiet dn-daemon; then
            log_success "DartNode daemon is running!"
        else
            log_warn "Daemon may not have started. Check with: systemctl status dn-daemon"
        fi
    fi

    DAEMON_INSTALLED=true
}

function configure_daemon_with_token() {
    # Fetch config from API if not already fetched
    if [ -z "$TOKEN_NODE_NAME" ]; then
        if ! fetch_token_config "$REGISTRATION_TOKEN"; then
            log_error "Failed to fetch configuration from API"
            return 1
        fi
    fi

    local config_dir="/etc/dn-daemon"
    mkdir -p "$config_dir"

    # Use database ID if available (re-setup), otherwise use node name temporarily
    # The node_id will be updated to the DB ID after registration in Step 7
    local DAEMON_NODE_ID="${TOKEN_NODE_DB_ID:-${TOKEN_NODE_NAME}}"

    log_info "Writing daemon configuration..."
    echo ""
    echo "  Node ID: ${DAEMON_NODE_ID}"
    echo "  Node Name: ${TOKEN_NODE_NAME}"
    echo "  Mode: ${TOKEN_MODE:-pve}"
    echo "  Redis Host: ${TOKEN_REDIS_HOST}"
    echo "  Redis Port: ${TOKEN_REDIS_PORT}"
    echo "  Redis DB: ${TOKEN_REDIS_DB}"
    echo "  Redis Password: ${TOKEN_REDIS_PASSWORD:0:8}..."
    echo "  VNC/SSH Secret: ${TOKEN_AUTH_SECRET:0:12}..."
    echo ""

    # Write main config
    cat > "${config_dir}/config.json" <<-EOF
{
    "mode": "${TOKEN_MODE:-pve}",
    "redis_host": "${TOKEN_REDIS_HOST}",
    "redis_port": ${TOKEN_REDIS_PORT},
    "redis_password": "${TOKEN_REDIS_PASSWORD}",
    "redis_db": ${TOKEN_REDIS_DB},
    "node_id": "${DAEMON_NODE_ID}",
    "auto_update": true,
    "update_interval": 300,
    "health_reporting": true,
    "debug": false
}
EOF

    # Write PVE-specific config
    cat > "${config_dir}/pve.json" <<-EOF
{
    "poll_interval": 5,
    "stats_ttl": 120,
    "storage_ttl": 300,
    "qmp_socket_dir": "/var/run/qemu-server",
    "backup_path": "/mnt/pve/dn-backups",
    "backup_interval": 300,
    "worker_count": 4,
    "vm_timeout": 8,
    "collect_timeout": 60,
    "vnc_enabled": true,
    "vnc_port": 5700,
    "vnc_socket_dir": "/var/run/qemu-server",
    "vnc_auth_secret": "${TOKEN_AUTH_SECRET}",
    "ssh_enabled": true,
    "ssh_auth_secret": "${TOKEN_AUTH_SECRET}"
}
EOF

    log_success "Daemon configuration saved to ${config_dir}/"
    log_info "Config files: config.json, pve.json"
}

# ============================================
# API Registration
# ============================================

function register_with_api() {
    echo ""
    echo -e "${BOLD}=== Register Node with DartNode ===${NC}"
    echo ""

    # If we have a token, use it
    if [ -n "$REGISTRATION_TOKEN" ]; then
        register_with_token
        return $?
    fi

    log_info "No registration token provided."
    log_info "You can get a node token from the admin panel."
    echo ""

    NODE_TOKEN=$(prompt_choice "Enter node registration token (or press Enter to skip)" "")

    if [ -z "$NODE_TOKEN" ]; then
        log_info "Node registration skipped. You can register manually later."
        return 0
    fi

    REGISTRATION_TOKEN="$NODE_TOKEN"
    register_with_token
}

function register_with_token() {
    log_info "Registering node with DartNode..."

    # Validate token first
    if ! validate_token "$REGISTRATION_TOKEN"; then
        return 1
    fi

    # Get system information
    log_step "Gathering system information..."

    HOSTNAME=$(hostname)
    WAN_IP=$(ip -4 addr show vmbr0 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1 || echo "")
    LAN_IP=$(ip -4 addr show vmbr1 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1 || echo "")

    if [ -z "$WAN_IP" ]; then
        # Try to get any IP
        WAN_IP=$(hostname -I | awk '{print $1}')
    fi

    log_info "Detected WAN IP: $WAN_IP"
    WAN_IP=$(prompt_choice "Confirm node IP" "$WAN_IP")

    # Create Proxmox API user and token
    log_step "Creating Proxmox API credentials..."

    API_PASSWORD=$(openssl rand -hex 12)

    # Remove existing user/token if present, then create fresh
    pveum user token remove dn-api@pve dartnode-web 2>/dev/null || true
    pveum user delete dn-api@pve 2>/dev/null || true
    pveum user add dn-api@pve --password "$API_PASSWORD" 2>/dev/null || true
    pveum aclmod / -user dn-api@pve -role Administrator 2>/dev/null || true

    log_step "Creating API token..."

    TOKEN_JSON=$(pveum user token add dn-api@pve dartnode-web --comment 'DartNode API Token' --privsep=0 --output-format json 2>/dev/null)

    # Parse the token value from JSON - try multiple methods
    if command -v python3 &> /dev/null; then
        API_SECRET=$(echo "$TOKEN_JSON" | python3 -c "import sys,json; print(json.load(sys.stdin).get('value',''))" 2>/dev/null || echo "")
    fi

    # Fallback to grep if python didn't work
    if [ -z "$API_SECRET" ]; then
        API_SECRET=$(echo "$TOKEN_JSON" | grep -oP '"value"\s*:\s*"\K[^"]+' 2>/dev/null || echo "")
    fi

    # Another fallback using sed
    if [ -z "$API_SECRET" ]; then
        API_SECRET=$(echo "$TOKEN_JSON" | sed -n 's/.*"value"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' 2>/dev/null || echo "")
    fi

    # Debug output
    echo "[DEBUG] TOKEN_JSON: $TOKEN_JSON"
    echo "[DEBUG] Extracted API_SECRET: ${API_SECRET:0:20}..."

    if [ -z "$API_SECRET" ]; then
        log_error "Failed to create API token!"
        log_error "Token JSON output: $TOKEN_JSON"
        return 1
    fi

    pveum aclmod / -token 'dn-api@pve!dartnode-web' -role Administrator 2>/dev/null || true

    log_success "API credentials created!"

    # Install SSH public key
    log_step "Installing DartNode SSH key..."

    # Rotated automation key, 4096-bit (SHA256:8tD9kxY6VpvLlDqIgCCGfG8wsarcr0wHEbn2GmNVFrs).
    # Must stay in sync with ProxmoxService::SSH_PUB_KEY / SSH_KEY_B64.
    DN_SSH_KEY="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQCkdPdPY6YK8bRlTSexQIo8LTbIw1a7k0tg+ulrU6aowPRP4DXdSCrLGmwU4qepvYg91ZorePE8D+PnBx0dj4Ff25TzudC+mWaxXINNYZdpLSHYFafAABdgVWS9zlTM2L1c+XbEhYutzhza9FlDkkJcDQNC4hDDov7P6GyFQIhNKdoD8Mg/uiV/gLg6/lsde3d2CxBPsQmVS4N+RXte/nlKwKA/G+iLok7v2bjcRzwCCY2hjAjh3iYyNENkSFuUh8iHE8GUBNj89D7cHikKvIUFZ0Df/gtTLRyQXqz/kPUk+PFrMb997apUoPYLmq+a8+MHSgd1Yw/pQVrT4MG9RvtModGQPWrdFvIgc3b3pIr4BveJkKmNg9UYnun48/OLEKToj5oREsebnp2aeq5gD6dm2GFYQ+twO/AEDxTqIHD/vm/uIJRkPAKp7RlOO4a/CldqS9UOrs/c2j3NDNGyQemk9vwq3GTC8vV4lP5lUpAVpteD44kUUDuDbrUIdALWVBaxdzAyzboBsQIq54ZLUwqvuaGqjnAofbmAP2RAb2d9VcX5oTSwJLyjpkNHBjaLpu31xbJMj5B4LVdjPjpP0X2gQW3gBoNrZCwiF8ePm62tLTZhQ5A769/26I4SQ7k4VZ0UUwdnG5KOtrK1qGYXn3xRTXgMvfVjbQr25PyCqq38wQ== dartnode-automation-2026"

    # Unique body of the CURRENT key — match on key material, not the comment.
    # The old check grepped for "dartnode", so a host still carrying the REVOKED
    # "dartnode"-commented key would be treated as already-provisioned and never
    # get the rotated key (and the revoked key would never be removed).
    DN_SSH_KEY_BODY="$(echo "$DN_SSH_KEY" | awk '{print $2}')"

    mkdir -p ~/.ssh
    chmod 700 ~/.ssh
    touch ~/.ssh/authorized_keys
    chmod 600 ~/.ssh/authorized_keys

    # Purge the revoked key (2048-bit, ...Byq2P) if present, then install the
    # current key if its exact body isn't already authorized.
    sed -i '/AAABAQCfetdLfmwhiKEZ/d' ~/.ssh/authorized_keys 2>/dev/null || true
    if ! grep -qF "$DN_SSH_KEY_BODY" ~/.ssh/authorized_keys 2>/dev/null; then
        echo "$DN_SSH_KEY" >> ~/.ssh/authorized_keys
    fi

    log_success "SSH key installed!"

    # Register with API
    log_step "Completing registration with DartNode API..."

    local storage_name="${STORAGE_NAME:-VM_LVM}"
    local storage_type="raw"
    # nodeUsername is the Proxmox API user (dn-api@pve)
    # nodePassword is the password for that user
    # apiTokenId is the full token ID (user!tokenname)
    # apiSecret is the token value from Proxmox
    local node_username="dn-api@pve"
    local api_token_id="dn-api@pve!dartnode-web"

    # Debug output
    log_info "Credentials being sent:"
    echo "  Node Username (API User): $node_username"
    echo "  Node Password: ${API_PASSWORD:0:8}..."
    echo "  API Token ID: $api_token_id"
    echo "  API Secret: ${API_SECRET:0:20}..."
    echo "  Storage Name: $storage_name"

    local response
    response=$(curl -sS -X POST "${API_BASE}/node/register" \
        -H "Content-Type: application/json" \
        -d "{\"token\":\"${REGISTRATION_TOKEN}\",\"nodeIp\":\"${WAN_IP}\",\"apiTokenId\":\"${api_token_id}\",\"apiSecret\":\"${API_SECRET}\",\"nodeUsername\":\"${node_username}\",\"nodePassword\":\"${API_PASSWORD}\",\"storageName\":\"${storage_name}\",\"storageType\":\"${storage_type}\"}" 2>/dev/null)

    if echo "$response" | grep -q '"success":true'; then
        local node_id node_name
        node_id=$(json_get "$response" "nodeId")
        node_name=$(json_get "$response" "nodeName")

        log_success "Node registered successfully!"
        echo ""
        log_info "Node DB ID: $node_id"
        log_info "Node Name: $node_name"

        # Update daemon config with the database ID as node_id
        # This ensures the daemon uses the DB ID for Redis keys
        if [ -n "$node_id" ] && [ -f "/etc/dn-daemon/config.json" ]; then
            log_step "Updating daemon config with database ID..."
            if command -v python3 &> /dev/null; then
                python3 -c "
import json
with open('/etc/dn-daemon/config.json', 'r') as f:
    cfg = json.load(f)
cfg['node_id'] = '$node_id'
with open('/etc/dn-daemon/config.json', 'w') as f:
    json.dump(cfg, f, indent=4)
" 2>/dev/null
            else
                sed -i "s/\"node_id\": \"[^\"]*\"/\"node_id\": \"$node_id\"/" /etc/dn-daemon/config.json 2>/dev/null
            fi
            log_success "Daemon config updated with node_id: $node_id"

            # Restart daemon to pick up new node_id
            systemctl restart dn-daemon 2>/dev/null || dartnode restart 2>/dev/null || true
            log_info "Daemon restarted with new node ID"
        fi

        REGISTERED_WITH_API=true
    else
        local error_msg
        error_msg=$(json_get "$response" "error")
        log_error "Registration failed: ${error_msg:-Unknown error}"
        return 1
    fi
}

# ============================================
# Summary
# ============================================

function show_summary() {
    echo ""
    echo -e "${BOLD}============================================${NC}"
    echo -e "${BOLD}       Deployment Summary${NC}"
    echo -e "${BOLD}============================================${NC}"
    echo ""

    if [ "$NETWORK_CONFIGURED" == "true" ]; then
        log_success "Network bridges configured (vmbr0, vmbr1)"
    else
        log_warn "Network bridges: Not configured"
    fi

    if [ "$RAID_CONFIGURED" == "true" ]; then
        log_success "RAID array created: $RAID_DEVICE"
    else
        log_warn "RAID array: Not configured"
    fi

    if [ "$LVM_CONFIGURED" == "true" ]; then
        log_success "LVM configured with vg0"
    else
        log_warn "LVM: Not configured"
    fi

    if [ "$STORAGE_CONFIGURED" == "true" ]; then
        log_success "Proxmox storage types configured"
    else
        log_warn "Proxmox storage: Not configured"
    fi

    if [ "$DAEMON_INSTALLED" == "true" ]; then
        log_success "DartNode daemon installed"
    else
        log_warn "DartNode daemon: Not installed"
    fi

    if [ "$REGISTERED_WITH_API" == "true" ]; then
        log_success "Registered with DartNode"
    else
        log_warn "DartNode registration: Not completed"
    fi

    echo ""
    echo -e "${BOLD}============================================${NC}"
    echo -e "${BOLD}       Next Steps${NC}"
    echo -e "${BOLD}============================================${NC}"
    echo ""

    if [ "$NETWORK_CONFIGURED" == "true" ]; then
        echo "1. Reboot the node or run 'ifreload -a' to apply network changes"
        echo ""
    fi

    if [ "$DAEMON_INSTALLED" == "true" ] && [ "$REGISTERED_WITH_API" != "true" ]; then
        echo "2. Configure the DartNode daemon:"
        echo "   dartnode config"
        echo ""
    fi

    if [ "$REGISTERED_WITH_API" != "true" ]; then
        echo -e "${CYAN}3. Add this node to DartNode Admin Panel:${NC}"
        echo ""
        echo -e "   ${BOLD}https://vader.dartnode.net${NC}"
        echo ""
        echo "   - Navigate to Infrastructure > Nodes"
        echo "   - Click 'Provision New Node'"
        echo "   - Enter this node's details and connection info"
        echo ""
    fi

    if [ "$NETWORK_CONFIGURED" == "true" ] || [ "$RAID_CONFIGURED" == "true" ]; then
        echo -e "${YELLOW}NOTE: A reboot is recommended to apply all changes.${NC}"
        echo ""
        if confirm "Reboot now?"; then
            log_info "Rebooting in 5 seconds..."
            sleep 5
            reboot
        fi
    fi
}

# ============================================
# Interactive Deployment
# ============================================

function run_interactive() {
    header

    log_info "This script will configure a Proxmox node for DartNode."
    echo ""

    # If no token provided, ask for one
    if [ -z "$REGISTRATION_TOKEN" ]; then
        echo -e "${BOLD}A registration token is required to proceed.${NC}"
        echo -e "${DIM}Get your token from: https://vader.dartnode.net/infra/vps/nodes/new${NC}"
        echo ""
        REGISTRATION_TOKEN=$(prompt_choice "Enter your registration token" "")

        if [ -z "$REGISTRATION_TOKEN" ]; then
            log_error "No token provided. Cannot continue."
            exit 1
        fi
    fi

    log_success "Registration token: ${REGISTRATION_TOKEN:0:16}..."
    echo ""
    if ! validate_token "$REGISTRATION_TOKEN"; then
        log_error "Invalid token. Exiting."
        exit 1
    fi
    fetch_token_config "$REGISTRATION_TOKEN"
    echo ""

    echo -e "${BOLD}┌─────────────────────────────────────────────────────────────┐${NC}"
    echo -e "${BOLD}│                    Deployment Steps                         │${NC}"
    echo -e "${BOLD}├─────────────────────────────────────────────────────────────┤${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}1.${NC} Configure network bridges (vmbr0, vmbr1)               ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}2.${NC} Set up RAID array (RAID-1, RAID-5, RAID-6, RAID-10)   ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}3.${NC} Configure LVM (Physical Volume + Volume Group)        ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}4.${NC} Configure Proxmox storage types                       ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}5.${NC} Configure NFS storage (optional)                      ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}6.${NC} Install DartNode stats daemon                         ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}7.${NC} Register with DartNode API                            ${BOLD}│${NC}"
    echo -e "${BOLD}└─────────────────────────────────────────────────────────────┘${NC}"
    echo ""
    echo -e "${DIM}Each step is optional - you can skip any step during the process.${NC}"

    if ! confirm "Ready to begin deployment?"; then
        log_info "Deployment cancelled."
        exit 0
    fi

    # Step 1: Network bridges
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 1 of 7: Network Configuration${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    setup_network_bridges

    # Step 2: RAID array
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 2 of 7: RAID Configuration${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    setup_raid_array

    # Step 3: LVM
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 3 of 7: LVM Configuration${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    setup_lvm

    # Step 4: Proxmox storage
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 4 of 7: Proxmox Storage Configuration${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    setup_proxmox_storage

    # Step 5: NFS storage
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 5 of 7: NFS Storage Configuration${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    setup_nfs_storage

    # Step 6: DartNode daemon
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 6 of 7: DartNode Daemon Installation${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    install_dartnode_daemon

    # Step 7: Register with API
    echo ""
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  STEP 7 of 7: DartNode Registration${NC}"
    echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    register_with_api

    # Show summary
    show_summary
}

# ============================================
# Restore Deployment (lost boot drive, data intact)
# ============================================

function run_restore() {
    header

    echo -e "${BOLD}${YELLOW}RESTORE MODE${NC}"
    echo ""
    log_info "For a node that lost its OS/boot drive (Proxmox configs gone) while the"
    log_info "customer data disks (RAID/LVM with vm-*-disk-* volumes) survived."
    echo ""
    log_warn "This mode NEVER creates/wipes RAID arrays or LVM volumes — it only"
    log_warn "assembles and re-registers what already exists."
    echo ""

    if [ -z "$REGISTRATION_TOKEN" ]; then
        echo -e "${BOLD}A registration token is required to proceed.${NC}"
        echo -e "${DIM}Create a re-setup token from the node's page: https://vader.dartnode.net/infra/vps/nodes${NC}"
        echo ""
        REGISTRATION_TOKEN=$(prompt_choice "Enter your registration token" "")

        if [ -z "$REGISTRATION_TOKEN" ]; then
            log_error "No token provided. Cannot continue."
            exit 1
        fi
    fi

    if ! validate_token "$REGISTRATION_TOKEN"; then
        log_error "Invalid token. Exiting."
        exit 1
    fi
    fetch_token_config "$REGISTRATION_TOKEN"
    echo ""

    # Re-registration matches on the Proxmox hostname (VPSHost.nodeId), and every
    # service pins data['nodeid'] to that hostname. A mismatch would register a
    # brand-new node and orphan every service.
    CURRENT_HOSTNAME=$(hostname)
    if [ -n "$TOKEN_NODE_NAME" ] && [ "$CURRENT_HOSTNAME" != "$TOKEN_NODE_NAME" ]; then
        log_error "Hostname mismatch: this host is '$CURRENT_HOSTNAME' but the token is for '$TOKEN_NODE_NAME'."
        log_error "Set the hostname to match before restoring:"
        log_error "  hostnamectl set-hostname $TOKEN_NODE_NAME   (then re-check /etc/hosts and pvecm status)"
        exit 1
    fi

    if [ -z "$TOKEN_NODE_DB_ID" ] || [ "$TOKEN_NODE_DB_ID" == "null" ]; then
        log_warn "No existing node record found in DartNode for '$TOKEN_NODE_NAME'."
        log_warn "Restore expects a node that is already in the DB (its services reference it)."
        if ! confirm "Continue anyway?" "n"; then
            exit 1
        fi
    fi

    echo -e "${BOLD}┌─────────────────────────────────────────────────────────────┐${NC}"
    echo -e "${BOLD}│                     Restore Steps                           │${NC}"
    echo -e "${BOLD}├─────────────────────────────────────────────────────────────┤${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}1.${NC} Configure network bridges (vmbr0, vmbr1)               ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}2.${NC} Assemble existing RAID array (non-destructive)        ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}3.${NC} Activate existing LVM volume group                    ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}4.${NC} Re-register Proxmox storage on the existing VG        ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}5.${NC} Configure NFS storage (cloud-init images, backups)    ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}6.${NC} Install DartNode stats daemon                         ${BOLD}│${NC}"
    echo -e "${BOLD}│${NC}  ${CYAN}7.${NC} Re-register with DartNode API                         ${BOLD}│${NC}"
    echo -e "${BOLD}└─────────────────────────────────────────────────────────────┘${NC}"
    echo ""
    echo -e "${DIM}After this completes, restore the VMs from the admin panel:${NC}"
    echo -e "${DIM}  vader.dartnode.net > Infra > VPS Nodes > ${TOKEN_NODE_NAME:-<node>} > Tools > Restore VMs${NC}"

    if ! confirm "Ready to begin restore?"; then
        log_info "Restore cancelled."
        exit 0
    fi

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 1 of 7: Network Configuration ━━${NC}"
    setup_network_bridges

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 2 of 7: RAID Recovery ━━${NC}"
    assemble_raid_array

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 3 of 7: LVM Recovery ━━${NC}"
    import_existing_lvm

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 4 of 7: Proxmox Storage Recovery ━━${NC}"
    restore_proxmox_storage

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 5 of 7: NFS Storage Configuration ━━${NC}"
    setup_nfs_storage

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 6 of 7: DartNode Daemon Installation ━━${NC}"
    install_dartnode_daemon

    echo ""
    echo -e "${BOLD}${CYAN}━━ STEP 7 of 7: DartNode Registration ━━${NC}"
    register_with_api

    show_summary

    echo ""
    echo -e "${BOLD}${GREEN}Node restore (host side) complete.${NC}"
    echo -e "Next: open ${BOLD}https://vader.dartnode.net/infra/vps/nodes${NC} > this node > ${BOLD}Tools > Restore VMs${NC}"
    echo -e "to scan for surviving disks and rebuild the VM configurations."
}

function usage() {
    echo "Usage: dn-deploy.sh [command] [options]"
    echo ""
    echo "Commands:"
    echo "  (none)     - Run interactive deployment"
    echo "  restore    - Restore a node after boot-drive loss (non-destructive:"
    echo "               assembles existing RAID, activates existing LVM VG,"
    echo "               re-registers storage — never wipes data disks)"
    echo "  network    - Configure network bridges only"
    echo "  raid       - Configure RAID array only"
    echo "  raid-assemble - Assemble an existing RAID array only (non-destructive)"
    echo "  lvm        - Configure LVM only"
    echo "  lvm-import - Activate an existing LVM VG only (non-destructive)"
    echo "  storage    - Configure Proxmox storage types only"
    echo "  nfs        - Configure NFS storage only"
    echo "  daemon     - Install DartNode daemon only"
    echo "  register   - Register with DartNode only"
    echo "  help       - Show this help"
    echo ""
    echo "Options:"
    echo "  --token <token>  - Registration token from DartNode admin panel"
    echo "  --restore        - Same as the 'restore' command"
    echo ""
    echo "Restore a node whose boot drive died but data disks survived:"
    echo "  curl -sSL https://pkg.dev.snaju.com/dn/dn-deploy.sh | bash -s -- --restore --token YOUR_TOKEN"
    echo ""
    echo "Interactive deployment (recommended):"
    echo "  curl -sSL https://pkg.dev.snaju.com/dn/dn-deploy.sh | bash"
    echo ""
    echo "With registration token:"
    echo "  curl -sSL https://pkg.dev.snaju.com/dn/dn-deploy.sh | bash -s -- --token YOUR_TOKEN"
    echo ""
}

# ============================================
# Parse Arguments
# ============================================

COMMAND=""
while [[ $# -gt 0 ]]; do
    case $1 in
        --token)
            REGISTRATION_TOKEN="$2"
            shift 2
            ;;
        --token=*)
            REGISTRATION_TOKEN="${1#*=}"
            shift
            ;;
        --restore)
            RESTORE_MODE=true
            shift
            ;;
        --)
            shift
            ;;
        *)
            if [ -z "$COMMAND" ]; then
                COMMAND="$1"
            fi
            shift
            ;;
    esac
done

COMMAND=${COMMAND:-interactive}

# --restore flag overrides the default interactive flow
if [ "$RESTORE_MODE" == "true" ] && [ "$COMMAND" == "interactive" ]; then
    COMMAND="restore"
fi

# ============================================
# Main
# ============================================

check_root
check_proxmox
fix_proxmox_repos

# Standalone commands still benefit from API-provided config (fleet NFS shares,
# expected storage name) when a token was passed — fetch it up front. The
# interactive/restore flows fetch it themselves after validating.
case ${COMMAND} in
    interactive|restore|help|--help|-h) ;;
    *)
        if [ -n "$REGISTRATION_TOKEN" ]; then
            fetch_token_config "$REGISTRATION_TOKEN" || true
        fi
        ;;
esac

case ${COMMAND} in
    interactive)
        run_interactive
        ;;
    restore)
        RESTORE_MODE=true
        run_restore
        ;;
    network)
        header
        setup_network_bridges
        ;;
    raid)
        header
        setup_raid_array
        ;;
    raid-assemble)
        header
        assemble_raid_array
        ;;
    lvm)
        header
        setup_lvm
        ;;
    lvm-import)
        header
        import_existing_lvm
        ;;
    storage)
        header
        setup_proxmox_storage
        ;;
    nfs)
        header
        setup_nfs_storage
        ;;
    daemon)
        header
        install_dartnode_daemon
        ;;
    register)
        header
        register_with_api
        ;;
    help|--help|-h)
        header
        usage
        ;;
    *)
        header
        echo "Unknown command: $COMMAND"
        echo ""
        usage
        exit 1
        ;;
esac
