#!/bin/bash

set -m # Enable Job Control

trap 'kill $(jobs -p)' SIGINT

# Reset
NC='\033[0m'              # Text Reset

# Regular Colors
Red='\033[0;31m'          # Red
Green='\033[0;32m'        # Green
Cyan='\033[0;36m'         # Cyan

if [[ -z $(which hcloud) ]]; then
    echo -e "${Red}hcloud could not be found in \$PATH!${NC}

Please put hcloud in \$PATH ($PATH),
install it with your package manager
or go to https://github.com/hetznercloud/cli/releases to download it."
    exit 1
fi

if [[ -z $(which jq) ]]; then
    echo -e "${Red}jq could not be found in \$PATH!${NC}

Please install jq to use this script."
    exit 1
fi

if [[ -z $(which nmap) ]]; then
    echo -e "${Red}nmap could not be found in \$PATH!${NC}

Please install nmap to use this script."
    exit 1
fi

usage() {
    if hcloud context list | grep -q -v "ACTIVE"; then
        types=$(hcloud server-type list -o columns=name,cores,cpu_type,memory,storage_type,architecture | grep -v arm | sed -e 's/^/    /')
        keys=$(hcloud ssh-key list -o columns=name,fingerprint,age | sed -e 's/^/    /')
        contexts="    Available contexts:
$(hcloud context list | sed -e 's/^/    /')"
    else
        types="No hcloud context, can’t get server types"
        keys="No hcloud context, can’t get SSH keys"
        contexts="No hcloud context available.
You can create one with the following command:
    hcloud create context name_of_the_context
Or let this script create one during execution."
    fi

    cat << EOF
$(basename "$0") (c) Framasoft 2023, WTFPL

USAGE
    $(basename "$0") [-h] [-d] [-s <int>] [-n <int>] [-t <vps type>] [-c <hcloud context>] -k <ssh-key>

OPTIONS
    -h                  Print this help and exit
    -d                  Delete all servers
    -dy                 Delete all servers without confirmation
    -s <int>            How many VPS you want to start.
                        Default: 1
                        Maximum should be: limit (hcloud).
                        Default: 1
    -n <int>            How many nodes you want to start on each VPS.
                        Default: 1
    -t <vps type>       The type of VPS to start.
                        Default: cpx21.
                        See below
    -c <hcloud context> Name of the hcloud context
                        Default: selenium-peertube.
                        See below
    -k <ssh-key>        The ssh key used to connect to the VPS.
                        MANDATORY, no default.Starting node
                        See below.
    -e <string>         The path to the environment file to be copied and used on the VPS.
                        Default: .env

$types

HCLOUD CONTEXT
    It’s the name of the project you want to create your VPS in.

$contexts

SSH KEYS
    You must have a ssh key registered on Hetzner to use this script.
    To create a key:
        hcloud ssh-key create --name my-key --public-key-from-file ~/.ssh/id_ed25519.pub

    The ssh keys currently registered on Hetzner are:
$keys
EOF
    exit "$1"
}

delete_server() {
    echo -e "${Cyan}$(hcloud server delete "$1")${NC}"
}

create_nodes_server() {
    i="$1"
    TYPE="$2"
    KEY="$3"
    REGION="$4"
    SERVER_NAME="$REGION-node-$i"
    hcloud server create --start-after-create --name "$SERVER_NAME" --image debian-12 --type "$TYPE" --location "$REGION" --ssh-key "$KEY" > /dev/null
    echo -e "${Cyan}VPS n°$i created and started${NC}"
}

start_nodes() {
    i="$1"
    REGION=$(hcloud server list -o json | jq -r '.[] | select(.name | contains("node-'$i'")) | .datacenter.location.name')
    SERVER_NAME="$REGION-node-$i"
    SERVER_IP=$(hcloud server ip "$SERVER_NAME")
    while [[ $(nmap -p 22 "$SERVER_IP" | grep -c open) -eq 0 ]]; do
        sleep 1
    done
    SSH_CONN="root@$SERVER_IP"
    scp -o "LogLevel=ERROR" -o "UserKnownHostsFile /dev/null" -o "StrictHostKeyChecking no" -o "VerifyHostKeyDNS no" start-nodes.sh "${SSH_CONN}:" > /dev/null
    scp -o "LogLevel=ERROR" -o "UserKnownHostsFile /dev/null" -o "StrictHostKeyChecking no" -o "VerifyHostKeyDNS no" "$ENV_FILE" "${SSH_CONN}:" > /dev/null
    ssh -o "LogLevel=ERROR" -o "UserKnownHostsFile /dev/null" -o "StrictHostKeyChecking no" -o "VerifyHostKeyDNS no" "$SSH_CONN" "/root/start-nodes.sh -n \"$NODES\"" > /dev/null
    echo -e "${Cyan}Nodes created on VPS n°${i}${NC}"
}

CONTEXT=selenium-peertube
SERVERS=1
NODES=1
TYPE=cpx21
DELETE=0
N_STRING=node
FORCE_DELETION=0
ENV_FILE=.env

while getopts "hds:n:t:k:c:y" option; do
    case $option in
        h)
            usage 0
            ;;
        d)
            DELETE=1
            ;;
        s)
            SERVERS=$OPTARG
            ;;
        n)
            NODES=$OPTARG
            if [[ $NODES -gt 1 ]]; then
                N_STRING=nodes
            fi
            ;;
        t)
            TYPE=$OPTARG
            ;;
        k)
            KEY=$OPTARG
            ;;
        c)
            CONTEXT=$OPTARG
            ;;
        y)
            FORCE_DELETION=1
            ;;
        *)
            usage 1
            ;;
    esac
done

if [[ $(hcloud context active) != "$CONTEXT" ]]; then
    echo -e "${Cyan}Hcloud context is not '$CONTEXT'!${NC}"
    if hcloud context list | grep -q -F "$CONTEXT"; then
        echo -e "${Green}Selecting hcloud context ${CONTEXT}${NC}"
        hcloud context use "$CONTEXT"
    else
        echo -e "${Red}Hcloud context ${CONTEXT} does not exist.${NC}
${Cyan}Will now try to create the context ${CONTEXT}${NC}"
        hcloud context create "$CONTEXT"
    fi

    exit 1
fi

if [[ $DELETE -eq 1 ]]; then
    SERVERS=$(hcloud server list -o json)
    if [[ $SERVERS == 'null' ]]; then
        echo -e "${Cyan}No VPS to delete.${NC}"
        exit 0
    fi
    NAMES=$(echo "$SERVERS" | jq -r '.[] | .name' | sort -h)
    echo -e "${Red}You are about to delete the following VPS${NC}:"
    echo "$NAMES"
    if [[ $FORCE_DELETION -eq 1 ]]; then
        confirm="yes"
    else
        echo -e -n "${Cyan}Please confirm the deletion by typing '${NC}${Red}yes${NC}': "
        read -r confirm
    fi
    if [[ $confirm == 'yes' ]]; then
        for i in $NAMES; do
            echo -e "${Cyan}Starting server $i deletion${NC}"
            delete_server "$i" &
        done
        # Wait for all delete_server jobs to finish
        while true; do
            fg > /dev/null 2>&1
            [ $? == 1 ] && break
        done
        if [[ $(hcloud server list -o json) == '[]' ]]; then
            echo -e "${Green}All servers have been deleted${NC}"
        else
            echo -e "${Red}Some servers have not been deleted:${NC}"
            hcloud server list
        fi
    else
        echo "Deletion cancelled."
    fi
    exit 0
fi

if [[ -z $KEY ]]; then
    echo -e "${Red}You must choose a ssh key!${NC}\n"
    usage 1
fi

KEY_FOUND=0
for i in $(hcloud ssh-key list -o json | jq -r '.[] | .name'); do
    if [[ $i == "$KEY" ]]; then
        KEY_FOUND=1
        break
    fi
done

if [[ $KEY_FOUND -eq 0 ]]; then
    echo -e "${Red}The chosen ssh key is not registered on Hetzner!${NC}\n"
    usage 1
fi

if hcloud server list | grep -q -v NAME; then
    echo -e "${Red}There already are servers in the context! Exiting.${NC}\nList of the servers:"
    hcloud server list
    exit 1
fi

if [[ ! -f "$ENV_FILE" ]]; then
    echo -e "${Red}Environment file '$ENV_FILE' does not exist!${NC}"
    exit 1
fi

echo -e "${Green}Creating $SERVERS VPS${NC}"
REGIONS=($(hcloud location list -o json | jq -r '.[] | select(.name != "fsn1") | .name' | shuf))
for i in $(seq 1 "$SERVERS"); do
    REGION=${REGIONS[$((i % ${#REGIONS[@]}))]}
    echo -e "${Cyan}Creating VPS n°$i in $REGION"
    create_nodes_server "$i" "$TYPE" "$KEY" "$REGION" &
done

# Wait for all create_nodes_server jobs to finish
while true; do
    fg > /dev/null 2>&1
    [ $? == 1 ] && break
done

echo -e "${Green}Starting nodes on $SERVERS VPS ($NODES $N_STRING each)${NC}"
for i in $(seq 1 "$SERVERS"); do
    echo -e "${Cyan}Starting $N_STRING on VPS n°$i${NC}"
    start_nodes "$i" &
done

echo -e "${Green}Waiting for all nodes to be started${NC}"

# Wait for all start_nodes jobs to finish
while true; do
    fg > /dev/null 2>&1
    [ $? == 1 ] && break
done

echo -e "${Green}All the servers and nodes have been created and started!

Number of servers: $SERVERS
Number of nodes per server: $NODES
Type of the servers:
  nodes servers: $TYPE

You can remove all servers with the following command
    $0 -d${NC}"