#!/bin/bash

# Prevent execution with non-root privileges
if [ "$EUID" -ne 0 ]; then
    echo "Error: This script must be run as root or with sudo." >&2
    exit 1
fi

echo "========================================="
echo "   WireGuard Auto-Installation Script"
echo "========================================="

# Detect the operating system
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$ID
    VERSION_ID=$VERSION_ID
elif type lsb_release >/dev/null 2>&1; then
    OS=$(lsb_release -si | tr '[:upper:]' '[:lower:]')
else
    OS=$(uname -s | tr '[:upper:]' '[:lower:]')
fi

echo "Detected OS: $OS ($VERSION_ID)"

# Function to install WireGuard on Debian/Ubuntu
install_debian_ubuntu() {
    echo "Updating package lists..."
    apt-get update -y
    
    echo "Installing WireGuard and dependencies..."
    apt-get install -y wireguard wireguard-tools iptables resolvconf
}

# Function to install WireGuard on CentOS/RHEL/Rocky/AlmaLinux
install_rhel() {
    echo "Installing EPEL repository..."
    if [ "$OS" = "centos" ] || [ "$OS" = "rhel" ]; then
        yum install -y epel-release
        # For CentOS 8 / RHEL 8, we might also need elrepo or copr depending on version, 
        # but modern Rocky/Alma/RHEL 8/9 have wireguard in baseos/epel or kernel >= 5.6.
    fi
    
    echo "Installing WireGuard..."
    if type dnf >/dev/null 2>&1; then
        dnf install -y elrepo-release epel-release 2>/dev/null || true
        dnf install -y wireguard-tools wireguard-dkms iptables || dnf install -y wireguard-tools iptables
    else
        yum install -y wireguard-tools wireguard-dkms iptables || yum install -y wireguard-tools iptables
    fi
}

# Function to install WireGuard on Fedora
install_fedora() {
    echo "Installing WireGuard..."
    dnf install -y wireguard-tools iptables
}

# Function to install WireGuard on Arch Linux
install_arch() {
    echo "Installing WireGuard..."
    pacman -Sy --noconfirm wireguard-tools
}

# Run installation based on OS
case "$OS" in
    ubuntu|debian|raspbian)
        install_debian_ubuntu
        ;;
    centos|rhel|rocky|almalinux)
        install_rhel
        ;;
    fedora)
        install_fedora
        ;;
    arch)
        install_arch
        ;;
    *)
        echo "Unsupported OS: $OS"
        echo "Attempting generic package installation..."
        if type apt-get >/dev/null 2>&1; then
            install_debian_ubuntu
        elif type dnf >/dev/null 2>&1 || type yum >/dev/null 2>&1; then
            install_rhel
        else
            echo "Failed to auto-detect package manager. Please install wireguard manually."
            exit 1
        fi
        ;;
esac

# Check if wg command is installed successfully
if command -v wg >/dev/null 2>&1; then
    echo "========================================="
    echo "WireGuard has been successfully installed!"
    echo "WireGuard version: $(wg --version)"
    echo "You can configure WireGuard under /etc/wireguard/"
    echo "========================================="
else
    echo "Error: WireGuard installation failed or wg command not found." >&2
    exit 1
fi