#!/bin/bash

#==========================
# Set up the environment
#==========================
set -e                  # exit on error
set -o pipefail         # exit on pipeline error
set -u                  # treat unset variable as error

#==========================
# Basic Information
#==========================
export LC_ALL=C
export LANG=en_US.UTF-8
export DEBIAN_FRONTEND=noninteractive
export SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"

#==========================
# Color
#==========================
Green="\033[32m"
Red="\033[31m"
Yellow="\033[33m"
Blue="\033[36m"
Font="\033[0m"
GreenBG="\033[42;37m"
RedBG="\033[41;37m"
OK="${Green}[  OK  ]${Font}"
ERROR="${Red}[FAILED]${Font}"
WARNING="${Yellow}[ WARN ]${Font}"

#==========================
# Print Colorful Text
#==========================
function print_ok() {
  echo -e "${OK} ${Blue} $1 ${Font}"
}

function print_error() {
  echo -e "${ERROR} ${Red} $1 ${Font}"
}

function print_warn() {
  echo -e "${WARNING} ${Yellow} $1 ${Font}"
}

#==========================
# Judge function
#==========================
function judge() {
  if [[ 0 -eq $? ]]; then
    print_ok "$1 succeeded"
    sleep 0.2
  else
    print_error "$1 failed"
    exit 1
  fi
}

prepare_host()
{
    print_ok "Update apt-get"
    sudo apt-get update
    judge "Update apt-get"

    print_ok "Install sshpass"
    sudo apt-get install -y sshpass
    judge "Install sshpass"
}

wait_server_till_can_ssh()
{
    userName=$1
    password=$2
    serverName=$3

    print_ok "Waiting for server to be ready: ssh $userName@$serverName"
    while true; do
        set +e
        sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "echo 'Server is ready'"
        if [ $? -eq 0 ]; then
            break
        fi
        print_warn "Server is not ready yet. Retrying..."
        sleep 5
    done

    print_ok "Server is ready to connect via ssh"
    set -e
}

prepare_server()
{
    userName=$1
    if [ -z "$userName" ]; then
        print_error "Please provide username. Usage: prepare_server.sh '<username>' '<password>' '<serverName>' '<desiredHostname>' '<desiredUsername>'"
        exit 1
    fi

    password=$2
    if [ -z "$password" ]; then
        print_error "Please provide password for user $userName. Usage: prepare_server.sh '<username>' '<password>' '<serverName>' '<desiredHostname>' '<desiredUsername>'"
        exit 1
    fi

    serverName=$3
    if [ -z "$serverName" ]; then
        print_error "Please provide server name. Usage: prepare_server.sh '<username>' '<password>' '<serverName>' '<desiredHostname>' '<desiredUsername>'"
        exit 1
    fi

    desiredHostname=$4
    if [ -z "$desiredHostname" ]; then
        echo "Please provide desired hostname. Usage: prepare_server.sh '<username>' '<password>' '<serverName>' '<desiredHostname>' '<desiredUsername>'"
        exit 1
    fi

    desiredUsername=$5
    if [ -z "$desiredUsername" ]; then
        print_error "Please provide desired username. Usage: prepare_server.sh '<username>' '<password>' '<serverName>' '<desiredHostname>' '<desiredUsername>'"
        exit 1
    fi

    prepare_host
    ssh-keygen -f "/home/anduin/.ssh/known_hosts" -R $serverName

    wait_server_till_can_ssh $userName $password $serverName

    print_ok "Ensure Server is Ubuntu 22.04" # Accept 22.04.1, 22.04.2, etc
    osVersion=$(sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "lsb_release -r | awk '{print \$2}'")
    if [ "$osVersion" != "22.04" ]; then
        print_error "Server is not Ubuntu 22.04. Please use Ubuntu 22.04"
        exit 1
    fi

    print_ok "Changing hostname for $serverName to $desiredHostname"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sudo hostnamectl set-hostname $desiredHostname"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sleep 3"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sudo reboot" || true
    sleep 5
    print_ok "Hostname changed to $desiredHostname"
    print_warn "Server is rebooting..."

    wait_server_till_can_ssh $userName $password $serverName

    print_ok "Creating a new user..."
    alreadyExist=$(sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "cat /etc/passwd | grep -w $desiredUsername | wc -l")
    if [ $alreadyExist -gt 0 ]; then
        print_ok "User $desiredUsername already exists."
    else
        print_ok "Creating user $desiredUsername"
        sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sudo adduser $desiredUsername --gecos 'First Last,RoomNumber,WorkPhone,HomePhone' --disabled-password"
    fi
    judge "User $desiredUsername created"

    print_ok "Setting password for user $desiredUsername"
    userPassword=$(uuidgen)
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "echo $desiredUsername:$userPassword | sudo chpasswd"
    judge "Password set for user $desiredUsername"

    print_ok "Adding user $desiredUsername to sudo group"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sudo usermod -aG sudo $desiredUsername"
    judge "User $desiredUsername created with password $userPassword"

    print_ok "Allowing user $desiredUsername to run sudo without password"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sudo mkdir -p /etc/sudoers.d"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "sudo touch /etc/sudoers.d/$desiredUsername"
    sshpass -p $password ssh -o StrictHostKeyChecking=no $userName@$serverName "echo '$desiredUsername ALL=(ALL) NOPASSWD:ALL' | sudo tee -a /etc/sudoers.d/$desiredUsername"
    judge "User $desiredUsername can run sudo without password"

    # If ~/ssh/id_rsa.pub does not exist, create it
    if [ ! -f ~/.ssh/id_rsa.pub ]; then
        print_warn "Creating ssh keys on local machine"
        ssh-keygen -t rsa -N "" -f ~/.ssh/id_rsa
    fi

    print_ok "Copying ssh keys with ssh-copy-id"
    sshpass -p $userPassword ssh-copy-id -i ~/.ssh/id_rsa.pub $desiredUsername@$serverName
    print_ok "SSH keys copied"

    wait_server_till_can_ssh $desiredUsername $userPassword $serverName

    print_ok "Disabling root login, password login and enabling ssh key login"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sed -i 's/PermitRootLogin yes/PermitRootLogin no/g' /etc/ssh/sshd_config"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sed -i 's/PasswordAuthentication yes/PasswordAuthentication no/g' /etc/ssh/sshd_config"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sed -i 's/PubkeyAuthentication no/PubkeyAuthentication yes/g' /etc/ssh/sshd_config"
    # Uncomment those lines if they are commented
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sed -i 's/#PermitRootLogin no/PermitRootLogin no/g' /etc/ssh/sshd_config"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sed -i 's/#PasswordAuthentication no/PasswordAuthentication no/g' /etc/ssh/sshd_config"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sed -i 's/#PubkeyAuthentication yes/PubkeyAuthentication yes/g' /etc/ssh/sshd_config"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo systemctl restart sshd"
    judge "Disable root login, password login and enabled ssh key login"

    print_ok "Server is ready for $desiredUsername to login. Deleting other users..."
    otherUsers=$(ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "cat /etc/passwd | grep -v nologin | grep -v false | grep -v root | grep -v sync | grep -v $desiredUsername | cut -d: -f1")
    for otherUser in $otherUsers; do
        print_warn "Deleting user $otherUser..."
        ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "sudo pkill -u $otherUser" || true
        ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "sudo deluser --remove-home $otherUser"
    done

    print_ok "Resetting machine-id"
    ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "sudo rm /etc/machine-id"
    ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "sudo rm /var/lib/dbus/machine-id"
    ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "sudo systemd-machine-id-setup"
    ssh -o StrictHostKeyChecking=no  $desiredUsername@$serverName "sudo cp /etc/machine-id /var/lib/dbus/machine-id"
    judge "Machine-id reset"

    print_ok "Enabling ufw firewall"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt-get install -y ufw"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo ufw allow OpenSSH"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "echo 'y' | sudo ufw enable"
    judge "Ufw firewall enabled"

    print_ok "Enabling BBR if not enabled"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sysctl net.ipv4.tcp_available_congestion_control | grep -q bbr || echo 'net.core.default_qdisc=fq' | sudo tee -a /etc/sysctl.conf"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sysctl net.ipv4.tcp_available_congestion_control | grep -q bbr || echo 'net.ipv4.tcp_congestion_control=bbr' | sudo tee -a /etc/sysctl.conf"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo sysctl -p"
    judge "BBR enabled"

    print_ok "Selecting best mirror"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "curl -s https://gist.aiursoft.cn/anduin/879917820a6c4b268fc12c21f1b3fe7a/raw/HEAD/mirror.sh | bash"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt update"
    judge "Best mirror selected"

    print_ok "Installing latest kernel..."
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt install -y linux-generic-hwe-22.04"
    judge "Latest kernel installed"

    print_ok "Rebooting server..."
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sleep 3"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo reboot" || true
    sleep 5
    print_warn "Server is rebooting..."

    wait_server_till_can_ssh $desiredUsername $userPassword $serverName

    print_ok "Installing updates"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt update"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt upgrade -y"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt autoremove -y"
    judge "Updates installed"

    print_ok "Rebooting server..."
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sleep 3"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo reboot" || true
    sleep 5
    print_warn "Server is rebooting..."

    wait_server_till_can_ssh $desiredUsername $userPassword $serverName

    print_ok "Set CPU to performance mode"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt install -y linux-tools-common linux-tools-$(uname -r)"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo cpupower frequency-info"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo cpupower frequency-set -g performance" || true
    judge "CPU set to performance mode"

    print_ok "Set timezone to GMT"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo timedatectl set-timezone GMT"
    judge "Timezone set to GMT"

    print_ok "Removing snap..."
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo systemctl disable --now snapd"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt purge -y snapd"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo rm -rf /snap /var/snap /var/lib/snapd /var/cache/snapd /usr/lib/snapd ~/snap"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "cat << EOF | sudo tee -a /etc/apt/preferences.d/no-snap.pref
Package: snapd
Pin: release a=*
Pin-Priority: -10
EOF"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo chown root:root /etc/apt/preferences.d/no-snap.pref"
    judge "Snap removed"

    print_ok "Autoremoving apt packages"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt autoremove -y --purge"
    judge "Apt packages autoremoved"

    print_ok "Benchmarking server..."
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sudo apt install -y sysbench"
    ssh -o StrictHostKeyChecking=no $desiredUsername@$serverName "sysbench cpu --threads=\$(nproc) run"
    judge "Server benchmarked"

    print_ok "Server is ready for use"
    print_ok "ssh $desiredUsername@$serverName"
}

# To use this function:
# Arg1: username
# Arg2: password
# Arg3: servername
# Arg4: Desired hostname
# Arg5: Desired username
prepare_server "$@"