#!/bin/bash
clear
# ------------------ Display Device ------------------
echo -e "\e[32m$(lsblk)\e[0m"

# ------------------ User Input ------------------
read -rp "Install Disk    (e.g. /dev/sda): " my_disk
read -p "Enable Encryption?     (YES/NO): " encrypt_disk
read -p "Specify ROOT size?     (YES/NO): " root_size
read -p "BOOT size in MB      (e.g. 500): " fs_boot
read -p "SWAP size in GB        (e.g. 2): " fs_swap
read -p "Enable EFIstub?        (YES/NO): " efi1
read -p "Essential packages?    (YES/NO): " packages
read -p "Use my mirrorlist?     (YES/NO): " my_source
read -p "Hostname             (e.g. bob): " name1

# ------------------ Specify ROOT Size ------------------
if [[ $root_size == YES ]];then
    read -p "ROOT size in GB      (e.g. 20) : " fs_root
fi

#List size of partitions
echo -e "\e[32mThe         boot         size  = ${fs_boot}MB\e[0m"
echo -e "\e[32mThe         swap         size  = ${fs_swap}GB\e[0m"
if [[ $root_size == YES ]];then
    echo -e "\e[32mThe         root         size  = ${fs_root}GB\e[0m"
else
    echo -e "\e[32mThe         root         size  = remaining space\e[0m"
fi

#Confirmation
read -rp "Are you sure?(YES/NO) " confirmation
if [[ $confirmation != YES ]];then
    echo -e "\e[31mAborted...\e[0m"
    exit 1
fi

# ------------------ Create boot partition ------------------
echo -e "\e[32mCreating boot partition...\e[0m"
printf "n\np\n\n\n+${fs_boot}M\nw\n" | fdisk "$my_disk" >/dev/null 2>&1
sleep 3

# ------------------ Create partitions ------------------
#Whether Encrypt
if [[ $encrypt_disk == YES ]];then
    #Specify ROOT Size
    if [[ $root_size == YES ]];then
        total=$((fs_swap + fs_root))
        printf "n\np\n\n\n+${total}G\nw\n" | fdisk "$my_disk" >/dev/null 2>&1
        sleep 3            #Encryption Specified root
    else
        printf "n\np\n\n\n\nw\n" | fdisk "$my_disk" >/dev/null 2>&1
        sleep 3                      #Encryption
    fi
    #Define encryption partition
    echo -e "\e[32m$(lsblk)\e[0m"
    read -rp "Encrypt Partition: " encryption_path
    read -rp "Boot    Partition: " boot_path
    #Encrypt partition
    echo -e "\e[32mEncrypting $encryption_path\e[0m"
    cryptsetup luksFormat "$encryption_path" --type luks2 --pbkdf argon2id -s 512 -h sha512 --iter-time 10000 -c aes-xts-plain64
    #Open encryption partition
    echo -e "\e[32mOpening $encryption_path\e[0m"
    cryptsetup luksOpen "$encryption_path" crypt_device
    #Create logical volume and volume group
    pvcreate /dev/mapper/crypt_device
    vgcreate OS /dev/mapper/crypt_device
    lvcreate -L "${fs_swap}G" OS -n SWAP
    lvcreate -l 100%FREE OS -n ROOT
    echo -e "\e[32mWait 9 seconds......\e[0m"
    #Format partitions
    mkfs.fat -F32 "$boot_path"
    sleep 3
    mkfs.ext4 -q /dev/mapper/OS-ROOT
    sleep 3
    mkswap /dev/mapper/OS-SWAP
    sleep 3
    #Mount partitions
    mount /dev/mapper/OS-ROOT /mnt
    mkdir -p /mnt/boot
    mount "$boot_path" /mnt/boot
    swapon /dev/mapper/OS-SWAP
else
    #Specify ROOT Size
    if [[ $root_size == YES ]];then
        printf "n\np\n\n\n+${fs_swap}G\nw\n" | fdisk "$my_disk" >/dev/null 2>&1          #swap
        printf "n\np\n\n\n+${fs_root}G\nw\n" | fdisk "$my_disk" >/dev/null 2>&1          #root
        sleep 3
    else
        printf "n\np\n\n\n+${fs_swap}G\nw\n" | fdisk "$my_disk" >/dev/null 2>&1          #swap
        printf "n\np\n\n\n\nw\n" | fdisk "$my_disk" >/dev/null 2>&1                      #root
        sleep 3
    fi
    #Define partitions of boot swap root
    echo -e "\e[32m$(lsblk)\e[0m"
    read -rp "Boot                  Partition: " boot_path
    read -rp "Swap                  Partition: " swap_path
    read -rp "Root                  Partition: " root_path
    echo -e "\e[32mWait 9 seconds......\e[0m"
    #Format partitions
    mkfs.fat -F32 "$boot_path"
    sleep 3
    mkfs.ext4 -q "$root_path"
    sleep 3
    mkswap "$swap_path"
    sleep 3
   #Mount partitions
    mount "$root_path" /mnt
    mkdir -p /mnt/boot
    mount "$boot_path" /mnt/boot
    swapon "$swap_path"
fi

#List partitioins infomation
echo -e "\e[33mFormat and mount successful\e[0m"
echo -e "\e[32m$(lsblk)\e[0m"

#Change my source
if [[ $my_source == YES ]];then
    curl https://sh.lihanzhang.cn/mirrorlist -so /etc/pacman.d/mirrorlist
fi

#Install software packages

echo -e "\e[32mBegin to install packges......\e[0m"

sleep 1
#GRUB
if [[ $efi1 = YES ]]; then
    p1=''
else
    p1='grub'
fi
#lvm2
if [[ $encrypt_disk = YES ]]; then
    p2='lvm2'
else
    p2=''
fi
#Install packages
if [[ $packages = YES ]]; then
    pacstrap -K /mnt base linux-lts linux-firmware efibootmgr networkmanager openssh vim sudo zsh zsh-autosuggestions zsh-syntax-highlighting zsh-completions terminus-font ${p1} ${p2}
else
    pacstrap -K /mnt base linux-lts linux-firmware efibootmgr networkmanager openssh vim sudo zsh zsh-autosuggestions zsh-syntax-highlighting zsh-completions \
    terminus-font reflector certbot unzip zip nmap screen wget go git base-devel hexedit duf docker docker-compose docker-buildx netcat nginx \
    fail2ban hyfetch qrencode ${p1} ${p2}
fi
genfstab -U /mnt >> /mnt/etc/fstab

#Download system configuration
curl https://sh.lihanzhang.cn/arch/arch_config.sh -o /mnt/arch_config.sh
chmod +x /mnt/arch_config.sh
#--------------------------------------Create info
part_n=$(echo $boot_path | awk '{print substr($0,length($0),1)}')
echo "$name1 Hostname" > /root/info
echo "$my_disk The Disk. The next is Boot partition" >> /root/info
echo "$part_n Boot number" >> /root/info
#--------------------------------------
if [[ $encrypt_disk == YES ]]; then
    euuid2=$(blkid | grep $encryption_path | awk -F\" '{print $2}' )
    echo '1 Enable encrypt' >> /root/info
    echo "$euuid2" >> /root/info
else
    echo '0 Disable encrypt' >> /root/info
    echo '0 Disable e-uuid' >> /root/info
fi
#--------------------------------------
if [[ $efi1 == YES ]]; then
    echo '1 Enable EFI stub' >> /root/info
else
    echo '0 Disable EFI stub' >> /root/info
fi
#--------------------------------------
if [[ $packages == YES ]]; then
    echo '1 Enable Mini install' >> /root/info
else
    echo '0 Disable Mini install' >> /root/info
fi
#--------------------------------------
if [[ $encrypt_disk == YES ]]; then
    echo '0 No Root path' >> /root/info
else
    echo "$root_path" >> /root/info
fi
cp /root/info /mnt/root/info
#Complete
echo -e "\e[32mAll software packages installed\e[0m"
echo -e "\e[32mExrcute arch-chroot /mnt\e[0m"
echo -e "\e[32mExecute /arch_config.sh\e[0m"