#!/usr/bin/env bash
set -euo pipefail

print_usage() {
    echo -e "\033[1;34mUsage: $0 [OPTIONS]\033[0m"
    echo -e "\033[1;34mOptions:\033[0m"
    echo -e "  --help, -h            Show this help message."
    echo -e "  --nonfree             Install non-free drivers."
    echo -e "  --free                Install free drivers only."
    echo -e "  --noconfirm           Do not prompt for confirmation during installation."
    echo -e "  --check               Print any missing hardware profiles."
    echo -e "\033[1;34mEither --nonfree or --free must be specified.\033[0m"
}

print_log() {
    if [[ ! -v CHECK_ONLY ]]; then
        echo -e "${1}"
    fi
}

print_error() {
    if [[ ! -v CHECK_ONLY ]]; then
        echo -e "${1}" >&2
    fi
}

# Parse CLI options
PARAMETERS=("$@")
PARSED_OPTIONS=$(getopt --options="h" --longoptions="help,nonfree,free,pacmanconfig:,noconfirm,sync,check" --name "$0" -- "${PARAMETERS[@]}")
# shellcheck disable=SC2181
if [[ $? -ne 0 ]]; then
  print_error "\033[1;31m\nFailed to parse CLI options\n\033[0m"
fi
eval set -- "$PARSED_OPTIONS"
while true; do
  case "$1" in
    --noconfirm)
    PACMAN_NOCONFIRM=1
    shift
    ;;
    --free)
    NONFREE=0
    shift
    ;;
    --nonfree)
    NONFREE=1
    shift
    ;;
    --pacmanconfig)
    PACMAN_CONFIG="$2"
    shift 2
    ;;
    --sync)
    PACMAN_SYNC=1
    shift
    ;;
    --check)
    CHECK_ONLY=1
    shift
    ;;
    --help|-h)
    print_usage
    exit 0
    ;;
    --)
    shift
    break
    ;;
    *)
    echo "Programming error"
    exit 3
    ;;
  esac
done

if [[ ! -v NONFREE ]]; then
    print_usage
    exit 0
fi

# Double check that we are running as root
if [[ $EUID -ne 0 ]] && [[ ! -v CHECK_ONLY ]]; then
    echo -e "\033[1;31mThe Garuda Hardware Tool must be run as root.\033[0m" >&2
    exit 1
fi

# Packages that are pending installation
packages=()

print_log "\033[1;34mDetecting hardware...\033[0m"

if [[ -x /usr/bin/Xorg ]]; then
    packages+=("garuda-hardware-profile-standard-x11")
else
    packages+=("garuda-hardware-profile-standard")
fi

VIRT=false
NVIDIA=false
NVIDIA_CLOSED=false

# Check if virtual machine is in use
if systemd-detect-virt -q; then
    VIRT=true
fi

# Detect all GPU devices
gpu_devices="$(lspci -mmnn -d '*:*:03xx')"

# Any GPU devices found?
if [[ -n "$gpu_devices" ]]; then
    # 00:02.0 VGA compatible controller [0300]: Intel Corporation Alder Lake-P GT2 [Iris Xe Graphics] [8086:46a6] (rev 0c)
    # 01:00.0 VGA compatible controller [0300]: NVIDIA Corporation GA104 [Geforce RTX 3070 Ti Laptop GPU] [10de:24a0] (rev a1)
    # ${BASH_REMATCH[1]} == pci id
    # ${BASH_REMATCH[2]} == manufacturer name
    # ${BASH_REMATCH[3]} == vendor id
    # ${BASH_REMATCH[4]} == model name
    # ${BASH_REMATCH[6]} == model name 2
    # ${BASH_REMATCH[7]} == model id
    regex='^([^ ]+) "[^"]+\[03..\]" "([^[]+) \[([0-9a-f]{4})\]" "([^[]+)( \[([^]]+)\])? \[([0-9a-f]{4})\]"'

    while read -r line; do
        if [[ $line =~ $regex ]]; then
            pci_id="${BASH_REMATCH[1]}"
            manufacturer="${BASH_REMATCH[2]}"
            vendor_id="${BASH_REMATCH[3]}"
            model_1=${BASH_REMATCH[4]}
            model_2="${BASH_REMATCH[6]}"

            print_log "\033[1;32mDetected GPU: Product: $manufacturer $model_1 $model_2, PCI ID: $pci_id, Vendor ID: $vendor_id\033[0m"

            # Detect nvidia
            if [[ "$vendor_id" == "10de" ]]; then
                # Detect if family is between (and including) volta and maxwell
                if [[ "$model_1" =~ ^([A-Z]{2}) ]]; then
                    family="${BASH_REMATCH[1]}"
                    case "$family" in
                        "GT" | "GF" | "GK")
                            ;; # Older than Maxwell, no support
                        "GV" | "GP" | "GM")
                            NVIDIA_CLOSED=true
                            ;;
                        *)
                            NVIDIA=true
                            ;;
                    esac
                fi
            fi
        fi
    done <<< "$gpu_devices"
fi

if [[ "$VIRT" == true ]]; then
    # If virtual machine, add the VM profile
    print_log "\033[1;32mVirtual machine detected, added: \033[1;34mgaruda-hardware-profile-vm\033[0m\033[1;32m.\033[0m"
    packages+=("garuda-hardware-profile-vm")
fi

if [[ "$NONFREE" == "1" ]]; then
    if [[ "$NVIDIA_CLOSED" == true ]]; then
        print_log "\033[1;32mNVIDIA GPU detected, added: \033[1;34mgaruda-hardware-profile-nvidia-closed\033[0m\033[1;32m.\033[0m"
        packages+=("garuda-hardware-profile-nvidia-closed")
    elif [[ "$NVIDIA" == true ]]; then
        print_log "\033[1;32mNVIDIA GPU detected, added: \033[1;34mgaruda-hardware-profile-nvidia\033[0m\033[1;32m.\033[0m"
        packages+=("garuda-hardware-profile-nvidia")
    fi
elif [[ "$NVIDIA" == true ]] || [[ "$NVIDIA_CLOSED" == true ]]; then
    print_log "\033[1;31mNVIDIA GPU detected, but --free specified. Skipping NVIDIA driver installation.\033[0m"
fi

PACMAN_PARAMS=()
PACMAN_INSTALL_PARAMS=("--needed")
if [[ -v PACMAN_NOCONFIRM ]]; then
    PACMAN_INSTALL_PARAMS+=("--noconfirm")
fi
if [[ -v PACMAN_CONFIG ]]; then
    PACMAN_PARAMS+=("--config" "$PACMAN_CONFIG")
fi

if [[ ${#packages[@]} -gt 0 ]]; then
    if [[ -v PACMAN_SYNC ]]; then
        pacman "${PACMAN_PARAMS[@]}" -Sy
    fi

    existing_packages=($(comm -12 <(pacman ${PACMAN_PARAMS[@]} -Slq | sort) <(printf "%s\n" "${packages[@]}" | sort)))
    if [[ ${#existing_packages[@]} -gt 0 ]]; then
        # Packages that are not yet installed
        pending=($(comm -13 <(pacman ${PACMAN_PARAMS[@]} -Qq | sort) <(printf "%s\n" "${existing_packages[@]}" | sort)))
        if [[ ${#pending[@]} -gt 0 ]]; then
            print_log "\033[1;34mInstalling packages: ${pending[*]}\033[0m"
            if [[ -v CHECK_ONLY ]]; then
                printf "%s\n" "${pending[@]}"
            else
                pacman "${PACMAN_INSTALL_PARAMS[@]}" "${PACMAN_PARAMS[@]}" --needed -S "${pending[@]}"
            fi
        else
            print_error "\033[1;32mAll required packages are already installed: ${existing_packages[*]}\033[0m"
        fi
    else
        print_log "\033[1;31mNo valid packages found for installation.\033[0m"
    fi
else
    print_error "\033[1;32mNo hardware profiles detected, nothing to install.\033[0m"
fi
