#!/bin/bash

#
# Default variables
#

upgrade_packages=false
script_dir="$(
    cd -- "$(dirname "$0")" >/dev/null 2>&1 || exit 1
    pwd -P
)"

#
# Argument parsing and help
#

function usage() {
    echo "Usage: $(basename -- "$0") [OPTIONS]"
    echo
    echo "Options:"
    echo "  -h      Display help."
    echo "  -u      Upgrade existing packages."
    exit 0
}

while getopts ":uh" arg; do
    case $arg in
    h) usage ;;
    u) upgrade_packages=true ;;
    :)
        echo "$0: Must supply an argument to -$OPTARG." >&2
        exit 1
        ;;
    ?)
        echo "Invalid option: -${OPTARG}."
        echo
        usage
        exit 2
        ;;
    esac
done

#
# Distribution-agnostic functions
#

function _acpib() {
    acpi -b | grep -Fv ' 0%'
}

function has_battery() {
    [[ $(_acpib | wc -l) -gt 0 ]] && return 0 || return 1
}

# Returns a list of packages for one or more targets.
# The packages file may contain multiple rows with the same target name.
function get_packages() {
    local pkg_file="$script_dir/.installer/packages"

    for target in "$@"; do
        while read -r row; do
            for package in $row; do
                echo "$package"
            done
        done <<<"$(grep "^${target:?}:" "$pkg_file" | cut -d: -f2 | cut -d# -f1)"
    done
}

# Return a list of targets for a given package manager
function construct_target_list() {
    local package_manager=${1:?package_manager not set} targets=()

    targets+=("$package_manager-any")
    targets+=("$package_manager-host-$(hostname | cut -d. -f1)")
    has_battery && targets+=("$package_manager-type-hasbattery")
    echo "${targets[@],,}" # Return lowercase
}

function setup_pip_packages() {
    # shellcheck disable=SC2046
    sudo pip3 install $(get_packages $(construct_target_list pip)) --no-input || return 1
}

function install_sddm_aerial_theme() {
    local theme_dir="/usr/share/sddm/themes/aerial"

    [ -d "$theme_dir" ] && return

    echo "Installing SDDM Aerial theme."

    git clone https://github.com/3ximus/aerial-sddm-theme.git /tmp/aerial
    sudo mv /tmp/aerial "$theme_dir"
    sudo chown -R root:root "$theme_dir"
}

function configure_sddm() {
    local sddm_config_dir="/etc/sddm.conf.d"
    local theme_dir="/usr/share/sddm/themes/aerial"
    local theme_user_config="$theme_dir/theme.conf.user"

    echo "Setting up '$sddm_config_dir'."

    if ! [ -d "$sddm_config_dir" ]; then
        sudo mkdir "$sddm_config_dir"
        sudo chown root:root -R "$sddm_config_dir"
    fi

    {
        echo "[Theme]"
        echo "Current=aerial"
    } | sudo tee "$sddm_config_dir/theme" >/dev/null

    {
        echo "[X11]"
        echo "DisplayCommand=/usr/share/sddm/scripts/Xsetup"
    } | sudo tee "$sddm_config_dir/xsetup" >/dev/null

    echo "Setting up '$theme_dir'."
    [ -f "$theme_user_config" ] && sudo mv "$theme_user_config" "$theme_user_config.disabled"
    return 0
}

function configure_ufw() {
    echo "Setting up UFW."

    sudo systemctl enable ufw.service || return 1
    sudo systemctl start ufw.service || return 2
    sudo ufw enable || return 3
    sudo ufw allow ssh || return 4
}

function install_picom() {
    local tmp=/tmp/picom

    echo "Installing Picom."

    rm -Rf "$tmp"
    git clone https://github.com/yshui/picom.git "$tmp" || return 1
    cd "$tmp" || return 2
    git submodule update --init --recursive || return 3
    meson --buildtype=release . build || return 4
    ninja -C build || return 5
    sudo ninja -C build install || return 6
    cd || return 7
    rm -Rf "$tmp"
}

function setup_homebrew() {
    local dir="$HOME/homebrew"

    echo "Setting up Homebrew."

    if [ ! -d "$dir" ]; then
        echo "Homebrew is already installed."
        mkdir "$dir" || return 1
        curl -L https://github.com/Homebrew/brew/tarball/master | tar xz --strip 1 -C "$dir" || return 2
    fi

    echo "Updating Homebrew."
    eval "$("$dir/bin/brew" shellenv)"
    brew update --force
    chmod -R go-w "$(brew --prefix)/share/zsh"
}

function setup_brew_formulas() {
    brew update || return 1

    # shellcheck disable=SC2046
    for package in $(get_packages $(construct_target_list brew)); do
        if brew list --full-name | grep -Eq "(^| )$package($| )" >/dev/null; then
            if $upgrade_packages; then
                echo "Package '$package' will be upgraded."
                brew upgrade "$package" || return 2
            else
                echo "Package '$package' already exists."
                continue
            fi
        else
            echo "Package '$package' will be installed."
            brew install "$package" || return 2
        fi
    done
}

function setup_krew() {
    cd "$(mktemp -d)" &&
        OS="$(uname | tr '[:upper:]' '[:lower:]')" &&
        ARCH="$(uname -m | sed -e 's/x86_64/amd64/' -e 's/\(arm\)\(64\)\?.*/\1\2/' -e 's/aarch64$/arm64/')" &&
        KREW="krew-${OS}_${ARCH}" &&
        curl -fsSLO "https://github.com/kubernetes-sigs/krew/releases/latest/download/${KREW}.tar.gz" &&
        tar zxvf "${KREW}.tar.gz" &&
        ./"${KREW}" install krew
    export PATH="${KREW_ROOT:-$HOME/.krew}/bin:$PATH"
}

function setup_krew_plugins() {
    local targets=("krew-any")

    if ! command -v kubectl-krew >/dev/null; then
        echo "WARN: Krew not installed: kubectl-krew"
        return 1
    fi

    # shellcheck disable=SC2046
    for package in $(get_packages $(construct_target_list krew)); do
        if kubectl-krew list | grep -q "^$package$" >/dev/null; then
            if $upgrade_packages; then
                echo "Package '$package' will be upgraded."
            else
                echo "Package '$package' already exists."
                continue
            fi
        else
            echo "Package '$package' will be installed."
        fi

        kubectl-krew install "$package" || return 2
    done
}

function build_alacritty() {
    local tmp_dir="/tmp/alacritty"

    echo "Installing Alacritty."
    git clone https://github.com/alacritty/alacritty.git $tmp_dir
    sudo cargo build --release --manifest-path=$tmp_dir/Cargo.toml
    sudo cp $tmp_dir/target/release/alacritty /usr/local/bin
    sudo cp $tmp_dir/extra/logo/alacritty-term.svg /usr/share/pixmaps/Alacritty.svg
    sudo desktop-file-install $tmp_dir/extra/linux/Alacritty.desktop
    sudo update-desktop-database
    sudo rm -Rf /tmp/alacritty
}

function change_shell() {
    sudo chsh -s "$(which zsh)" "$USER"
}

#
# Setup for Arch-like systems
#

function prepare_arch() {
    which pamac >/dev/null && return 0
    sudo pacman -S pamac --noconfirm || return 1
}

function setup_arch_with_pamac() {
    # shellcheck disable=SC2046
    sudo pamac install $(get_packages $(construct_target_list pacman)) --no-confirm || return 1
}

function setup_arch_with_yay() {
    # shellcheck disable=SC2046
    for package in $(get_packages $(construct_target_list aur)); do
        if pacman -Qs "^$package$" >/dev/null; then
            if $upgrade_packages; then
                echo "Package '$package' will be upgraded."
            else
                echo "Package '$package' already exists."
                continue
            fi
        else
            echo "Package '$package' will be installed."
        fi

        sudo pamac build "$package" --no-confirm || return 1
    done
}

function setup_arch() {
    prepare_arch || return 1
    setup_arch_with_pamac || return 2
    setup_arch_with_yay || return 3
    setup_pip_packages || return 4
    configure_sddm || return 5
    configure_ufw || return 6
    setup_homebrew || return 7
    setup_brew_formulas || return 8
    setup_krew || return 9
    setup_krew_plugins || return 10
    change_shell || return 11
}

#
# Setup for Debian-like systems
#

function setup_debian_repo_vscodium() {
    echo "Setting up repository for VSCodium."
    wget -qO - https://gitlab.com/paulcarroty/vscodium-deb-rpm-repo/raw/master/pub.gpg |
        gpg --dearmor |
        sudo dd of=/usr/share/keyrings/vscodium-archive-keyring.gpg
    echo 'deb [ signed-by=/usr/share/keyrings/vscodium-archive-keyring.gpg ] https://download.vscodium.com/debs vscodium main' |
        sudo tee /etc/apt/sources.list.d/vscodium.list
}

function setup_debian_base_with_apt() {
    sudo apt-get update
    # shellcheck disable=SC2046
    sudo apt-get install $(get_packages apt-base) -y || return 1
}

function setup_debian_with_apt() {
    sudo apt-get update
    # shellcheck disable=SC2046
    sudo apt-get install $(get_packages $(construct_target_list apt)) -y || return 2
}

function setup_debian_with_git() {
    install_picom || return 1
    install_sddm_aerial_theme || return 2
    build_alacritty || return 3
}

function setup_debian() {
    setup_debian_base_with_apt || return 1
    setup_debian_repo_vscodium || return 2
    setup_debian_with_apt || return 3
    setup_pip_packages || return 4
    setup_debian_with_git || return 5
    configure_sddm || return 6
    configure_ufw || return 7
    setup_homebrew || return 8
    setup_brew_formulas || return 9
    setup_krew || return 10
    setup_krew_plugins || return 11
    change_shell || return 12
}

if [ "$EUID" -eq 0 ]; then
    echo "Don't run this script as root." 1>&2
    exit 2
fi

ID_LIKE="$(grep ID_LIKE= </etc/os-release | cut -d= -f2)"

echo "Setting up ${ID_LIKE^:?}-like OS."

if [[ "$ID_LIKE" == "arch" ]]; then
    setup_arch || echo "Setup failed: $?"
elif [[ "$ID_LIKE" == "debian" ]]; then
    setup_debian || echo "Setup failed: $?"
else
    echo "ERROR: Unsupported system: ID_LIKE=$ID_LIKE"
    exit 3
fi

echo "Adding $USER to 'video' group."
sudo usermod -aG video "$USER"

echo "Setup finished."