#!/bin/bash

#
# Default variables
#

upgrade_packages=false
os_id="$(grep ^ID= </etc/os-release | cut -d= -f2)"
os_id_like="$(grep ^ID_LIKE= </etc/os-release | cut -d= -f2)"
script_dir="$(
    cd -- "$(dirname "$0")" >/dev/null 2>&1 || exit 1
    pwd -P
)"
opt_dir=/opt/dotfiles
python_venv=$opt_dir/pyenv

[ -z "$KREW_ROOT" ] && export KREW_ROOT="${XDG_DATA_HOME:-$HOME/.local/share}/krew"

#
# 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() {
    echo "Installing Python packages into virtualenv $python_venv"

    ! [ -d "$opt_dir" ] && sudo mkdir "$opt_dir"
    sudo python3 -m venv "$python_venv"

    # shellcheck disable=SC1091
    source "$python_venv/bin/activate"

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

    deactivate
}

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="${XDG_DATA_HOME:-$HOME/.local/share}/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() {
    set -e
    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/bin:$PATH"
    set +e
}

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 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 add_debian_keyring() {
    local url="${1:?Missing key URL}" name="${2:?Missing key name}"
    local file="/usr/share/keyrings/$name.gpg"
    echo "Adding keyring for $name from $url."
    wget -qO - "$url" | gpg --dearmor | sudo dd of="$file"
    sudo chmod a+r "$file"
}

function setup_debian_repo_docker() {
    echo "Setting up repository for Docker."
    add_debian_keyring "https://download.docker.com/linux/$os_id/gpg" docker || return 1
    echo "deb [ arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker.gpg ] https://download.docker.com/linux/$os_id $(lsb_release -cs) stable" |
        sudo tee /etc/apt/sources.list.d/docker.list
}

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

function setup_debian_repo_azure() {
    echo "Setting up repository for Azure CLI."
    add_debian_keyring "https://packages.microsoft.com/keys/microsoft.asc" "microsoft" || return 1
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/microsoft.gpg] https://packages.microsoft.com/repos/azure-cli/ $(lsb_release -cs) main" |
        sudo tee /etc/apt/sources.list.d/azure-cli.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
}

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

#
# Main
#

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

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

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

for group in video docker; do
    echo "Adding $USER to '$group' group."
    sudo usermod -aG "$group" "$USER"
done

echo "Setup finished."