#!/bin/bash

# find distro
if command -v lsb_release &>/dev/null; then
    DISTRO=$(lsb_release -si | tr '[:upper:]' '[:lower:]')
    elif [[ -f /etc/os-release ]]; then
    DISTRO=$(source /etc/os-release && echo $ID | tr '[:upper:]' '[:lower:]')
fi

#function name: check if whiptail is installed
#description: check if whiptail is installed, if not install it
#parameters: none
#return: none
function check_whiptail() {
    if ! command -v whiptail &>/dev/null; then
        if [[ $DISTRO =~ (debian|ubuntu) ]]; then
            sudo apt install whiptail -y
            elif [[ $DISTRO == "fedora" ]]; then
            sudo dnf install newt -y
            elif [[ $DISTRO == "arch" ]]; then
            sudo pacman -S --needed newt --noconfirm
        fi
    fi
}

function setupDesktopEnvironment() {
    DE=$(whiptail --title "Desktop Environment" --menu "Choose a desktop environment" 15 60 8 \
        "1" "GNOME" \
        "2" "KDE" \
        "3" "XFCE" \
        "4" "Cinnamon" \
        "5" "MATE" \
        "6" "Budgie" \
        "7" "Sway" \
        "8" "Hyprland" 3>&1 1>&2 2>&3)
    case $DE in
        1) installGNOME ;;
        2) installKDE ;;
        3) installXFCE ;;
        4) installCinnamon ;;
        5) installMATE ;;
        6) installBudgie ;;
        7) installSway ;;
        8) installHyprland ;;
    esac
}

function installGNOME() {
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install gnome-shell task-gnome-desktop gnome-tweaks gnome-extensions gnome-software-plugin-flatpak -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf group install gnome-desktop -y
        sudo dnf in gnome-tweaks gnome-extensions -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed gnome gnome-tweaks gnome-extensions gnome-packagekit --noconfirm
    fi
}

function installKDE() {
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install kde-standard kdeconnect plasma-discover-backend-flatpak -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf group install kde-desktop -y
        sudo dnf in kde-connect -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed plasma-meta kdeconnect gwenview --noconfirm
    fi
}

function installXFCE() {
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install xfce4 xfce4-goodies -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf in @xfce-desktop-environment -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed xfce4 xfce4-goodies --noconfirm
    fi
}

function installCinnamon() {
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install task-cinnamon-desktop -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf in @cinnamon-desktop-environment -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed cinnamon --noconfirm
    fi
}

function installMATE() {
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install mate-desktop-environment -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf in @mate-desktop-environment -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed mate mate-extra --noconfirm
    fi
}

function installBudgie() {
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install --install-suggests budgie-desktop -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf group install budgie-desktop budgie-desktop-apps -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed budgie budgie-desktop --noconfirm
    fi
}

function installSway() {
    echo "Warning: Experimental Script."
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install sway swaybg swayidle swaylock waybar xdg-desktop-portal-wlr xwayland -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf in sway-desktop-environment sway swaybg swayidle swaylock waybar -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed sway swaybg swayidle swaylock waybar --noconfirm
    fi
}

function installHyprland() {
    echo "Warning: Experimental Script."
    if [[ $DISTRO =~ (debian|ubuntu) ]]; then
        sudo apt install hyprland -y
        elif [[ $DISTRO == "fedora" ]]; then
        sudo dnf in hyprland -y
        elif [[ $DISTRO == "arch" ]]; then
        sudo pacman -S --needed hyprland --noconfirm
    fi
}

echo -e "\e[1;31mWARNING: This script assumes a standard Debian/Fedora/Arch install.\e[0m"
read -r -p "Are you sure you want to continue? [N/y] " response
if [[ "$response" =~ ^[Yy]$ ]]; then
    echo "Continuing..."
    check_whiptail
    setupDesktopEnvironment
else
    echo "Script execution aborted."
    exit 1
fi