#!/bin/bash

set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
source "$SCRIPT_DIR/common.sh"

arch::print_banner() {
    local host="$1"
    printf -- ' \e[H\e[2J\n           \e[0;34m.\n          \e[0;34m/ \\\n         \e[0;34m/   \\\\      \e[1;37m               #     \e[1;34m| *\n        \e[0;34m/^.   \\\\     \e[1;37m a##e #%%" a#"e 6##%%  \e[1;34m| | |-^-. |   | \\ /\n       \e[0;34m/  .-.  \\\\    \e[1;37m.oOo# #   #    #  #  \e[1;34m| | |   | |   |  X\n      \e[0;34m/  (   ) _\\\\   \e[1;37m%%OoO# #   %%#e" #  #  \e[1;34m| | |   | ^._.| / \\ \e[0;37mTM\n     \e[1;34m/ _.~   ~._^\\\\\n    \e[1;34m/.^         ^.\\\\ \e[0;37mTM\n\n\n    \e[1;32mCMtec '%s' install/Update script.\e[0;37m\n' "$host"
}

arch::ensure_prerequisites() {
    if ! pacman -Qs inetutils >/dev/null; then
        sudo pacman -Syy --noconfirm inetutils
    fi
}

arch::maybe_create_snapshot() {
    if pacman -Qs timeshift >/dev/null; then
        common::log_step "Creating backup/snapshot"
        sudo timeshift --create --comments "Update script"
    fi
}

arch::maybe_enable_multilib() {
    if roles::enabled GAME; then
        if ! grep -q "^\\[multilib\\]" /etc/pacman.conf; then
            common::log_step "Enabling multilib"
            sudo tee -a /etc/pacman.conf >/dev/null <<'EOT'

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

arch::update_system_packages() {
    common::log_step "Updating pacman packages"
    sudo pacman -Suyy --noconfirm
    sudo pacman --noconfirm --needed -S base-devel linux-headers git
    common::symlink "$PROJECT_ROOT/gitconfig" "$HOME/.gitconfig"
}

arch::install_graphics_stack() {
    if roles::enabled NVIDIA_GPU || roles::enabled NVIDIA_1080_GPU; then
        common::log_step "Installing NVIDIA drivers"
        local new_kernel="yes"
        if pacman -Qs nvidia-open-dkms >/dev/null; then
            new_kernel="no"
        elif pacman -Qs nvidia-dkms >/dev/null; then
            new_kernel="no"
        fi

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

        if roles::enabled GAME; 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 0
        fi
    else
        common::log_step "Installing Intel drivers"
        sudo pacman --noconfirm --needed -S mesa intel-media-driver
    fi
}

arch::update_aur_packages() {
    common::log_step "Updating AUR packages"
    if pacman -Qs yay >/dev/null; then
        yay --noconfirm --aur
    else
        if [[ -d "$HOME/yay-bin" ]]; then
            rm -rf "$HOME/yay-bin"
        fi
        (
            cd "$HOME"
            git clone https://aur.archlinux.org/yay-bin.git
            cd yay-bin
            makepkg --noconfirm -si
        )
        rm -rf "$HOME/yay-bin"
    fi
}

arch::update_flatpaks() {
    common::log_step "Updating Flatpak packages"
    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
}

arch::install_initial_packages() {
    common::log_step "Installing initial system packages"
    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
    common::link_shell_profiles arch
    common::ensure_local_bin
    common::link_update_shortcut
    common::link_start_scripts
    sudo localectl set-locale LANG=en_US.UTF-8
}

arch::configure_hyprland() {
    common::log_step "Installing hyprland"
    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 "$HOME/.local/bin"
    mkdir -p "$HOME/.config/hypr"

    if [[ -f "$PROJECT_ROOT/config/hypr/hyprland_${HOSTNAME}.conf" ]]; then
        common::symlink "$PROJECT_ROOT/config/hypr/hyprland_${HOSTNAME}.conf" "$HOME/.config/hypr/hyprland_extra.conf"
    else
        : >"$HOME/.config/hypr/hyprland_extra.conf"
    fi

    if [[ -f "$PROJECT_ROOT/config/hypr/hypridle_${HOSTNAME}.conf" ]]; then
        common::symlink "$PROJECT_ROOT/config/hypr/hypridle_${HOSTNAME}.conf" "$HOME/.config/hypr/hypridle.conf"
    else
        common::symlink "$PROJECT_ROOT/config/hypr/hypridle.conf" "$HOME/.config/hypr/hypridle.conf"
    fi

    common::symlink "$PROJECT_ROOT/config/hypr/hyprland.conf" "$HOME/.config/hypr/hyprland.conf"
    common::symlink "$PROJECT_ROOT/wrappedhl" "$HOME/.local/bin/wrappedhl"
    common::symlink "$PROJECT_ROOT/config/hypr/hyprpaper.conf" "$HOME/.config/hypr/hyprpaper.conf"
    common::symlink "$PROJECT_ROOT/config/gtk-3.0" "$HOME/.config"
    common::symlink "$PROJECT_ROOT/config/qt5ct" "$HOME/.config"
    common::symlink "$PROJECT_ROOT/config/qt6ct" "$HOME/.config"
    common::symlink "$PROJECT_ROOT/config/dunst" "$HOME/.config"
    common::symlink "$PROJECT_ROOT/config/walker" "$HOME/.config"

    common::log_step "Installing waybar"
    sudo pacman --noconfirm --needed -S waybar
    mkdir -p "$HOME/.config/waybar"
    common::symlink "$PROJECT_ROOT/config/waybar/config" "$HOME/.config/waybar/config"
    common::symlink "$PROJECT_ROOT/config/waybar/style.css" "$HOME/.config/waybar/style.css"

    common::log_step "Installing autologin"
    yay --noconfirm -S --needed --aur pam_autologin
    sudo /bin/cp -rf "$PROJECT_ROOT/login" /etc/pam.d/login
    sudo /bin/cp -rf "$PROJECT_ROOT/getty@.service" /usr/lib/systemd/system/getty@.service
    sudo touch /etc/security/autologin.conf

    common::log_step "Installing udisks2"
    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

    common::log_step "Installing Audio"
    sudo pacman --noconfirm --needed -S pipewire pipewire-pulse pipewire-alsa pavucontrol helvum
    mkdir -p "$HOME/.config/pipewire"
    common::symlink "$PROJECT_ROOT/config/pipewire/pipewire.conf" "$HOME/.config/pipewire/pipewire.conf"

    common::log_step "Installing Caffeine"
    yay --noconfirm -S --needed --aur caffeine-ng
}

arch::configure_terminal_utilities() {
    common::log_step "Installing terminal utilities"
    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 "$HOME/.config/nvim"
    mkdir -p "$HOME/.config/alacritty"
    mkdir -p "$HOME/.config/lazygit"
    mkdir -p "$HOME/.local/bin"

    common::link_tmux_conf
    common::link_lua_check
    common::link_starship_config
    common::link_nvim_config
    common::link_alacritty_config
    common::link_lazygit_config
    common::symlink "$PROJECT_ROOT/start_nvim.sh" "$HOME/.local/bin/start_nvim"

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

    common::install_fzf

    common::log_step "Installing LSP servers"
    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

    common::log_step "Installing Rust embedded rp2040"
    (
        cd "$HOME"
        rustup target add thumbv6m-none-eabi
        cargo install elf2uf2-rs --locked
        cargo install probe-run
        cargo install flip-link
        mkdir -p "$HOME/.local/bin"
        common::symlink "$PROJECT_ROOT/pico-load.sh" "$HOME/.local/bin/pico-load"
    )

    common::log_step "Installing VirtualHere client"
    common::install_virtualhere_client start
}

arch::install_desktop_base() {
    common::log_step "Installing desktop base utilities"
    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

    common::log_step "Installing other fonts"
    sudo pacman --noconfirm --needed -S noto-fonts poppler-data adobe-source-code-pro-fonts
}

arch::install_desktop_work() {
    common::log_step "Installing desktop work utilities"
    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
}

arch::install_music_tools() {
    common::log_step "Installing music utilities"
    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

    common::user_systemctl --now enable pipewire{,-pulse}.{socket,service} filter-chain.service
    pw-metadata -n settings 0 clock.min-quantum 128
}

arch::install_game_tools() {
    common::log_step "Installing game utilities"
    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 "$HOME/Games/opentrack" ]]; then
        mkdir -p "$HOME/Games"
        (
            cd "$HOME/Games"
            git clone https://github.com/opentrack/opentrack
            cd opentrack/
            mkdir build
            cd build
            cmake ..
            ccmake .
            make
            make install
            cat <<'DESKTOP' | sudo tee /usr/share/applications/opentrack.desktop >/dev/null
[Desktop Entry]
Version=3.1.0
Type=Application
Name=Opentrack
Exec=$HOME/Games/opentrack/build/install/bin/opentrack -platform xcb
Icon=$HOME/Games/opentrack/contrib/cute-octopus-vector-material_15-1831.jpg
Terminal=false
StartupNotify=true
DESKTOP
            sudo update-desktop-database /usr/share/applications
        )
    fi

    if [[ "$HOSTNAME" == "SIMONBOX" ]]; then
        common::log_step "Installing Simon specific game utilities"
        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::enabled VR; then
        common::log_step "Installing VR utilities"
        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
}

arch::install_lab_tools() {
    common::log_step "Installing lab utilities"
    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 "$HOME"
        wget -O spm6103_viewer.py "https://git.cmtec.se/cm/spm6103_viewer/-/raw/main/spm6103_viewer.py?ref_type=heads&inline=false"
    )
}

arch::install_bt_tools() {
    common::log_step "Installing BT utilities"
    sudo pacman -S --noconfirm --needed bluez bluez-utils blueman
    common::sudo_systemctl enable bluetooth.service
    common::sudo_systemctl start bluetooth.service
}

arch::final_message() {
    printf -- '\033[32m \n\n***** Update complete! Please reboot. *****\n\n\033[37m'
}

arch::run() {
    sudo -v
    arch::ensure_prerequisites

    HOSTNAME=$(common::uppercase_hostname)
    export HOSTNAME
    roles::load "$HOSTNAME"

    arch::print_banner "$HOSTNAME"
    roles::print_summary

    sleep 1

    arch::maybe_create_snapshot
    arch::maybe_enable_multilib
    arch::update_system_packages
    arch::install_graphics_stack
    arch::update_aur_packages
    arch::update_flatpaks
    arch::install_initial_packages

    if roles::enabled HYPERLAND; then
        arch::configure_hyprland
    fi

    if roles::enabled TERMINAL; then
        arch::configure_terminal_utilities
    fi

    if roles::enabled DESKTOP_BASE; then
        arch::install_desktop_base
    fi

    if roles::enabled DESKTOP_WORK; then
        arch::install_desktop_work
    fi

    if roles::enabled MUSIC; then
        arch::install_music_tools
    fi

    if roles::enabled GAME; then
        arch::install_game_tools
    fi

    if roles::enabled LAB; then
        arch::install_lab_tools
    fi

    if roles::enabled BT; then
        arch::install_bt_tools
    fi

    arch::final_message
}