#!/bin/bash
set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Configuration
WG_INTERFACE="wg-cluster"
WG_PORT=51820
WG_NETWORK="10.10.0.0/24"
GLUSTER_BRICK_PATH="/gluster/cluster"
GLUSTER_MOUNT_PATH="/data/storage"
GLUSTER_VOLUME="cluster-volume"

echo -e "${GREEN}================================${NC}"
echo -e "${GREEN}    Nullpoint Cluster Setup${NC}"
echo -e "${GREEN}================================${NC}\n"

# Check if running as root
if [ "$EUID" -ne 0 ]; then 
    echo -e "${RED}Please run as root${NC}"
    exit 1
fi

# Install required packages
echo -e "${YELLOW}[+] Installing required packages...${NC}"
dnf install -y wireguard-tools glusterfs-server glusterfs-client || exit 1

# Enable and start GlusterFS
systemctl enable glusterd
systemctl start glusterd

# Create directories
echo -e "${YELLOW}[+] Creating directories...${NC}"
mkdir -p "$GLUSTER_BRICK_PATH"
mkdir -p "$GLUSTER_MOUNT_PATH"
mkdir -p /data

# Function to generate WireGuard keys
generate_wg_keys() {
    local private_key=$(wg genkey)
    local public_key=$(echo "$private_key" | wg pubkey)
    local preshared_key=$(wg genpsk)
    echo "$private_key:$public_key:$preshared_key"
}

# Function to get next available IP
get_next_ip() {
    local base_ip="10.10.0"
    local next_num=2
    
    if [ -f /etc/wireguard/${WG_INTERFACE}.conf ]; then
        # Find highest IP in use
        existing_ips=$(grep -E "AllowedIPs|Address" /etc/wireguard/${WG_INTERFACE}.conf | grep -oE "10\.10\.0\.[0-9]+" | cut -d. -f4 | sort -n | tail -1)
        if [ ! -z "$existing_ips" ]; then
            next_num=$((existing_ips + 1))
        fi
    fi
    
    echo "${base_ip}.${next_num}"
}

# Function to setup firewall rules
setup_firewall() {
    echo -e "${YELLOW}[+] Configuring firewall...${NC}"
    
    # WireGuard
    firewall-cmd --permanent --add-port=${WG_PORT}/udp
    
    # GlusterFS ports
    firewall-cmd --permanent --add-service=glusterfs
    firewall-cmd --permanent --add-port=24007-24008/tcp  # GlusterFS Daemon
    firewall-cmd --permanent --add-port=49152-49200/tcp  # Brick ports
    
    # Allow traffic from WireGuard network
    firewall-cmd --permanent --zone=trusted --add-source=${WG_NETWORK}
    
    firewall-cmd --reload
}

# Create new cluster
create_cluster() {
    echo -e "${GREEN}[*] Creating new cluster...${NC}\n"
    
    # Generate keys
    keys=$(generate_wg_keys)
    private_key=$(echo "$keys" | cut -d: -f1)
    public_key=$(echo "$keys" | cut -d: -f2)
    preshared_key=$(echo "$keys" | cut -d: -f3)
    
    # Create WireGuard config
    cat > /etc/wireguard/${WG_INTERFACE}.conf <<EOF
[Interface]
Address = 10.10.0.1/24
ListenPort = ${WG_PORT}
PrivateKey = ${private_key}
SaveConfig = false

# Peers will be added here
EOF
    
    chmod 600 /etc/wireguard/${WG_INTERFACE}.conf
    
    # Start WireGuard
    systemctl enable wg-quick@${WG_INTERFACE}
    systemctl start wg-quick@${WG_INTERFACE}
    
    # Setup firewall
    setup_firewall
    
    # Create GlusterFS volume (single brick for now)
    echo -e "${YELLOW}[+] Creating GlusterFS volume...${NC}"
    
    # Create brick directory
    mkdir -p "${GLUSTER_BRICK_PATH}/brick1"
    
    # Create volume with single brick (will be converted to replica when nodes join)
    gluster volume create ${GLUSTER_VOLUME} $(hostname):${GLUSTER_BRICK_PATH}/brick1 force 2>/dev/null || true
    gluster volume start ${GLUSTER_VOLUME} 2>/dev/null || true
    
    # Mount the volume locally
    mount -t glusterfs localhost:/${GLUSTER_VOLUME} ${GLUSTER_MOUNT_PATH}
    
    # Add to fstab for persistence
    grep -q "${GLUSTER_VOLUME}" /etc/fstab || echo "localhost:/${GLUSTER_VOLUME} ${GLUSTER_MOUNT_PATH} glusterfs defaults,_netdev 0 0" >> /etc/fstab
    
    # Get external IP for other nodes to connect
    external_ip=$(ip route get 1.1.1.1 | awk '{print $7; exit}')
    
    echo -e "\n${GREEN}════════════════════════════════════════${NC}"
    echo -e "${GREEN}Cluster created successfully!${NC}"
    echo -e "${GREEN}════════════════════════════════════════${NC}\n"
    echo -e "Share these details with nodes joining the cluster:\n"
    echo -e "${YELLOW}Cluster IP:${NC} ${external_ip}"
    echo -e "${YELLOW}WireGuard Port:${NC} ${WG_PORT}"
    echo -e "${YELLOW}Public Key:${NC} ${public_key}"
    echo -e "${YELLOW}Preshared Key:${NC} ${preshared_key}"
    echo -e "\n${YELLOW}Cluster Secret (for easy sharing):${NC}"
    echo -e "${GREEN}${external_ip}:${WG_PORT}:${public_key}:${preshared_key}${NC}"
    echo -e "\n${YELLOW}Status:${NC}"
    echo "  - WireGuard interface: ${WG_INTERFACE} (10.10.0.1)"
    echo "  - GlusterFS volume: ${GLUSTER_VOLUME}"
    echo "  - Mount point: ${GLUSTER_MOUNT_PATH}"
}

# Join existing cluster
join_cluster() {
    echo -e "${GREEN}[*] Joining existing cluster...${NC}\n"
    
    # Get cluster details
    read -p "Enter cluster node IP: " cluster_ip
    read -p "Enter cluster secret (or press enter to input separately): " cluster_secret
    
    if [ -z "$cluster_secret" ]; then
        read -p "Enter WireGuard port [${WG_PORT}]: " wg_port_input
        wg_port=${wg_port_input:-$WG_PORT}
        read -p "Enter cluster public key: " cluster_public_key
        read -p "Enter preshared key: " preshared_key
    else
        # Parse secret
        cluster_ip=$(echo "$cluster_secret" | cut -d: -f1)
        wg_port=$(echo "$cluster_secret" | cut -d: -f2)
        cluster_public_key=$(echo "$cluster_secret" | cut -d: -f3)
        preshared_key=$(echo "$cluster_secret" | cut -d: -f4)
    fi
    
    # Generate local keys
    keys=$(generate_wg_keys)
    private_key=$(echo "$keys" | cut -d: -f1)
    public_key=$(echo "$keys" | cut -d: -f2)
    
    # Get next available IP
    my_ip=$(get_next_ip)
    
    echo -e "${YELLOW}[+] Configuring WireGuard (IP: ${my_ip})...${NC}"
    
    # Create WireGuard config
    cat > /etc/wireguard/${WG_INTERFACE}.conf <<EOF
[Interface]
Address = ${my_ip}/24
ListenPort = ${WG_PORT}
PrivateKey = ${private_key}
SaveConfig = false

[Peer]
PublicKey = ${cluster_public_key}
PresharedKey = ${preshared_key}
Endpoint = ${cluster_ip}:${wg_port}
AllowedIPs = ${WG_NETWORK}
PersistentKeepalive = 25
EOF
    
    chmod 600 /etc/wireguard/${WG_INTERFACE}.conf
    
    # Start WireGuard
    systemctl enable wg-quick@${WG_INTERFACE}
    systemctl restart wg-quick@${WG_INTERFACE}
    
    # Setup firewall
    setup_firewall
    
    # Wait for WireGuard connection
    echo -e "${YELLOW}[+] Waiting for WireGuard connection...${NC}"
    sleep 3
    
    # Test connection to cluster
    if ! ping -c 1 -W 2 10.10.0.1 > /dev/null 2>&1; then
        echo -e "${RED}Failed to connect to cluster via WireGuard!${NC}"
        echo "Please check the cluster details and firewall settings."
        exit 1
    fi
    
    echo -e "${GREEN}[✓] WireGuard connection established${NC}"
    
    # Add this node to the cluster node's WireGuard config
    echo -e "${YELLOW}[+] Requesting cluster to add this node as peer...${NC}"
    
    # SSH to cluster node and add peer (requires SSH key setup)
    ssh_cmd="wg set ${WG_INTERFACE} peer ${public_key} preshared-key <(echo ${preshared_key}) allowed-ips ${WG_NETWORK} persistent-keepalive 25"
    
    echo -e "${YELLOW}Run this command on the cluster node (10.10.0.1) to add this peer:${NC}"
    echo -e "${GREEN}sudo wg set ${WG_INTERFACE} peer ${public_key} preshared-key <(echo ${preshared_key}) allowed-ips ${WG_NETWORK} persistent-keepalive 25${NC}"
    echo -e "${GREEN}sudo bash -c 'echo \"[Peer]\" >> /etc/wireguard/${WG_INTERFACE}.conf'${NC}"
    echo -e "${GREEN}sudo bash -c 'echo \"PublicKey = ${public_key}\" >> /etc/wireguard/${WG_INTERFACE}.conf'${NC}"
    echo -e "${GREEN}sudo bash -c 'echo \"PresharedKey = ${preshared_key}\" >> /etc/wireguard/${WG_INTERFACE}.conf'${NC}"
    echo -e "${GREEN}sudo bash -c 'echo \"AllowedIPs = ${WG_NETWORK}\" >> /etc/wireguard/${WG_INTERFACE}.conf'${NC}"
    echo -e "${GREEN}sudo bash -c 'echo \"PersistentKeepalive = 25\" >> /etc/wireguard/${WG_INTERFACE}.conf'${NC}"
    
    read -p "Press enter once you've added this peer to the cluster node..." 
    
    # Join GlusterFS cluster
    echo -e "${YELLOW}[+] Joining GlusterFS cluster...${NC}"
    
    # Probe the cluster
    gluster peer probe 10.10.0.1
    
    # Wait for peer to be connected
    sleep 3
    
    # Create brick directory
    mkdir -p "${GLUSTER_BRICK_PATH}/brick1"
    
    # Add brick to existing volume and increase replica count
    echo -e "${YELLOW}[+] Adding brick to GlusterFS volume...${NC}"
    
    # Get current replica count
    replica_count=$(gluster volume info ${GLUSTER_VOLUME} 2>/dev/null | grep "Number of Bricks" | grep -oE "[0-9]+" | head -1)
    new_replica_count=$((replica_count + 1))
    
    # Add brick with increased replica count
    gluster volume add-brick ${GLUSTER_VOLUME} replica ${new_replica_count} $(hostname):${GLUSTER_BRICK_PATH}/brick1 force
    
    # Mount the volume
    mount -t glusterfs 10.10.0.1:/${GLUSTER_VOLUME} ${GLUSTER_MOUNT_PATH}
    
    # Add to fstab
    grep -q "${GLUSTER_VOLUME}" /etc/fstab || echo "10.10.0.1:/${GLUSTER_VOLUME} ${GLUSTER_MOUNT_PATH} glusterfs defaults,_netdev 0 0" >> /etc/fstab
    
    echo -e "\n${GREEN}════════════════════════════════════════${NC}"
    echo -e "${GREEN}Successfully joined cluster!${NC}"
    echo -e "${GREEN}════════════════════════════════════════${NC}\n"
    echo -e "${YELLOW}Node details:${NC}"
    echo "  - WireGuard IP: ${my_ip}"
    echo "  - Public Key: ${public_key}"
    echo "  - GlusterFS mounted at: ${GLUSTER_MOUNT_PATH}"
}

# Show cluster status
show_status() {
    echo -e "\n${YELLOW}=== Cluster Status ===${NC}\n"
    
    if [ -f /etc/wireguard/${WG_INTERFACE}.conf ]; then
        echo -e "${GREEN}WireGuard Status:${NC}"
        wg show ${WG_INTERFACE}
        echo ""
    else
        echo -e "${RED}WireGuard not configured${NC}\n"
    fi
    
    echo -e "${GREEN}GlusterFS Status:${NC}"
    gluster peer status
    echo ""
    gluster volume status ${GLUSTER_VOLUME} 2>/dev/null || echo "Volume ${GLUSTER_VOLUME} not found"
    echo ""
    
    echo -e "${GREEN}Mounted at:${NC} ${GLUSTER_MOUNT_PATH}"
    df -h ${GLUSTER_MOUNT_PATH} 2>/dev/null || echo "Not mounted"
}

# Main menu
echo "What would you like to do?"
echo "  1) Create new cluster"
echo "  2) Join existing cluster"
echo "  3) Show cluster status"
echo "  4) Exit"
echo ""
read -p "Enter choice [1-4]: " choice

case $choice in
    1)
        create_cluster
        ;;
    2)
        join_cluster
        ;;
    3)
        show_status
        ;;
    4)
        echo "Exiting..."
        exit 0
        ;;
    *)
        echo -e "${RED}Invalid choice${NC}"
        exit 1
        ;;
esac