#!/bin/bash

set -euo pipefail

# User inputs
HOSTNAME=""
ROOT_PASSWORD=""
USERNAME=""
USER_PASSWORD=""
WANT_AUTOLOGIN=""

# Cleanup on failure
cleanup() {
    local exit_code=$?
    if [ $exit_code -ne 0 ]; then
        echo "An error occurred during installation. Unmounting..."
        umount -A -R /mnt 2>/dev/null || true
    fi
}
trap cleanup EXIT

install_deps() {
    echo "You will need commands like mkfs.ext4, sgdisk and pacstrap."
    printf "Try to install them with pacman (if you are on arch linux)? [Y/n] "
    read -r CONFIRMATION
    case "$CONFIRMATION" in
        [nN]*)
            echo "Operation canceled."
            return
            ;;
        "" | [yY]*)
            echo "Installing dependencies"
            sudo pacman -S --needed --noconfirm dosfstools gptfdisk arch-install-scripts
            echo ""
            ;;
        *)
            echo "Invalid input. Defaulting to 'No'."
            exit 1
            ;;
    esac
}

send_status_message() {
    local term_width
    term_width=$(tput cols 2>/dev/null || echo 80)
    local msg="$1"
    local edge
    edge=$(printf '%*s' "$term_width" '' | tr ' ' '=')
    printf "\n%s\n" "$edge"
    printf "%*s\n" $(((${#msg} + term_width) / 2)) "$msg"
    printf "%s\n\n" "$edge"
}

get_config_inputs() {
    send_status_message "CONFIGURATION: Getting user inputs"

    while [ -z "$WANT_AUTOLOGIN" ]; do
        printf "Do you want to enable autologin? (TTY only) [y/n]: "
        read -r WANT_AUTOLOGIN
        case "$WANT_AUTOLOGIN" in
            [yY]*) WANT_AUTOLOGIN="y" ;;
            [nN]*) WANT_AUTOLOGIN="n" ;;
            *) WANT_AUTOLOGIN="" ;;
        esac
    done
    echo ""

    while [ -z "$HOSTNAME" ]; do
        printf "Enter hostname: "
        read -r HOSTNAME
        # Validate hostname
        if [[ ! "$HOSTNAME" =~ ^[a-zA-Z0-9.-]+$ ]]; then
            echo "Invalid hostname. Use alphanumeric, dots, and hyphens only."
            HOSTNAME=""
            continue
        fi
    done
    echo ""

    while [ -z "$ROOT_PASSWORD" ]; do
        printf "Enter root password: "
        read -rs ROOT_PASSWORD
        printf "\nConfirm root password: "
        read -rs CONFIRM_ROOT_PASSWORD
        echo ""

        if [ "$ROOT_PASSWORD" != "$CONFIRM_ROOT_PASSWORD" ] || [ -z "$ROOT_PASSWORD" ]; then
            echo "Root passwords do not match or are empty."
            ROOT_PASSWORD=""
            continue
        fi
    done
    echo ""

    while [ -z "$USERNAME" ]; do
        printf "Enter username: "
        read -r USERNAME
        # Validate username
        if [[ ! "$USERNAME" =~ ^[a-z_][a-z0-9_-]*$ ]]; then
            echo "Invalid username. Must start with lowercase/underscore and contain no spaces."
            USERNAME=""
            continue
        fi
    done
    echo ""

    while [ -z "$USER_PASSWORD" ]; do
        printf "Enter user password: "
        read -rs USER_PASSWORD
        printf "\nConfirm user password: "
        read -rs CONFIRM_USER_PASSWORD
        echo ""

        if [ "$USER_PASSWORD" != "$CONFIRM_USER_PASSWORD" ] || [ -z "$USER_PASSWORD" ]; then
            echo "User passwords do not match or are empty."
            USER_PASSWORD=""
            continue
        fi
    done
}

# Function to get drive name
DISK=""
get_drive() {
    send_status_message "DRIVE SETUP: Selecting and confirming target drive"

    printf '%-8s %-8s %-30s %s\n' NAME SIZE MODEL
    lsblk # -do NAME,SIZE,MODEL
    printf "\nEnter the drive name (e.g., sda, nvme0n1): "
    read -r DISK
    if [ -z "${DISK}" ]; then
        echo "No drive name entered. Exiting."
        exit 1
    fi

    DISK="/dev/${DISK}"
    if [ ! -b "$DISK" ]; then
        echo "Error: ${DISK} is not a valid block device."
        exit 1
    fi

    PART_PREFIX="$DISK"
    if [[ "$DISK" =~ [0-9]$ ]]; then
        PART_PREFIX="${DISK}p"
    fi

    printf "WARNING: All data on %s will be erased. Are you sure? [Y/n] " "${DISK}"
    read -r CONFIRMATION
    case "$CONFIRMATION" in
        [nN]*)
            echo "Operation canceled."
            exit 0
            ;;
        "" | [yY]*)
            echo "Proceeding with the operation on ${DISK}."
            printf "Wiping drive in "
            for i in 3 2 1; do
                printf '%d...' "$i"
                sleep 1
            done
            echo ""
            ;;
        *)
            echo "Invalid input. Defaulting to 'No'."
            exit 1
            ;;
    esac

    echo "Drive ${DISK} confirmed for installation."
}

setup_partitions() {
    send_status_message "PARTITIONING: Creating partitions and filesystems"

    if [ ! -d "/sys/firmware/efi" ]; then
        echo "BIOS systems are not supported."
        exit 1
    fi

    echo "Ensuring $DISK is not in use..."
    grep "${DISK}" /proc/swaps | cut -d' ' -f1 | xargs -r swapoff || true
    grep "${DISK}" /proc/mounts | cut -d' ' -f2 | xargs -r umount -l || true
    wipefs -a "${DISK}" || true
    sgdisk -Z "${DISK}" || true

    sgdisk -a 2048 -o "${DISK}" # new GPT disk 2048 alignment

    # Create partitions
    sgdisk -n 1::+1G --typecode=1:ef00 --change-name=1:"BOOT" "${DISK}"
    sgdisk -n 2::-0 --typecode=2:8300 --change-name=2:"ROOT" "${DISK}"
    # sgdisk -n 3::+4G --typecode=3:8200 --change-name=3:"SWAP" "${DISK}"

    partprobe "${DISK}"
    sleep 2

    # Create filesystems
    mkfs.fat -F32 -n "BOOT" "${PART_PREFIX}1"
    mkfs.ext4 -F -L "ROOT" "${PART_PREFIX}2"
    # mkswap -f -L "SWAP" "${PART_PREFIX}3"

    # Mount
    mount -t ext4 "${PART_PREFIX}2" /mnt
    mount --mkdir -t vfat "${PART_PREFIX}1" /mnt/boot
    # swapon "${PART_PREFIX}3"
}

# Setup systemdboot as the bootloader
setup_bootloader() {
    send_status_message "BOOTLOADER: Installing and configuring systemd-boot"

    arch-chroot /mnt bootctl install

    BOOTLOADER_ENTRY="/mnt/boot/loader/entries/arch.conf"

    PARTUUID=$(blkid -s PARTUUID -o value "${PART_PREFIX}2")

    if [ -z "$PARTUUID" ]; then
        echo "Error: Unable to find PARTUUID for ${PART_PREFIX}2."
        exit 3
    fi

    echo "Creating bootloader entry with PARTUUID=$PARTUUID"

    cat <<EOF > "$BOOTLOADER_ENTRY"
title   Arch Linux
linux   /vmlinuz-linux
initrd  /initramfs-linux.img
options root=PARTUUID=$PARTUUID rw splash
EOF

    echo "default arch.conf" > /mnt/boot/loader/loader.conf
}

setup_autologin() {
    if [ "$WANT_AUTOLOGIN" != "y" ]; then
        return
    fi

    send_status_message "CONFIG: Configuring autologin for $USERNAME"

    cat > "/mnt/etc/systemd/system/autologin@.service" <<EOF
[Unit]
Description=Autologin on %I
After=systemd-user-sessions.service
Conflicts=getty@%i.service

[Service]
ExecStart=-/sbin/agetty --autologin $USERNAME --noclear %I \$TERM
Type=idle
Restart=always
TTYPath=/dev/%I
TTYReset=yes
TTYVHangup=yes
TTYVTDisallocate=yes

[Install]
WantedBy=multi-user.target
EOF

    arch-chroot /mnt systemctl disable lemurs
    arch-chroot /mnt systemctl enable autologin@tty1.service
}

# >------ Main execution flow ------<

# Install dependencies on the host
install_deps
echo ""

# Ask for hostname, username, user and root passwords
get_config_inputs
echo ""

# Ask for the drive to install on
get_drive
echo ""

# Partition, create filesystems and mount
setup_partitions

# Then install the kernel and some required packages
send_status_message "BASE INSTALLATION: Installing base system with pacstrap"
pacstrap -K /mnt base linux linux-firmware networkmanager lemurs sudo
if [ $? -ne 0 ]; then
    echo "Error: Failed to install base system."
    exit 4
fi

send_status_message "SYSTEM CONFIG: Generating fstab"
genfstab -U /mnt >> /mnt/etc/fstab

# Setup base locales
send_status_message "SYSTEM CONFIG: Setting timezone, locales, and hostname"
ln -sf /usr/share/zoneinfo/Europe/Paris /mnt/etc/localtime
echo 'en_US.UTF-8 UTF-8' >> /mnt/etc/locale.gen
arch-chroot /mnt locale-gen
echo "LANG=en_US.UTF-8" > /mnt/etc/locale.conf
echo "$HOSTNAME" > /mnt/etc/hostname

send_status_message "USER SETUP: Creating user and setting passwords"
# Set root password
arch-chroot /mnt chpasswd <<EOF
root:$ROOT_PASSWORD
EOF

# Create user and set password
arch-chroot /mnt useradd -m "$USERNAME"
arch-chroot /mnt chpasswd <<EOF
$USERNAME:$USER_PASSWORD
EOF

# Setup wheel group
arch-chroot /mnt usermod -aG wheel "$USERNAME"
echo "%wheel ALL=(ALL:ALL) ALL" > /mnt/etc/sudoers.d/10-wheel
chmod 440 /mnt/etc/sudoers.d/10-wheel

send_status_message "SERVICES: Enabling NetworkManager and lemurs"
# Setup services and display manager
arch-chroot /mnt systemctl enable NetworkManager lemurs

# Configure autologin for lemurs if requested
setup_autologin

setup_bootloader

send_status_message "CLEANUP: Unmounting all filesystems"
umount -A -R /mnt # Unmount everything for safety

send_status_message "Installation finished. Reboot now."