#!/bin/bash

update_system() {
    echo "Updating system and installing mandatory tools"
    apt-get update
    apt-get install sudo nmon tmux restic tcpdump nano iputils-ping net-tools -y
}

install_docker() {
    echo "Installing Docker"
    apt-get install ca-certificates curl gnupg lsb-release -y
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
    apt-get update
    apt-get install docker-ce docker-ce-cli containerd.io -y
    systemctl start docker
    echo "Docker installation completed"
}

install_caddy() {
    echo "Installing Caddy"
    apt install -y debian-keyring debian-archive-keyring apt-transport-https
    curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/gpg.key' | sudo gpg --dearmor -o /usr/share/keyrings/caddy-stable-archive-keyring.gpg
    curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/debian.deb.txt' | sudo tee /etc/apt/sources.list.d/caddy-stable.list
    apt-get update
    apt-get install caddy -y
    echo "Caddy installation completed."
}

setup_tf_users() {
    echo "Setting up TFUsers"
    wget https://docs.grid.tf/threefold_public/public/raw/branch/master/add-tf-users.sh
    sh add-tf-users.sh
    rm add-tf-users.sh
}

configure_ssh() {
    echo "Configuring SSH"
    ubuntu_version=$(lsb_release -rs)

    if [[ "$ubuntu_version" == "22.04" ]]; then
        echo "Detected Ubuntu 22.04 — replacing sshd_config and restarting ssh service"
        wget -q https://docs.grid.tf/threefold_public/public/raw/branch/master/sshd_config -O /etc/ssh/sshd_config
        systemctl restart ssh

    elif [[ "$ubuntu_version" == "24.04" ]]; then
        echo "Detected Ubuntu 24.04 — updating ssh.socket for port 34022"
        wget -q https://docs.grid.tf/threefold_public/public/raw/branch/master/sshd_config -O /etc/ssh/sshd_config
        mkdir -p /etc/systemd/system/ssh.socket.d
        cat > /etc/systemd/system/ssh.socket.d/port.conf <<EOF
[Socket]
ListenStream=
ListenStream=34022
EOF

        systemctl daemon-reload
        systemctl restart ssh.socket
        systemctl enable ssh.socket
    else
        echo "Unsupported Ubuntu version: $ubuntu_version"
        return 1
    fi
}


update_system
setup_tf_users
configure_ssh

while getopts ":dc" opt; do
  case ${opt} in
    d )
      install_docker
      ;;
    c )
      install_caddy
      ;;
    \? )
      echo "Invalid option: $OPTARG" 1>&2
      ;;
  esac
done
shift $((OPTIND -1))

echo "Preping VM Completed."