#!/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 MUSIC 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"
    ROLES[MUSIC]="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"
    ROLES[BT]="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
    printf -- '\033[33m Creating backup/snapshot\n\033[37m'
    sudo timeshift --create --comments "Update script"
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-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_arch ~/.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 wayvnc nm-connection-editor usbutils plymouth dracut dunst hyprpaper hypridle hyprland hyprlock xdg-desktop-portal-hyprland polkit-gnome xorg-xhost gnome-keyring qt6ct qt6ct gnome-themes-extra qt5-wayland qt6-wayland lxappearance qt5-tools adwaita-fonts gnome-disk-utility hyprpaper tk

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

    mkdir -p ~/.local/bin
    mkdir -p ~/.config/hypr

    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

    if [ -f ~/linuxbox/config/hypr/hypridle_$HOSTNAME.conf ]; then
        ln -sf ~/linuxbox/config/hypr/hypridle_$HOSTNAME.conf ~/.config/hypr/hypridle.conf
    else
        ln -sf ~/linuxbox/config/hypr/hypridle.conf ~/.config/hypr/hypridle.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/hyprpaper.conf ~/.config/hypr/hyprpaper.conf
    ln -sf ~/linuxbox/config/gtk-3.0 ~/.config
    ln -sf ~/linuxbox/config/qt5ct ~/.config
    ln -sf ~/linuxbox/config/qt6ct ~/.config
    ln -sf ~/linuxbox/config/dunst ~/.config
    ln -sf ~/linuxbox/config/walker ~/.config

    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

    printf -- '\033[33m Installing Caffeine\n\033[37m'
    yay --noconfirm -S --needed --aur caffeine-ng
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 dysk lazygit alacritty btop ranger tmux fd ttf-nerd-fonts-symbols ttf-roboto-mono-nerd gdu ruby bottom go php luarocks composer jdk-openjdk julia nodejs npm
    yay --noconfirm -S --needed --aur neovim-git-bin

    sudo npm install -g neovim

    cargo install tree-sitter-cli ripgrep eza
    cargo install zoxide --locked
    cargo install starship --locked

    mkdir -p ~/.config/nvim
    mkdir -p ~/.config/alacritty
    mkdir -p ~/.config/lazygit
    mkdir -p ~/.local/bin

    ln -sf ~/linuxbox/tmux.conf ~/.tmux.conf
    ln -sf ~/linuxbox/luacheckrc ~/.luacheckrc
    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
    ln -sf ~/linuxbox/config/lazygit/config.yml ~/.config/lazygit/config.yml
    ln -sf ~/linuxbox/start_nvim.sh ~/.local/bin/start_nvim

    if [ -d ~/.local/share/nvim/lazy ]; then
        nvim --headless "+Lazy! sync" +qa
    fi

    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 LSP servers\n\033[37m'
    sudo pacman --noconfirm --needed -S prettier stylua python-black shfmt lua-language-server bash-language-server ccls vscode-html-languageserver vscode-json-languageserver marksman pyright yaml-language-server vscode-css-languageserver clang
    yay --noconfirm -S --needed --aur dockerfile-language-server
    rustup component add rust-analyzer clippy rustfmt

    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 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 cameractrls feh
    flatpak install -y flathub com.discordapp.Discord
    flatpak install -y flathub com.behringer.XAirEdit
    flatpak install -y flathub com.github.vikdevelop.timer
    flatpak install -y flathub io.github.efogdev.mpris-timer
    flatpak install -y flathub org.remmina.Remmina
    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
fi

if [ "${ROLES[DESKTOP_WORK]}" == "yes" ]; then
    printf -- '\033[33m Installing desktop work utilities\n\033[37m'
    sudo pacman --noconfirm --needed -S kicad kicad-library 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 music production utility
if [ "${ROLES[MUSIC]}" == "yes" ]; then
    printf -- '\033[33m Installing music utilities\n\033[37m'
    sudo pacman -S --noconfirm --needed pipewire pipewire-jack pipewire-alsa pipewire-pulse gamemode
    sudo pacman -S --noconfirm --needed qpwgraph realtime-privileges
    flatpak install -y flathub com.bitwig.BitwigStudio

    sudo usermod -a -G realtime,audio,gamemode $USER

    sudo tee /etc/security/limits.d/99-pipewire.conf <<'EOF'
@audio soft rtprio 95
@audio soft memlock unlimited
@audio hard rtprio 95
@audio hard memlock unlimited
EOF

    systemctl --user --now enable pipewire{,-pulse}.{socket,service} filter-chain.service
    pw-metadata -n settings 0 clock.min-quantum 128
fi

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

    if [ ! -d ~/Games/opentrack ]; then
        mkdir -p ~/Games
        cd ~/Games
        git clone https://github.com/opentrack/opentrack
        cd opentrack/
        mkdir build
        cd build
        cmake ..
        ccmake .
        make
        make install
        echo "[Desktop Entry]
Version=3.1.0
Type=Application
Name=Opentrack
Exec=~/Games/opentrack/build/install/bin/opentrack -platform xcb
Icon=~/Games/opentrack/contrib/cute-octopus-vector-material_15-1831.jpg
Terminal=false
StartupNotify=true" | sudo tee /usr/share/applications/opentrack.desktop >/dev/null
        sudo update-desktop-database /usr/share/applications
    fi

    if [ "$HOSTNAME" == "SIMONBOX" ]; then
        printf -- '\033[33m Installing Simon specific game utilities\n\033[37m'
        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

    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
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 ps7_libps3000a
    yay --noconfirm -S --needed --aur nrf-udev
    yay --noconfirm -S --needed --aur nrfconnect-appimage
    sudo pacman -S --noconfirm --needed tk python-pyserial
    cd ~ && wget -O spm6103_viewer.py "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'
