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

# Get absolute path to script (must be done before any cd commands)
SCRIPT_PATH=$(realpath "$0")

# Default Configuration (can be overridden by /etc/wg-admin/config.conf or environment variables)
SERVER_DOMAIN="${SERVER_DOMAIN:-}"
WG_PORT="${WG_PORT:-51820}"
VPN_IPV4_RANGE="${VPN_IPV4_RANGE:-10.10.69.0/24}"
VPN_IPV6_RANGE="${VPN_IPV6_RANGE:-fd69:dead:beef:69::/64}"
WG_INTERFACE="${WG_INTERFACE:-wg0}"
CONF_D_DIR="/etc/wireguard/conf.d"
SERVER_CONF="${CONF_D_DIR}/server.conf"
CLIENT_OUTPUT_DIR="/etc/wireguard/clients"
WG_CONFIG="/etc/wireguard/${WG_INTERFACE}.conf"
DNS_SERVERS="${DNS_SERVERS:-8.8.8.8, 8.8.4.4}"
LOG_FILE="/var/log/wireguard-admin.log"
MIN_DISK_SPACE_MB=100

# Load configuration file if it exists
load_config() {
    local config_file="/etc/wg-admin/config.conf"
    if [[ -f "$config_file" ]]; then
        while IFS='=' read -r key value; do
            [[ "$key" =~ ^[[:space:]]*# ]] && continue
            [[ -z "$key" ]] && continue
            key=$(echo "$key" | xargs)
            value=$(echo "$value" | xargs)
            export "$key=$value"
        done < "$config_file"
    fi
}

# Load config at script start
load_config

# Helper functions to parse CIDR ranges
get_ipv4_network() {
    local cidr="$1"
    echo "$cidr" | cut -d'/' -f1 | sed 's/0$//' | sed 's/\.$//'
}

get_ipv6_network() {
    local cidr="$1"
    echo "$cidr" | cut -d'/' -f1 | sed 's/::$//'
}


# Global variables for cleanup and rollback
TEMP_DIR=""
BACKUP_DIR=""
ROLLBACK_NEEDED=false

# Logging functions with timestamps
log() {
    local level="$1"
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[${timestamp}] [${level}] ${message}" | tee -a "${LOG_FILE}"
}

log_info() {
    log "INFO" "$@"
}

log_error() {
    log "ERROR" "$@" >&2
}

log_warn() {
    log "WARN" "$@"
}

# Cleanup trap handler
cleanup_handler() {
    local exit_code=$?

    # Remove temporary directories
    if [[ -n "${TEMP_DIR}" ]] && [[ -d "${TEMP_DIR}" ]]; then
        log_info "Cleaning up temporary directory: ${TEMP_DIR}"
        rm -rf "${TEMP_DIR}"
    fi

    # Rollback on failure if needed
    if [[ ${ROLLBACK_NEEDED} == true ]] && [[ ${exit_code} -ne 0 ]]; then
        log_error "Operation failed, attempting rollback..."
        rollback_installation
    fi

    exit ${exit_code}
}

# Set up traps for cleanup
trap cleanup_handler EXIT INT TERM HUP

usage() {
    echo "Usage: $0 <command> [options]"
    echo ""
    echo "Commands:"
    echo "  install              Install WireGuard VPN server"
    echo "  add <name> [--psk]   Add a new client (optional PSK for extra security)"
    echo "  list                 List all clients"
    echo "  remove <name>        Remove a client"
    echo "  show <name>          Show client configuration"
    echo "  qr <name>            Show QR code for client"
    echo ""
    echo "Options:"
    echo "  --psk                Use pre-shared key (PSK) for additional security"
    echo "  -h, --help           Show this help"
    echo ""
    echo "Examples:"
    echo "  $0 install"
    echo "  $0 add my-phone --psk"
    echo "  $0 list"
    exit 1
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        echo "ERROR: This script must be run as root. Please use 'sudo $0 $1'" >&2
        exit 1
    fi
}

# Pre-install validation
pre_install_validation() {
    log_info "Running pre-installation validation..."

    # Check root is already done by check_root()

    # Check disk space (at least MIN_DISK_SPACE_MB free)
    local free_space_kb=$(df / | awk 'NR==2 {print $4}')
    local free_space_mb=$((free_space_kb / 1024))
    if [[ ${free_space_mb} -lt ${MIN_DISK_SPACE_MB} ]]; then
        log_error "Insufficient disk space (${free_space_mb}MB free, ${MIN_DISK_SPACE_MB}MB required)"
        echo "ERROR: Insufficient disk space. At least ${MIN_DISK_SPACE_MB}MB free space is required." >&2
        echo "Action: Free up disk space or use a different filesystem." >&2
        exit 1
    fi
    log_info "Disk space validation passed (${free_space_mb}MB free)"

    # Check port availability
    if ss -ulnp | grep -q ":${WG_PORT}"; then
        log_error "Port ${WG_PORT} is already in use"
        echo "ERROR: WireGuard port ${WG_PORT} is already in use." >&2
        echo "Action: Stop the service using port ${WG_PORT} or change WG_PORT in the script." >&2
        echo "To find what's using the port: 'sudo ss -tulnp | grep ${WG_PORT}'" >&2
        exit 1
    fi
    log_info "Port ${WG_PORT} is available"

    # Check required commands will be installed
    log_info "Pre-installation validation passed"
}

# Rollback installation on failure
rollback_installation() {
    log_warn "Rolling back installation..."

    # Stop services
    systemctl stop wg-quick@wg0.service 2>/dev/null || true
    systemctl disable wg-quick@wg0.service 2>/dev/null || true

    # Restore from backup if exists
    if [[ -d "${BACKUP_DIR}" ]]; then
        log_info "Restoring from backup: ${BACKUP_DIR}"

        if [[ -f "${BACKUP_DIR}/wireguard.conf" ]]; then
            cp "${BACKUP_DIR}/wireguard.conf" "${WG_CONFIG}"
        fi

        if [[ -f "${BACKUP_DIR}/nftables.conf" ]]; then
            cp "${BACKUP_DIR}/nftables.conf" /etc/nftables.conf
            nft -f /etc/nftables.conf 2>/dev/null || true
        fi

        if [[ -d "${BACKUP_DIR}/conf.d" ]]; then
            rm -rf "${CONF_D_DIR}"
            cp -r "${BACKUP_DIR}/conf.d" "${CONF_D_DIR}"
            chmod 700 "${CONF_D_DIR}"
        fi
    fi

    # Bring down interface
    wg-quick down wg0 2>/dev/null || true

    log_warn "Rollback complete. Please review the logs and try again."
}


# Backup and rollback functions
BACKUP_DIR="/etc/wg-admin/backups"
BACKUP_RETENTION=10

backup_config() {
    local operation="${1:-manual}"
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    local backup_name="wg-backup-${operation}-${timestamp}"
    local backup_path="${BACKUP_DIR}/${backup_name}"

    log_info "Creating configuration backup: ${backup_name}"

    # Create backup directory if it doesn't exist
    mkdir -p "${BACKUP_DIR}"
    chmod 700 "${BACKUP_DIR}"

    # Create backup directory
    mkdir -p "${backup_path}"

    # Backup WireGuard configurations
    if [[ -d "/etc/wireguard" ]]; then
        cp -a /etc/wireguard "${backup_path}/" 2>/dev/null || true
    fi

    # Backup nftables configurations
    if [[ -f "/etc/nftables.conf" ]]; then
        cp -a /etc/nftables.conf "${backup_path}/" 2>/dev/null || true
    fi

    if [[ -d "/etc/nftables.d" ]]; then
        cp -a /etc/nftables.d "${backup_path}/" 2>/dev/null || true
    fi

    # Backup sysctl configuration
    if [[ -f "/etc/sysctl.d/99-wireguard.conf" ]]; then
        cp -a /etc/sysctl.d/99-wireguard.conf "${backup_path}/" 2>/dev/null || true
    fi

    # Create backup metadata
    cat > "${backup_path}/backup-info.txt" <<EOF
Backup created: ${timestamp}
Operation: ${operation}
Hostname: $(hostname)
WireGuard interface: ${WG_INTERFACE}
EOF

    # Set restrictive permissions
    chmod -R 600 "${backup_path}"

    # Apply retention policy - keep only last ${BACKUP_RETENTION} backups
    apply_retention_policy

    log_info "Backup created successfully: ${backup_path}"
    echo "${backup_path}"
}

restore_config() {
    local backup_path="$1"

    if [[ -z "${backup_path}" ]]; then
        log_error "Backup path required"
        echo "ERROR: Backup path required"
        echo "Usage: restore_config <backup_path>"
        exit 1
    fi

    if [[ ! -d "${backup_path}" ]]; then
        log_error "Backup not found: ${backup_path}"
        echo "ERROR: Backup not found: ${backup_path}"
        exit 1
    fi

    if [[ ! -f "${backup_path}/backup-info.txt" ]]; then
        log_error "Invalid backup directory (missing backup-info.txt)"
        echo "ERROR: Invalid backup directory (missing backup-info.txt)"
        exit 1
    fi

    echo "=== WARNING: Configuration Restore ==="
    echo "This will replace current WireGuard and firewall configurations"
    echo "Backup: ${backup_path}"
    echo ""
    read -p "Continue? (yes/no): " confirm

    if [[ "${confirm}" != "yes" ]]; then
        log_info "Restore cancelled by user"
        echo "Restore cancelled"
        exit 0
    fi

    log_info "Restoring configuration from: ${backup_path}"

    # Stop WireGuard service
    if systemctl is-active --quiet wg-quick@wg0.service 2>/dev/null; then
        log_info "Stopping WireGuard service..."
        systemctl stop wg-quick@wg0.service 2>/dev/null || true
    fi

    # Stop nftables service
    if systemctl is-active --quiet nftables.service 2>/dev/null; then
        log_info "Stopping nftables service..."
        systemctl stop nftables.service 2>/dev/null || true
    fi

    # Create a backup of current state before restore
    backup_config "pre-restore-$(date +%s)"

    # Restore WireGuard configurations
    if [[ -d "${backup_path}/wireguard" ]]; then
        log_info "Restoring WireGuard configurations..."
        rm -rf /etc/wireguard
        cp -a "${backup_path}/wireguard" /etc/
    fi

    # Restore nftables configurations
    if [[ -f "${backup_path}/nftables.conf" ]]; then
        log_info "Restoring nftables configuration..."
        cp -a "${backup_path}/nftables.conf" /etc/
    fi

    if [[ -d "${backup_path}/nftables.d" ]]; then
        log_info "Restoring nftables.d directory..."
        rm -rf /etc/nftables.d
        cp -a "${backup_path}/nftables.d" /etc/
    fi

    # Restore sysctl configuration
    if [[ -f "${backup_path}/99-wireguard.conf" ]]; then
        log_info "Restoring sysctl configuration..."
        cp -a "${backup_path}/99-wireguard.conf" /etc/sysctl.d/
        sysctl -p /etc/sysctl.d/99-wireguard.conf 2>/dev/null || true
    fi

    # Start services
    log_info "Starting nftables service..."
    systemctl start nftables.service

    log_info "Starting WireGuard service..."
    systemctl start wg-quick@wg0.service

    echo ""
    log_info "Configuration restored successfully"
    echo "Configuration restored successfully from: ${backup_path}"
    echo ""
    echo "Note: A pre-restore backup was created for safety"
}

apply_retention_policy() {
    if [[ ! -d "${BACKUP_DIR}" ]]; then
        return
    fi

    # List all backups sorted by modification time (oldest first)
    local backups=($(ls -t "${BACKUP_DIR}" 2>/dev/null | grep -E '^wg-backup-'))

    # If we have more than retention count, remove oldest backups
    if [[ ${#backups[@]} -gt ${BACKUP_RETENTION} ]]; then
        log_info "Applying retention policy (keeping last ${BACKUP_RETENTION} backups)..."
        local to_remove=(${backups[@]:${BACKUP_RETENTION}})

        for backup in "${to_remove[@]}"; do
            log_info "Removing old backup: ${backup}"
            rm -rf "${BACKUP_DIR}/${backup}"
        done
    fi
}

# Validation Functions

validate_client_name() {
    local name="$1"
    if [[ -z "$name" ]]; then
        echo "ERROR: Client name cannot be empty"
        return 1
    fi
    if [[ ! "$name" =~ ^[a-zA-Z0-9_-]+$ ]]; then
        echo "ERROR: Client name must contain only alphanumeric characters, hyphens, and underscores"
        return 1
    fi
    if [[ ${#name} -gt 64 ]]; then
        echo "ERROR: Client name cannot exceed 64 characters"
        return 1
    fi
    return 0
}

validate_ip_availability() {
    local ipv4="$1"
    local ipv6="$2"

    # Check if IPv4 is already in use
    if [[ -n "$ipv4" ]]; then
        if grep -q "AllowedIPs = ${ipv4}/" "${CONF_D_DIR}"/client-*.conf 2>/dev/null; then
            echo "ERROR: IPv4 address ${ipv4} is already in use"
            return 1
        fi
    fi

    # Check if IPv6 is already in use
    if [[ -n "$ipv6" ]]; then
        if grep -q "AllowedIPs = ${ipv6}/" "${CONF_D_DIR}"/client-*.conf 2>/dev/null; then
            echo "ERROR: IPv6 address ${ipv6} is already in use"
            return 1
        fi
    fi

    return 0
}

validate_dns_servers() {
    local dns="$1"
    if [[ -z "$dns" ]]; then
        return 0  # Empty DNS is allowed
    fi

    # Split by comma and validate each DNS server
    IFS=',' read -ra dns_array <<< "$dns"
    for dns_server in "${dns_array[@]}"; do
        dns_server=$(echo "$dns_server" | xargs)  # Trim whitespace
        if [[ ! "$dns_server" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
            echo "ERROR: Invalid DNS server format: ${dns_server}"
            return 1
        fi
    done

    return 0
}

validate_port_range() {
    local port="$1"
    if [[ -z "$port" ]]; then
        echo "ERROR: Port cannot be empty"
        return 1
    fi
    if [[ ! "$port" =~ ^[0-9]+$ ]]; then
        echo "ERROR: Port must be a number"
        return 1
    fi
    if [[ "$port" -lt 1 ]] || [[ "$port" -gt 65535 ]]; then
        echo "ERROR: Port must be between 1 and 65535"
        return 1
    fi
    return 0
}

validate_config_syntax() {
    local config_file="$1"
    if [[ ! -f "$config_file" ]]; then
        echo "ERROR: Config file not found: ${config_file}"
        return 1
    fi

    local in_interface_section=false
    local in_peer_section=false
    local has_interface=false
    local has_peer=false
    local seen_public_keys=()

    while IFS= read -r line; do
        # Skip comments and empty lines
        [[ "$line" =~ ^[[:space:]]*# ]] && continue
        [[ -z "$line" ]] && continue

        # Check section headers
        if [[ "$line" =~ ^\[Interface\] ]]; then
            if [[ "$in_interface_section" == true ]]; then
                echo "ERROR: Duplicate [Interface] section in ${config_file}"
                return 1
            fi
            if [[ "$in_peer_section" == true ]]; then
                echo "ERROR: [Interface] section found after [Peer] section in ${config_file}"
                return 1
            fi
            in_interface_section=true
            has_interface=true
            continue
        fi

        if [[ "$line" =~ ^\[Peer\] ]]; then
            if [[ "$in_interface_section" == false ]] && [[ "$has_interface" == false ]]; then
                echo "ERROR: [Peer] section found before [Interface] section in ${config_file}"
                return 1
            fi
            in_peer_section=true
            has_peer=true
            continue
        fi

        # Validate keys within sections
        if [[ "$line" =~ ^[[:space:]]*PrivateKey[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local key="${BASH_REMATCH[1]}"
            if [[ "$key" =~ ^[[:space:]]*# ]]; then
                continue
            fi
            if [[ ! "$key" =~ ^[A-Za-z0-9+/]{42}[A-Za-z0-9+/=]{2}$ ]] && [[ ! "$key" =~ ^[A-Za-z0-9+/]{43}[A-Za-z0-9+/=]$ ]]; then
                echo "ERROR: Invalid PrivateKey format in ${config_file}"
                return 1
            fi
        fi

        if [[ "$line" =~ ^[[:space:]]*PublicKey[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local key="${BASH_REMATCH[1]}"
            if [[ "$key" =~ ^[[:space:]]*# ]]; then
                continue
            fi
            if [[ ! "$key" =~ ^[A-Za-z0-9+/]{42}[A-Za-z0-9+/=]{2}$ ]] && [[ ! "$key" =~ ^[A-Za-z0-9+/]{43}[A-Za-z0-9+/=]$ ]]; then
                echo "ERROR: Invalid PublicKey format in ${config_file}"
                return 1
            fi
            # Check for duplicate public keys
            for existing_key in "${seen_public_keys[@]}"; do
                if [[ "$existing_key" == "$key" ]]; then
                    echo "ERROR: Duplicate public key found in ${config_file}"
                    return 1
                fi
            done
            seen_public_keys+=("$key")
        fi

        if [[ "$line" =~ ^[[:space:]]*PresharedKey[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local key="${BASH_REMATCH[1]}"
            if [[ "$key" =~ ^[[:space:]]*# ]]; then
                continue
            fi
            if [[ ! "$key" =~ ^[A-Za-z0-9+/]{42}[A-Za-z0-9+/=]{2}$ ]] && [[ ! "$key" =~ ^[A-Za-z0-9+/]{43}[A-Za-z0-9+/=]$ ]]; then
                echo "ERROR: Invalid PresharedKey format in ${config_file}"
                return 1
            fi
        fi

        if [[ "$line" =~ ^[[:space:]]*Address[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local addresses="${BASH_REMATCH[1]}"
            IFS=',' read -ra addr_array <<< "$addresses"
            for addr in "${addr_array[@]}"; do
                addr=$(echo "$addr" | xargs)  # Trim whitespace
                if [[ "$addr" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+/[0-9]+$ ]]; then
                    # IPv4 CIDR validation
                    local ip="${addr%/*}"
                    local prefix="${addr#*/}"
                    if [[ "$prefix" -lt 0 ]] || [[ "$prefix" -gt 32 ]]; then
                        echo "ERROR: Invalid IPv4 prefix length: ${prefix}"
                        return 1
                    fi
                elif [[ "$addr" =~ ^[0-9a-fA-F:]+/[0-9]+$ ]]; then
                    # IPv6 CIDR validation (basic)
                    local prefix="${addr#*/}"
                    if [[ "$prefix" -lt 0 ]] || [[ "$prefix" -gt 128 ]]; then
                        echo "ERROR: Invalid IPv6 prefix length: ${prefix}"
                        return 1
                    fi
                else
                    echo "ERROR: Invalid Address format: ${addr}"
                    return 1
                fi
            done
        fi

        if [[ "$line" =~ ^[[:space:]]*AllowedIPs[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local allowed_ips="${BASH_REMATCH[1]}"
            IFS=',' read -ra ip_array <<< "$allowed_ips"
            for ip in "${ip_array[@]}"; do
                ip=$(echo "$ip" | xargs)  # Trim whitespace
                if [[ "$ip" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+/[0-9]+$ ]]; then
                    local prefix="${ip#*/}"
                    if [[ "$prefix" -lt 0 ]] || [[ "$prefix" -gt 32 ]]; then
                        echo "ERROR: Invalid IPv4 prefix in AllowedIPs: ${ip}"
                        return 1
                    fi
                elif [[ "$ip" =~ ^[0-9a-fA-F:]+/[0-9]+$ ]]; then
                    local prefix="${ip#*/}"
                    if [[ "$prefix" -lt 0 ]] || [[ "$prefix" -gt 128 ]]; then
                        echo "ERROR: Invalid IPv6 prefix in AllowedIPs: ${ip}"
                        return 1
                    fi
                else
                    echo "ERROR: Invalid AllowedIPs format: ${ip}"
                    return 1
                fi
            done
        fi

        if [[ "$line" =~ ^[[:space:]]*DNS[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local dns_servers="${BASH_REMATCH[1]}"
            IFS=',' read -ra dns_array <<< "$dns_servers"
            for dns in "${dns_array[@]}"; do
                dns=$(echo "$dns" | xargs)  # Trim whitespace
                if [[ ! "$dns" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
                    echo "ERROR: Invalid DNS server format: ${dns}"
                    return 1
                fi
            done
        fi

        if [[ "$line" =~ ^[[:space:]]*ListenPort[[:space:]]*=[[:space:]]*(.+)$ ]]; then
            local port="${BASH_REMATCH[1]}"
            if [[ ! "$port" =~ ^[0-9]+$ ]]; then
                echo "ERROR: ListenPort must be a number"
                return 1
            fi
            if [[ "$port" -lt 1 ]] || [[ "$port" -gt 65535 ]]; then
                echo "ERROR: ListenPort must be between 1 and 65535"
                return 1
            fi
        fi

    done < "$config_file"

    return 0
}

get_next_ipv4() {
    local network=$(get_ipv4_network "${VPN_IPV4_RANGE}")
    local used_ips=$(grep -h "AllowedIPs = ${network}" "${CONF_D_DIR}"/client-*.conf 2>/dev/null | cut -d' ' -f3 | cut -d'/' -f1 | sort -V | uniq)

    for i in {2..254}; do
        local ip="${network}${i}"
        if ! echo "$used_ips" | grep -q "^${ip}$"; then
            echo "${ip}"
            return
        fi
    done

    echo "ERROR: No available IPv4 addresses" >&2
    exit 1
}

get_next_ipv6() {
    local network=$(get_ipv6_network "${VPN_IPV6_RANGE}")
    local used_ips=$(grep -h "AllowedIPs = ${network}" "${CONF_D_DIR}"/client-*.conf 2>/dev/null | grep -o "${network}[0-9a-f]*" | sort | uniq)

    for i in {2..254}; do
        local ip=$(printf "${network}%x" $i)
        if ! echo "$used_ips" | grep -q "^${ip}$"; then
            echo "${ip}"
            return
        fi
    done

    echo "ERROR: No available IPv6 addresses" >&2
    exit 1
}

cmd_install() {
    check_root

    # Check if SERVER_DOMAIN is set
    if [[ -z "$SERVER_DOMAIN" ]]; then
        echo "ERROR: SERVER_DOMAIN is not set."
        echo ""
        echo "Please create a configuration file:"
        echo "  sudo mkdir -p /etc/wg-admin"
        echo "  sudo cp config.example /etc/wg-admin/config.conf"
        echo "  sudo nano /etc/wg-admin/config.conf"
        echo ""
        echo "Or set it via environment variable:"
        echo "  SERVER_DOMAIN=vpn.example.com sudo ./wireguard.sh install"
        exit 1
    fi

    log_info "=== WireGuard VPN Installation for Debian 13 ==="
    echo "=== WireGuard VPN Installation for Debian 13 ==="
    echo "Server: ${SERVER_DOMAIN}"
    echo "IPv4 VPN range: ${VPN_IPV4_RANGE}"
    echo "IPv6 VPN range: ${VPN_IPV6_RANGE}"
    echo "Port: ${WG_PORT}"
    echo ""

    # Auto-backup before install (only if config exists)
    if [[ -f "${WG_CONFIG}" ]] || [[ -f "/etc/nftables.conf" ]]; then
        backup_config "install-pre"
    fi

    # Enable rollback flag
    ROLLBACK_NEEDED=true

    # Create backup directory for potential rollback
    BACKUP_DIR=$(mktemp -d)
    log_info "Created backup directory: ${BACKUP_DIR}"

    # Backup existing configs if they exist
    if [[ -f "${WG_CONFIG}" ]]; then
        cp "${WG_CONFIG}" "${BACKUP_DIR}/wireguard.conf"
        log_info "Backed up existing WireGuard config"
    fi
    if [[ -f "/etc/nftables.conf" ]]; then
        cp /etc/nftables.conf "${BACKUP_DIR}/nftables.conf"
        log_info "Backed up existing nftables config"
    fi
    if [[ -d "${CONF_D_DIR}" ]]; then
        cp -r "${CONF_D_DIR}" "${BACKUP_DIR}/conf.d"
        log_info "Backed up existing client configs"
    fi

    # Reset and cleanup existing WireGuard installation
    log_info "Checking for existing WireGuard installation..."
    if systemctl is-enabled --quiet wg-quick@wg0.service 2>/dev/null || systemctl list-unit-files | grep -q wg-quick@wg0.service; then
        log_info "Stopping WireGuard service..."
        echo "Stopping WireGuard service..."
        systemctl stop wg-quick@wg0.service 2>/dev/null || true
        echo "Disabling WireGuard service..."
        systemctl disable wg-quick@wg0.service 2>/dev/null || true
    fi

    if systemctl is-enabled --quiet wg-load-clients.service 2>/dev/null || systemctl list-unit-files | grep -q wg-load-clients.service; then
        echo "Removing old WireGuard client loader service..."
        systemctl disable wg-load-clients.service 2>/dev/null || true
        systemctl stop wg-load-clients.service 2>/dev/null || true
        rm -f /etc/systemd/system/wg-load-clients.service
        systemctl daemon-reload 2>/dev/null || true
    fi

    # Also stop if running but not managed by systemd
    if wg show wg0 &>/dev/null; then
        echo "WireGuard interface is active, bringing down..."
        wg-quick down wg0 2>/dev/null || true
    fi

    if [[ -d "/etc/wireguard" ]]; then
        echo "Removing existing WireGuard configuration..."
        rm -rf /etc/wireguard
    fi

    if [[ -d "/etc/clients.d" ]]; then
        echo "Removing old client directory..."
        rm -rf /etc/clients.d
    fi

    if [[ -d "/etc/wireguard/conf.d" ]]; then
        echo "Removing existing config.d directory..."
        rm -rf /etc/wireguard/conf.d
    fi

    if [[ -d "/etc/wireguard/peer.d" ]]; then
        echo "Removing old peer.d directory..."
        rm -rf /etc/wireguard/peer.d
    fi

    if [[ -d "/root/wireguard-clients" ]]; then
        echo "Removing old client config directory..."
        rm -rf /root/wireguard-clients
    fi

    # Flush nftables rules
    if command -v nft &> /dev/null; then
        echo "Flushing nftables rules..."
        nft flush ruleset 2>/dev/null || true
    fi

    echo "Cleanup complete."

    # Install packages
    echo "Installing packages..."
    apt-get update
    apt-get install -y wireguard wireguard-tools qrencode nftables

    # Get public interface BEFORE configuring nftables
    PUBLIC_INTERFACE=$(ip route get 8.8.8.8 | grep -oP 'dev \K\S+' | head -1)
    echo "Public interface: ${PUBLIC_INTERFACE}"

    # Enable IP forwarding
    echo "Enabling IP forwarding..."
    cat > /etc/sysctl.d/99-wireguard.conf <<EOF
net.ipv4.ip_forward = 1
net.ipv6.conf.all.forwarding = 1
EOF
    sysctl -p /etc/sysctl.d/99-wireguard.conf

    # Configure nftables
    echo "Configuring nftables firewall..."
    mkdir -p /etc/nftables.d

    # Create wireguard config in /etc/nftables.d
    cat > /etc/nftables.d/wireguard.conf <<EOF
#!/usr/sbin/nft -f
# nftables configuration for WireGuard VPN

flush ruleset

table inet wireguard {
    chain prerouting {
        type filter hook prerouting priority -150;

        # Rate limiting for SSH (3 connections per minute, burst 5)
        tcp dport 22 ct state new limit rate 3/minute burst 5 packets accept
        tcp dport 22 ct state new limit rate 3/minute burst 5 packets drop

        # Rate limiting for WireGuard (10 packets per second)
        udp dport 51820 limit rate 10/second burst 20 packets accept
        udp dport 51820 limit rate 10/second burst 20 packets drop
    }

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

        iifname lo accept

        # Connection tracking bypass for WireGuard UDP traffic
        iifname "${PUBLIC_INTERFACE}" udp dport 51820 notrack

        ct state established,related accept
        ct state invalid drop

        # Allow SSH
        tcp dport 22 accept

        # Allow WireGuard UDP traffic (already tracked via notrack)
        udp dport 51820 accept

        # ICMPv4
        icmp type { echo-request, echo-reply } accept

        # ICMPv6 - ensure neighbor discovery is allowed
        icmpv6 type { echo-request, echo-reply, nd-neighbor-solicit, nd-neighbor-advert, nd-router-solicit, nd-router-advert } accept
    }

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

        ct state established,related accept
        iifname wg0 accept
        oifname wg0 accept
    }

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

        # Connection tracking bypass for WireGuard UDP traffic
        oifname "${PUBLIC_INTERFACE}" udp dport 51820 notrack
    }
}

table ip nat {
    chain postrouting {
        type nat hook postrouting priority 100; policy accept;

        # TCP MSS clamping for MTU issues (clamp to 1360)
        oifname "${PUBLIC_INTERFACE}" tcp flags syn tcp option maxseg size set 1360

        oifname "${PUBLIC_INTERFACE}" ip saddr ${VPN_IPV4_RANGE} masquerade
    }
}

table ip6 nat {
    chain postrouting {
        type nat hook postrouting priority 100; policy accept;

        # TCP MSS clamping for IPv6 MTU issues
        oifname "${PUBLIC_INTERFACE}" tcp flags syn tcp option maxseg size set 1360

        oifname "${PUBLIC_INTERFACE}" ip6 saddr ${VPN_IPV6_RANGE} masquerade
    }
}
EOF

    # Create /etc/nftables.conf that includes wireguard.conf
    # The nftables service loads /etc/nftables.conf by default
    cat > /etc/nftables.conf <<EOF
#!/usr/sbin/nft -f
# nftables configuration automatically generated by wireguard.sh
flush ruleset
include "/etc/nftables.d/wireguard.conf"
EOF

    chmod 600 /etc/nftables.conf
    chmod 600 /etc/nftables.d/wireguard.conf

    # Validate nftables configuration before applying
    echo "Validating nftables configuration..."
    if ! nft check -f /etc/nftables.conf; then
        echo "ERROR: nftables configuration validation failed"
        exit 1
    fi
    echo "nftables configuration is valid"

    nft -f /etc/nftables.conf

    # Enable nftables service to load rules on boot
    echo "Enabling nftables service..."
    systemctl enable nftables.service
    systemctl start nftables.service

    # Verify nftables is running
    if ! systemctl is-active --quiet nftables.service; then
        echo "ERROR: nftables service failed to start"
        echo "=== Service Status ==="
        systemctl status nftables.service
        exit 1
    fi

    echo "nftables service started successfully"

    # Generate server keys
    log_info "Generating server keys..."
    echo "Generating server keys..."
    mkdir -p /etc/wireguard
    cd /etc/wireguard

    # Generate keys with atomic write
    local temp_private=$(mktemp)
    local temp_public=$(mktemp)
    wg genkey > "$temp_private"
    wg pubkey < "$temp_private" > "$temp_public"
    SERVER_PRIVATE_KEY=$(cat "$temp_private")
    SERVER_PUBLIC_KEY=$(cat "$temp_public")

    # Move to final location with proper permissions
    mv "$temp_private" server_private.key
    mv "$temp_public" server_public.key
    chmod 600 server_private.key
    chmod 644 server_public.key
    log_info "Server keys generated and secured"

    # Create config.d directory
    mkdir -p "${CONF_D_DIR}"

    # Create server.conf in conf.d directory
    cat > "${SERVER_CONF}" <<EOF
[Interface]
PrivateKey = ${SERVER_PRIVATE_KEY}
Address = ${VPN_IPV4_RANGE%.*}.1/${VPN_IPV4_RANGE#*/}, ${VPN_IPV6_RANGE%:*}:1/${VPN_IPV6_RANGE#*/}
ListenPort = ${WG_PORT}
EOF

    # Set permissions
    chmod 600 "${SERVER_CONF}"
    chmod 700 "${CONF_D_DIR}"
    chmod 600 /etc/wireguard/server_private.key

    chmod +x "${SCRIPT_PATH}"

    # Load any existing clients BEFORE starting service
    echo "Loading configs from ${CONF_D_DIR}..."
    "${SCRIPT_PATH}" load-clients

    # Enable WireGuard service
    systemctl enable wg-quick@wg0.service

    # Start WireGuard service
    echo ""
    echo "Starting WireGuard..."
    # Bring down any existing wg0 interface before starting service
    if wg show wg0 &>/dev/null; then
        wg-quick down wg0 2>/dev/null || true
    fi
    systemctl start wg-quick@wg0.service

    # Wait for WireGuard to initialize
    sleep 2

    # Verify WireGuard is running
    if ! systemctl is-active --quiet wg-quick@wg0.service; then
        echo "ERROR: WireGuard service failed to start"
        echo "=== Service Status ==="
        systemctl status wg-quick@wg0.service
        exit 1
    fi

    # Verify correct listening port
    ACTUAL_PORT=$(wg show wg0 listen-port)
    if [[ "$ACTUAL_PORT" != "$WG_PORT" ]]; then
        echo "ERROR: WireGuard listening on port $ACTUAL_PORT instead of $WG_PORT"
        echo "=== Config File ==="
        grep ListenPort "${WG_CONFIG}"
        echo "=== Running Interface ==="
        wg show wg0
        exit 1
    fi

    echo "WireGuard started successfully on port $ACTUAL_PORT"

    echo ""
    log_info "=== Installation Complete ==="
    echo "=== Installation Complete ==="
    echo ""
    echo "Server Public Key: ${SERVER_PUBLIC_KEY}"
    echo "Endpoint: ${SERVER_DOMAIN}:${WG_PORT}"
    echo "VPN IPv4 Range: ${VPN_IPV4_RANGE}"
    echo "VPN IPv6 Range: ${VPN_IPV6_RANGE}"
    echo ""
    echo "Use '$0 add <name> [--psk]' to add clients"
    echo "Use '$0 list' to list clients"
    echo "Configs will be merged from: ${CONF_D_DIR}"
    echo "  - ${SERVER_CONF} (server interface)"
    echo "  - ${CONF_D_DIR}/client-*.conf (client peers)"
    echo ""
    echo "Check status: wg show"
    echo "System status: systemctl status wg-quick@wg0"
    echo ""
    echo "To manually reload configs: $0 load-clients"

    # Disable rollback flag - installation successful
    ROLLBACK_NEEDED=false

    # Clean up backup directory
    if [[ -d "${BACKUP_DIR}" ]]; then
        rm -rf "${BACKUP_DIR}"
        log_info "Cleaned up backup directory"
    fi

    log_info "Installation completed successfully"
}

cmd_add() {
    check_root

    local name=""
    local use_psk=false

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --psk)
                use_psk=true
                shift
                ;;
            *)
                if [[ -z "$name" ]]; then
                    name="$1"
                fi
                shift
                ;;
        esac
    done

    if [[ -z "$name" ]]; then
        log_error "Client name required"
        echo "ERROR: Client name required" >&2
        usage
    fi

    # Validate client name
    if ! validate_client_name "$name"; then
        exit 1
    fi

    if [[ -f "${CONF_D_DIR}/client-${name}.conf" ]]; then
        log_error "Client '${name}' already exists"
        echo "ERROR: Client '${name}' already exists" >&2
        exit 1
    fi

    if [[ ! -f "/etc/wireguard/server_public.key" ]]; then
        log_error "WireGuard server not installed. Run '$0 install' first"
        echo "ERROR: WireGuard server not installed. Run '$0 install' first" >&2
        exit 1
    fi

    local server_public_key=$(cat /etc/wireguard/server_public.key)

    log_info "Creating client '${name}'${use_psk:+ with PSK}"

    # Auto-backup before add
    backup_config "add-${name}"

    # Use global TEMP_DIR for trap cleanup
    TEMP_DIR=$(mktemp -d)
    log_info "Created temporary directory: ${TEMP_DIR}"

    pushd "$TEMP_DIR" > /dev/null
    wg genkey | tee client_private.key | wg pubkey > client_public.key
    local client_private_key=$(cat client_private.key)
    local client_public_key=$(cat client_public.key)

    # Generate PSK if requested
    local psk_line=""
    if [[ "$use_psk" == true ]]; then
        wg genpsk > client_psk.key
        local client_psk=$(cat client_psk.key)
        psk_line="PresharedKey = ${client_psk}"
        log_info "Generated pre-shared key for additional security"
    fi
    popd > /dev/null

    local client_ipv4=$(get_next_ipv4)
    local client_ipv6=$(get_next_ipv6)

    # Validate IP availability
    if ! validate_ip_availability "$client_ipv4" "$client_ipv6"; then
        exit 1
    fi

    # Validate DNS servers
    if ! validate_dns_servers "$DNS_SERVERS"; then
        exit 1
    fi

    mkdir -p "${CONF_D_DIR}"
    # Atomic write to temp file then move
    local temp_server_conf=$(mktemp)
    cat > "$temp_server_conf" <<HEREDOC
[Peer]
# ${name}
PublicKey = ${client_public_key}
${psk_line}
AllowedIPs = ${client_ipv4}/32, ${client_ipv6}/128
HEREDOC
    mv "$temp_server_conf" "${CONF_D_DIR}/client-${name}.conf"
    chmod 600 "${CONF_D_DIR}/client-${name}.conf"
    log_info "Created server config for '${name}'"

    local client_output="${CLIENT_OUTPUT_DIR}/${name}.conf"
    mkdir -p "${CLIENT_OUTPUT_DIR}"
    # Atomic write to temp file then move
    local temp_client_conf=$(mktemp)
    cat > "$temp_client_conf" <<HEREDOC
[Interface]
PrivateKey = ${client_private_key}
Address = ${client_ipv4}/24, ${client_ipv6}/64
DNS = ${DNS_SERVERS}

[Peer]
PublicKey = ${server_public_key}
${psk_line}
Endpoint = ${SERVER_DOMAIN}:${WG_PORT}
AllowedIPs = 0.0.0.0/0, ::/0
PersistentKeepalive = 25
HEREDOC
    mv "$temp_client_conf" "$client_output"
    chmod 600 "$client_output"
    log_info "Created client config for '${name}'"

    qrencode -t ansiutf8 < "$client_output" > "${CLIENT_OUTPUT_DIR}/${name}.qr"
    chmod 600 "${CLIENT_OUTPUT_DIR}/${name}.qr"

    # TEMP_DIR will be cleaned up by the trap handler
    log_info "Temporary directory will be cleaned up by trap handler"

    "${SCRIPT_PATH}" load-clients

    # Reset TEMP_DIR since cleanup is done
    TEMP_DIR=""

    log_info "Client '${name}' added successfully"
    echo "Client '${name}' added successfully${use_psk:+ with PSK}"
    echo ""
    echo "Client IPv4: ${client_ipv4}"
    echo "Client IPv6: ${client_ipv6}"
    echo "Config saved to: ${client_output}"
    echo "QR code saved to: ${CLIENT_OUTPUT_DIR}/${name}.qr"
    if [[ "$use_psk" == true ]]; then
        echo "Pre-shared key enabled for additional security"
    fi
}

cmd_list() {
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: This command must be run as root"
        exit 1
    fi

    echo "WireGuard Clients:"
    echo ""

    if [[ ! -d "${CONF_D_DIR}" ]] || [[ -z "$(ls -A ${CONF_D_DIR}/client-*.conf 2>/dev/null)" ]]; then
        echo "No clients found"
        return
    fi

    printf "%-20s %-20s %-35s %s\n" "Name" "IPv4" "IPv6" "Status"
    printf "%-20s %-20s %-35s %s\n" "----" "----" "----" "------"

    for client_file in "${CONF_D_DIR}"/client-*.conf; do
        local name=$(basename "${client_file}" .conf | sed 's/^client-//')
        local ipv4=$(grep "AllowedIPs = $(get_ipv4_network "${VPN_IPV4_RANGE}")" "${client_file}" | grep -o "$(get_ipv4_network "${VPN_IPV4_RANGE}")[0-9]*" || echo "N/A")
        local ipv6=$(grep "AllowedIPs = fd69:dead:beef:69::" "${client_file}" | grep -o 'fd69:dead:beef:69::[0-9a-f]*' || echo "N/A")
        local public_key=$(grep "PublicKey = " "${client_file}" | cut -d' ' -f3)

        local status="Disconnected"
        if wg show wg0 peers | grep -q "$public_key"; then
            status="Connected"
        fi

        printf "%-20s %-20s %-35s %s\n" "$name" "$ipv4" "$ipv6" "$status"
    done
}

cmd_remove() {
    check_root

    local name=$1

    if [[ -z "$name" ]]; then
        log_error "Client name required"
        echo "ERROR: Client name required" >&2
        usage
    fi

    # Validate client name
    if ! validate_client_name "$name"; then
        exit 1
    fi

    local client_file="${CONF_D_DIR}/client-${name}.conf"

    if [[ ! -f "$client_file" ]]; then
        log_error "Client '${name}' not found"
        echo "ERROR: Client '${name}' not found" >&2
        exit 1
    fi

    # Auto-backup before remove
    backup_config "remove-${name}"

    log_info "Removing client '${name}'..."
    rm "$client_file"
    rm -f "${CLIENT_OUTPUT_DIR}/${name}.conf"
    rm -f "${CLIENT_OUTPUT_DIR}/${name}.qr"

    "${SCRIPT_PATH}" load-clients

    log_info "Client '${name}' removed successfully"
    echo "Client '${name}' removed successfully"
}

cmd_show() {
    if [[ -z "$1" ]]; then
        echo "ERROR: Client name required"
        usage
    fi

    local name=$1
    local client_file="${CLIENT_OUTPUT_DIR}/${name}.conf"

    if [[ ! -f "$client_file" ]]; then
        echo "ERROR: Client '${name}' not found"
        exit 1
    fi

    cat "$client_file"
}

cmd_qr() {
    if [[ -z "$1" ]]; then
        echo "ERROR: Client name required"
        usage
    fi

    local name=$1
    local qr_file="${CLIENT_OUTPUT_DIR}/${name}.qr"

    if [[ ! -f "$qr_file" ]]; then
        echo "ERROR: QR code for '${name}' not found"
        exit 1
    fi

    cat "$qr_file"
}

cmd_load_clients() {
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: This command must be run as root"
        exit 1
    fi

    TEMP_CONFIG="${WG_CONFIG}.tmp"

    # Ensure conf.d directory exists with proper permissions
    mkdir -p "${CONF_D_DIR}"
    chmod 700 "${CONF_D_DIR}"

    # Verify server.conf exists
    if [[ ! -f "${SERVER_CONF}" ]]; then
        echo "ERROR: Server config not found at ${SERVER_CONF}. Run '$0 install' first"
        exit 1
    fi

    # Start with server.conf (Interface section)
    cat "${SERVER_CONF}" > "${TEMP_CONFIG}"
    echo "" >> "${TEMP_CONFIG}"

    # Append all client-*.conf files alphabetically (no comments for clean parsing)
    if compgen -G "${CONF_D_DIR}/client-*.conf" > /dev/null; then
        for client_file in "${CONF_D_DIR}"/client-*.conf; do
            cat "${client_file}" >> "${TEMP_CONFIG}"
            echo "" >> "${TEMP_CONFIG}"
        done
    fi

    # Set proper permissions
    chmod 600 "${TEMP_CONFIG}"

    # Verify temp file was created and is non-empty
    if [[ ! -s "${TEMP_CONFIG}" ]]; then
        echo "ERROR: Failed to generate configuration file"
        rm -f "${TEMP_CONFIG}"
        exit 1
    fi

    # Replace config file
    mv "${TEMP_CONFIG}" "${WG_CONFIG}"
    chmod 600 "${WG_CONFIG}"

    # Validate config syntax
    if ! validate_config_syntax "${WG_CONFIG}"; then
        echo "ERROR: Configuration file validation failed"
        rm -f "${WG_CONFIG}"
        exit 1
    fi

    # Check if WireGuard is running
    if wg show wg0 &>/dev/null; then
        # Reload WireGuard using systemctl
        if ! systemctl reload wg-quick@wg0.service; then
            echo "ERROR: Failed to reload WireGuard configuration"
            echo "=== Config file content ==="
            cat "${WG_CONFIG}"
            exit 1
        fi
        echo "WireGuard reloaded successfully"
    else
        # Start WireGuard
        if ! wg-quick up wg0; then
            echo "ERROR: Failed to start WireGuard"
            echo "=== Config file content ==="
            cat "${WG_CONFIG}"
            exit 1
        fi
        echo "WireGuard started successfully"
    fi
}

# Main entry point
case "${1:-}" in
    install)
        cmd_install
        ;;
    add)
        shift  # Remove 'add' from arguments
        cmd_add "$@"
        ;;
    list)
        cmd_list
        ;;
    remove)
        cmd_remove "${2:-}"
        ;;
    show)
        cmd_show "${2:-}"
        ;;
    qr)
        cmd_qr "${2:-}"
        ;;
    load-clients)
        cmd_load_clients
        ;;
    -h|--help|help|"")
        usage
        ;;
    *)
        log_error "Unknown command '${1}'"
        echo "ERROR: Unknown command '${1}'" >&2
        usage
        ;;
esac