#!/bin/bash
set -e
sudo -v

if ! pacman -Qs inetutils >/dev/null; then
    sudo pacman -Syy --noconfirm inetutils
fi

HOSTNAME=$(hostname | tr '[:lower:]' '[:upper:]')

# Declare associative array for roles
declare -A ROLES
for role in LAB BT GAME VR DESKTOP_BASE DESKTOP_WORK CODE NVIDIA_GPU NVIDIA_1080_GPU TERMINAL HYPERLAND; do
    ROLES["$role"]="no"
done

# Assign roles per hostname
case "$HOSTNAME" in
CMBOX)
    ROLES[DESKTOP_BASE]="yes"
    ROLES[DESKTOP_WORK]="yes"
    ROLES[CODE]="yes"
    ROLES[TERMINAL]="yes"
    ROLES[HYPERLAND]="yes"
    ;;
STEAMBOX)
    ROLES[GAME]="yes"
    ROLES[VR]="yes"
    ROLES[DESKTOP_BASE]="yes"
    ROLES[NVIDIA_GPU]="yes"
    ROLES[TERMINAL]="yes"
    ROLES[HYPERLAND]="yes"
    ;;
LABBOX)
    ROLES[DESKTOP_BASE]="yes"
    ROLES[CODE]="yes"
    ROLES[TERMINAL]="yes"
    ROLES[HYPERLAND]="yes"
    ROLES[LAB]="yes"
    ;;
SIMONBOX)
    ROLES[GAME]="yes"
    ROLES[VR]="yes"
    ROLES[DESKTOP_BASE]="yes"
    ROLES[NVIDIA_1080_GPU]="yes"
    ROLES[TERMINAL]="yes"
    ROLES[HYPERLAND]="yes"
    ROLES[BT]="yes"
    ;;
esac

echo -e ' \e[H\e[2J
           \e[0;34m.
          \e[0;34m/ \
         \e[0;34m/   \      \e[1;37m               #     \e[1;34m| *
        \e[0;34m/^.   \     \e[1;37m a##e #%" a#"e 6##%  \e[1;34m| | |-^-. |   | \ /
       \e[0;34m/  .-.  \    \e[1;37m.oOo# #   #    #  #  \e[1;34m| | |   | |   |  X
      \e[0;34m/  (   ) _\   \e[1;37m%OoO# #   %#e" #  #  \e[1;34m| | |   | ^._.| / \ \e[0;37mTM
     \e[1;34m/ _.~   ~._^\
    \e[1;34m/.^         ^.\ \e[0;37mTM


    \e[1;32mCMtec '$HOSTNAME' install/Update script.\e[0;37m
'
# Print enabled roles
printf -- '    \033[37m'
for role in "${!ROLES[@]}"; do
    if [ "${ROLES[$role]}" == "yes" ]; then
        printf '%s, ' "$role"
    fi
done
printf -- '\n\n\033[37m'

sleep 1

# Create backup/snapshot
if pacman -Qs timeshift >/dev/null; then
    CONFIG_FILE="/etc/timeshift/timeshift.json"
    if [ -f "$CONFIG_FILE" ]; then
        BTRFS_MODE=$(jq '.btrfs_mode' "$CONFIG_FILE")
        if [ "$BTRFS_MODE" == '"true"' ]; then
            printf -- '\033[33m Creating backup/snapshot\n\033[37m'
            sudo timeshift --create --comments "Update script"
        fi
    fi
fi

# Enable multilib (if applicable)
if [ "${ROLES[GAME]}" == "yes" ]; then
    if ! grep -q "^\[multilib\]" /etc/pacman.conf; then
        printf -- '\033[33m Enabling multilib\n\033[37m'
        sudo tee -a /etc/pacman.conf >/dev/null <<EOT

[multilib]
Include = /etc/pacman.d/mirrorlist
EOT
    fi
fi

# Update all pacman packages
printf -- '\033[33m Updating pacman packages\n\033[37m'
sudo pacman -Suyy --noconfirm
sudo pacman --noconfirm --needed -S base-devel linux-headers git
ln -sf ~/linuxbox/gitconfig ~/.gitconfig

# Install drivers (if applicable)
if [ "${ROLES[NVIDIA_GPU]}" == "yes" ] || [ "${ROLES[NVIDIA_1080_GPU]}" == "yes" ]; then
    # Install nvidia drivers
    printf -- '\033[33m Installing NVIDIA drivers\n\033[37m'

    if pacman -Qs nvidia-open-dkms >/dev/null; then
        NEW_KERNEL="no"
    elif pacman -Qs nvidia-dkms >/dev/null; then
        NEW_KERNEL="no"
    else
        NEW_KERNEL="yes"
    fi

    if [ "${ROLES[NVIDIA_GPU]}" == "yes" ]; then
        sudo pacman --noconfirm --needed -S nvidia-open-dkms
    elif [ "${ROLES[NVIDIA_1080_GPU]}" == "yes" ]; then
        sudo pacman --noconfirm --needed -S nvidia-dkms
    fi

    sudo pacman --noconfirm --needed -S nvidia-utils nvidia-settings opencl-nvidia cuda

    if [ "${ROLES[GAME]}" == "yes" ]; then
        sudo pacman --noconfirm --needed -S lib32-nvidia-utils
    fi

    if [ "$NEW_KERNEL" == "yes" ]; then
        printf -- '\033[33m \n\n***** Nvidia driver updated kernel! Please reboot and run update once more! *****\n\n\033[37m'
        exit
    fi
else
    # Install intel drivers
    printf -- '\033[33m Installing Intel drivers\n\033[37m'
    sudo pacman --noconfirm --needed -S mesa intel-media-driver
fi

# Update all AUR packages
printf -- '\033[33m Updating AUR packages\n\033[37m'
if pacman -Qs yay >/dev/null; then
    yay --noconfirm --aur
else
    if [ -d ~/yay-bin ]; then rm -rf ~/yay-bin; fi
    cd ~
    git clone https://aur.archlinux.org/yay-bin.git
    cd yay-bin
    makepkg --noconfirm -si
    cd ~
    if [ -d ~/yay-bin ]; then rm -rf ~/yay-bin; fi
fi

# Update all Flatpak packages
printf -- '\033[33m Updating Flatpak packages\n\033[37m'
if pacman -Qs flatpak >/dev/null; then
    flatpak update -y
else
    sudo pacman --noconfirm --needed -S flatpak
    flatpak remote-add --if-not-exists flathub https://flathub.org/repo/flathub.flatpakrepo
    flatpak install -y flathub com.github.tchx84.Flatseal
fi

# Install initial system packages
printf -- '\033[33m Installing initial system packages\n\033[37m'
sudo pacman --noconfirm --needed -S timeshift bc git-lfs cmake gawk wget gettext unzip curl inetutils python python-pip python-pipx python-pipenv python-pynvim rustup
yay --noconfirm -S --needed --aur downgrade
rustup update stable
rustup default stable
ln -sf ~/linuxbox/bashrc ~/.bashrc
mkdir -p ~/.local/bin
ln -sf ~/linuxbox/update_wrapper.sh ~/.local/bin/update
sudo localectl set-locale LANG=en_US.UTF-8

# Install hyprland
if [ "${ROLES[HYPERLAND]}" == "yes" ]; then
    printf -- '\033[33m Installing hyprland\n\033[37m'
    sudo pacman --noconfirm --needed -S usbutils plymouth dracut dunst hyprpaper hypridle hyprland hyprlock xdg-desktop-portal-hyprland polkit-gnome xorg-xhost gnome-keyring qt5ct qt6ct gnome-themes-extra qt5-wayland qt6-wayland lxappearance qt5-tools adwaita-fonts gnome-disk-utility

    yay --noconfirm -S --needed --aur adwaita-qt5-git
    yay --noconfirm -S --needed --aur adwaita-qt6-git
    yay --noconfirm -S --needed --aur hyprshot

    mkdir -p ~/.local/bin
    mkdir -p ~/.config/hypr
    mkdir -p ~/.config/gtk-3.0
    mkdir -p ~/.config/qt5ct
    mkdir -p ~/.config/qt6ct
    mkdir -p ~/.config/dunst

    if [ -f ~/linuxbox/config/hypr/hyprland_$HOSTNAME.conf ]; then
        ln -sf ~/linuxbox/config/hypr/hyprland_$HOSTNAME.conf ~/.config/hypr/hyprland_extra.conf
    else
        touch ~/.config/hypr/hyprland_extra.conf
    fi

    ln -sf ~/linuxbox/config/hypr/hyprland.conf ~/.config/hypr/hyprland.conf
    ln -sf ~/linuxbox/wrappedhl ~/.local/bin/wrappedhl
    ln -sf ~/linuxbox/config/hypr/hypridle.conf ~/.config/hypr/hypridle.conf
    ln -sf ~/linuxbox/config/hypr/hyprpaper.conf ~/.config/hypr/hyprpaper.conf
    ln -sf ~/linuxbox/config/gtk-3.0/settings.ini ~/.config/gtk-3.0/settings.ini
    ln -sf ~/linuxbox/config/qt5ct/qt5ct.conf ~/.config/qt5ct/qt5ct.conf
    ln -sf ~/linuxbox/config/qt6ct/qt6ct.conf ~/.config/qt6ct/qt6ct.conf
    ln -sf ~/linuxbox/config/dunst/dunstrc ~/.config/dunst/dunstrc

    printf -- '\033[33m Installing waybar\n\033[37m'
    sudo pacman --noconfirm --needed -S waybar
    mkdir -p ~/.config/waybar
    ln -sf ~/linuxbox/config/waybar/config ~/.config/waybar/config
    ln -sf ~/linuxbox/config/waybar/style.css ~/.config/waybar/style.css

    printf -- '\033[33m Installing autologin\n\033[37m'
    yay --noconfirm -S --needed --aur pam_autologin
    sudo /bin/cp -rf ~/linuxbox/login /etc/pam.d/login
    sudo /bin/cp -rf ~/linuxbox/getty@.service /usr/lib/systemd/system/getty@.service
    sudo touch /etc/security/autologin.conf

    printf -- '\033[33m Installing udisks2\n\033[37m'
    sudo pacman --noconfirm --needed -S udisks2 udiskie
    echo 'ENV{ID_FS_USAGE}=="filesystem|other|crypto", ENV{UDISKS_FILESYSTEM_SHARED}="1"' | sudo tee /lib/udev/rules.d/99-udisks2.rules
    echo 'D /media 0755 root root 0 -' | sudo tee /etc/tmpfiles.d/media.conf
    if [ "$HOSTNAME" == "STEAMBOX" ]; then
        echo "$USER ALL=(ALL:ALL) NOPASSWD: /usr/sbin/udisksctl mount -b /dev/sda1" | sudo tee /etc/sudoers.d/udisksctl
    fi

    printf -- '\033[33m Installing Audio\n\033[37m'
    sudo pacman --noconfirm --needed -S pipewire pipewire-pulse pipewire-alsa pavucontrol helvum
    mkdir -p ~/.config/pipewire
    ln -sf ~/linuxbox/config/pipewire/pipewire.conf ~/.config/pipewire/pipewire.conf

fi

# Install terminal utility
if [ "${ROLES[TERMINAL]}" == "yes" ]; then
    printf -- '\033[33m Installing terminal utilities\n\033[37m'
    if pacman -Qs gnu-free-fonts >/dev/null; then
        sudo pacman --noconfirm -R gnu-free-fonts
    fi
    sudo pacman --noconfirm --needed -S alacritty btop ranger tmux fd neovim ttf-nerd-fonts-symbols ttf-roboto-mono-nerd

    cargo install ripgrep eza
    cargo install zoxide --locked
    cargo install starship --locked

    mkdir -p ~/.config/nvim
    mkdir -p ~/.config/alacritty

    ln -sf ~/linuxbox/tmux.conf ~/.tmux.conf
    ln -sf ~/linuxbox/config/starship.toml ~/.config/starship.toml
    ln -sf ~/linuxbox/config/nvim/init.lua ~/.config/nvim/init.lua
    ln -sf ~/linuxbox/config/alacritty/alacritty.toml ~/.config/alacritty/alacritty.toml

    printf -- '\033[33m Installing fzf\n\033[37m'
    rm -rf ~/.fzf
    git clone --depth 1 https://github.com/junegunn/fzf.git ~/.fzf
    ~/.fzf/install --all
    source ~/.fzf.bash

    printf -- '\033[33m Installing fzfmenu\n\033[37m'
    mkdir -p ~/.local/bin
    ln -sf ~/linuxbox/fzfmenu.sh ~/.local/bin/fzfmenu

    printf -- '\033[33m Installing lazygit\n\033[37m'
    cd ~
    LAZYGIT_VERSION=$(curl -s "https://api.github.com/repos/jesseduffield/lazygit/releases/latest" | grep -Po '"tag_name": "v\K[^"]*')
    curl -Lo lazygit.tar.gz "https://github.com/jesseduffield/lazygit/releases/latest/download/lazygit_${LAZYGIT_VERSION}_Linux_x86_64.tar.gz"
    tar xf lazygit.tar.gz lazygit
    sudo install lazygit /usr/local/bin
    rm ~/lazygit && rm ~/lazygit.tar.gz
    mkdir -p ~/.config/lazygit
    ln -sf ~/linuxbox/config/lazygit/config.yml ~/.config/lazygit/config.yml

    printf -- '\033[33m Installing ble.sh\n\033[37m'
    if [ -d ~/ble.sh ]; then rm -rf ~/ble.sh; fi
    cd ~
    git clone --recursive --depth 1 --shallow-submodules https://github.com/akinomyoga/ble.sh.git
    make -C ble.sh install
    if [ -d ~/ble.sh ]; then rm -rf ~/ble.sh; fi
    ln -sf ~/linuxbox/blerc ~/.blerc
fi

# Install code utility
if [ "${ROLES[CODE]}" == "yes" ]; then
    printf -- '\033[33m Installing code utilities\n\033[37m'
    sudo pacman --noconfirm --needed -S code gdu ruby bottom go php luarocks composer jdk-openjdk julia nodejs npm platformio-core
    sudo npm install -g neovim
    cargo install tree-sitter-cli

    ln -sf ~/linuxbox/clang-format ~/.clang-format
    ln -sf ~/linuxbox/luacheckrc ~/.luacheckrc

    printf -- '\033[33m Installing Rust embedded rp2040\n\033[37m'
    cd ~
    rustup target add thumbv6m-none-eabi
    cargo install elf2uf2-rs --locked
    cargo install probe-run flip-link
    mkdir -p ~/.local/bin
    ln -sf ~/linuxbox/pico-load.sh ~/.local/bin/pico-load

    printf -- '\033[33m Installing teensy udev rules\n\033[37m'
    if [ ! -f /.dockerenv ]; then
        sudo rm -f /tmp/00-teensy.rules /etc/udev/rules.d/00-teensy.rules /lib/udev/rules.d/00-teensy.rules
        sudo wget -O /tmp/00-teensy.rules https://www.pjrc.com/teensy/00-teensy.rules
        sudo install -o root -g root -m 0664 /tmp/00-teensy.rules /lib/udev/rules.d/00-teensy.rules
        sudo udevadm control --reload-rules
        sudo udevadm trigger
    fi

    printf -- '\033[33m Installing VirtualHere client\n\033[37m'
    cd ~
    wget https://www.virtualhere.com/sites/default/files/usbclient/scripts/virtualhereclient.service
    wget https://www.virtualhere.com/sites/default/files/usbclient/vhclientx86_64
    wget https://www.virtualhere.com/sites/default/files/usbclient/vhuit64
    chmod +x ./vhclientx86_64
    chmod +x ./vhuit64
    sudo mv ./vhclientx86_64 /usr/sbin
    sudo mv ./vhuit64 /usr/sbin
    echo "$USER ALL=(ALL:ALL) NOPASSWD: /usr/sbin/vhclientx86_64" | sudo tee /etc/sudoers.d/$USER+vhclientx86_64
    echo "$USER ALL=(ALL:ALL) NOPASSWD: /usr/sbin/vhuit64" | sudo tee /etc/sudoers.d/$USER+vhuit64
    sudo mv virtualhereclient.service /etc/systemd/system/virtualhereclient.service
    sudo systemctl daemon-reload
    sudo systemctl enable virtualhereclient.service
    sudo systemctl start virtualhereclient.service

fi

# Install desktop utility
if [ "${ROLES[DESKTOP_BASE]}" == "yes" ]; then
    printf -- '\033[33m Installing desktop base utilities\n\033[37m'
    sudo pacman --noconfirm --needed -S feh
    flatpak install -y flathub io.github.equicord.equibop
    flatpak install -y flathub com.behringer.XAirEdit
    flatpak install -y flathub com.moonlight_stream.Moonlight
    flatpak install -y flathub com.github.vikdevelop.timer
    flatpak install -y flathub io.github.efogdev.mpris-timer
    yay --noconfirm -S --needed --aur brave-bin

    printf -- '\033[33m Installing other fonts\n\033[37m'
    sudo pacman --noconfirm --needed -S noto-fonts poppler-data adobe-source-code-pro-fonts

    printf -- '\033[33m Installing Sunshine\n\033[37m'
    if ! grep -q "^\[lizardbyte\]" /etc/pacman.conf; then
        sudo tee -a /etc/pacman.conf >/dev/null <<EOT

[lizardbyte]
SigLevel = Optional
Server = https://github.com/LizardByte/pacman-repo/releases/latest/download
EOT
        sudo pacman -Sy
    fi
    sudo pacman --noconfirm --needed -S lizardbyte/sunshine
    mkdir -p ~/.config/systemd/user
    cp -f ~/linuxbox/config/systemd/user/sunshine.service ~/.config/systemd/user/sunshine.service
    sudo udevadm control --reload-rules
    sudo udevadm trigger
    sudo setcap cap_sys_admin+p $(readlink -f $(which sunshine))
    sudo systemctl daemon-reload
fi

if [ "${ROLES[DESKTOP_WORK]}" == "yes" ]; then
    printf -- '\033[33m Installing desktop work utilities\n\033[37m'
    sudo pacman --noconfirm --needed -S kicad freecad
    flatpak install -y flathub org.kde.krita
    flatpak install -y flathub com.prusa3d.PrusaSlicer
    flatpak install -y flathub com.jgraph.drawio.desktop
    flatpak install -y flathub org.gimp.GIMP
fi

# Install game utility
if [ "${ROLES[GAME]}" == "yes" ]; then
    printf -- '\033[33m Installing game utilities\n\033[37m'
    sudo pacman -S --noconfirm --needed wine-staging
    sudo pacman -S --noconfirm --needed winetricks wine-mono wine-gecko qt5-tools
    sudo pacman -S --noconfirm --needed steam
    yay --noconfirm -S --needed --aur protonup-qt protontricks
    yay --noconfirm -S --needed --aur lug-helper
    yay --noconfirm -S --needed --aur opentrack-git
    yay --noconfirm -S --needed --aur openmpi

    if [ "$HOSTNAME" == "SIMONBOX" ]; then
        flatpak install -y flathub org.vinegarhq.Sober
        sudo pacman -S --noconfirm --needed fluidsynth gamemode gvfs libayatana-appindicator innoextract lib32-gamemode lib32-vkd3d python-pefile python-protobuf vulkan-icd-loader vkd3d lib32-vulkan-icd-loader vulkan-tools xorg-xgamma umu-launcher
        sudo pacman -S --noconfirm --needed lutris
    fi
fi

# Install VR utility
if [ "${ROLES[VR]}" == "yes" ]; then
    printf -- '\033[33m Installing VR utilities\n\033[37m'
    sudo pacman -S --noconfirm --needed cli11 glib2-devel nlohmann-json glew
    yay --noconfirm -S --needed --aur monado-vulkan-layers-git
    yay --noconfirm -S --needed --aur wlx-overlay-s-git
    yay --noconfirm -S --needed --aur xrgears
    yay --noconfirm -S --needed --aur envision-xr-git
fi

# Install lab utility
if [ "${ROLES[LAB]}" == "yes" ]; then
    printf -- '\033[33m Installing lab utilities\n\033[37m'
    yay --noconfirm -S --needed --aur ps7_libpicoipp ps7_libpicocv
    yay --noconfirm -S --needed --aur picoscope7
    yay --noconfirm -S --needed --aur ps7_libps2000a
    yay --noconfirm -S --needed --aur nrf-udev
    yay --noconfirm -S --needed --aur nrfconnect-appimage
    cd~ && wget "https://git.cmtec.se/cm/spm6103_viewer/-/raw/main/spm6103_viewer.py?ref_type=heads&inline=false"
fi

# Install BT utility
if [ "${ROLES[BT]}" == "yes" ]; then
    printf -- '\033[33m Installing BT utilities\n\033[37m'
    sudo pacman -S --noconfirm --needed bluez bluez-utils blueman
    sudo systemctl enable bluetooth.service
    sudo systemctl start bluetooth.service
fi

printf -- '\033[32m \n\n***** Update complete! Please reboot. *****\n\n\033[37m'