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

# Configuration
SERVER_DOMAIN="velkhana.calmcacil.dev"
WG_PORT="51820"
VPN_IPV4_RANGE="10.10.69.0/24"
VPN_IPV6_RANGE="fd69:dead:beef:69::/64"
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="8.8.8.8, 8.8.4.4"

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

usage() {
    echo "Usage: $0 <command> [options]"
    echo ""
    echo "Commands:"
    echo "  install              Install WireGuard VPN server"
    echo "  add <name>           Add a new client"
    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 "  -h, --help           Show this help"
    exit 1
}

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

get_next_ipv4() {
    local used_ips=$(grep -h "AllowedIPs = 10.10.69." "${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="10.10.69.${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 used_ips=$(grep -h "AllowedIPs = fd69:dead:beef:69::" "${CONF_D_DIR}"/client-*.conf 2>/dev/null | grep -o 'fd69:dead:beef:69::[0-9a-f]*' | sort | uniq)

    for i in {2..254}; do
        local ip=$(printf "fd69:dead:beef:69::%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

    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 ""

    # Reset and cleanup existing WireGuard installation
    echo "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
        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 input {
        type filter hook input priority 0; policy drop;

        iifname lo accept
        ct state established,related accept
        ct state invalid drop

        tcp dport 22 accept
        udp dport 51820 accept

        icmp type { echo-request, echo-reply } accept
        icmpv6 type { echo-request, echo-reply, nd-neighbor-solicit, nd-neighbor-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;
    }
}

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

        oifname "${PUBLIC_INTERFACE}" ip saddr 10.10.69.0/24 masquerade
    }
}

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

        oifname "${PUBLIC_INTERFACE}" ip6 saddr fd69:dead:beef:69::/64 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

    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
    echo "Generating server keys..."
    mkdir -p /etc/wireguard
    cd /etc/wireguard
    wg genkey | tee server_private.key | wg pubkey > server_public.key
    SERVER_PRIVATE_KEY=$(cat server_private.key)
    SERVER_PUBLIC_KEY=$(cat server_public.key)

    # 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 = 10.10.69.1/24, fd69:dead:beef:69::1/64
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 ""
    echo "=== Installation Complete ==="
    echo ""
    echo "Server Public Key: ${SERVER_PUBLIC_KEY}"
    echo "Endpoint: ${SERVER_DOMAIN}:${WG_PORT}"
    echo "VPN IPv4 Range: 10.10.69.0/24"
    echo "VPN IPv6 Range: fd69:dead:beef:69::/64"
    echo ""
    echo "Use '$0 add <name>' 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"
}

cmd_add() {
    check_root

    local name=$1

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

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

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

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

    local client_keys_dir=$(mktemp -d)
    pushd "$client_keys_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)
    popd > /dev/null

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

    mkdir -p "${CONF_D_DIR}"
    cat > "${CONF_D_DIR}/client-${name}.conf" <<HEREDOC
[Peer]
# ${name}
PublicKey = ${client_public_key}
AllowedIPs = ${client_ipv4}/32, ${client_ipv6}/128
HEREDOC
    chmod 600 "${CONF_D_DIR}/client-${name}.conf"

    local client_output="${CLIENT_OUTPUT_DIR}/${name}.conf"
    mkdir -p "${CLIENT_OUTPUT_DIR}"
    cat > "$client_output" <<HEREDOC
[Interface]
PrivateKey = ${client_private_key}
Address = ${client_ipv4}/24, ${client_ipv6}/64
DNS = ${DNS_SERVERS}

[Peer]
PublicKey = ${server_public_key}
Endpoint = ${SERVER_DOMAIN}:${WG_PORT}
AllowedIPs = 0.0.0.0/0, ::/0
PersistentKeepalive = 25
HEREDOC
    chmod 600 "$client_output"

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

    rm -rf "$client_keys_dir"

    "${SCRIPT_PATH}" load-clients

    echo "Client '${name}' added successfully"
    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"
}

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 = 10.10.69." "${client_file}" | grep -o '10.10.69\.[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
        echo "ERROR: Client name required"
        usage
    fi

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

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

    rm "$client_file"
    rm -f "${CLIENT_OUTPUT_DIR}/${name}.conf"
    rm -f "${CLIENT_OUTPUT_DIR}/${name}.qr"

    "${SCRIPT_PATH}" load-clients

    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}"

    # 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)
        cmd_add "${2:-}"
        ;;
    list)
        cmd_list
        ;;
    remove)
        cmd_remove "${2:-}"
        ;;
    show)
        cmd_show "${2:-}"
        ;;
    qr)
        cmd_qr "${2:-}"
        ;;
    load-clients)
        cmd_load_clients
        ;;
    -h|--help|help|"")
        usage
        ;;
    *)
        echo "ERROR: Unknown command '${1}'"
        usage
        ;;
esac