#!/usr/bin/env bash
# Arch Linux Installer
# Author: Awkee
# Created: 06/21/2025
# License: GPL3
##################################################

# 脚本路径 #
router_file=$0

CURL="curl -C - "
pac_cmd_ins="pacman -Sy --noconfirm --needed"
FZF="fzf --height 40% --reverse --border --no-info --inline-info --prompt >>> "

# 欢迎和再见提示信息
WELCOME="^_^ Arch linux toolbox"
SEE_YOU="^_^ Bye!"

####################### PROMPT FUNCTION #######################

#### 默认选项 #####
default_confirm="no"    # 是否显示提示确认，no-(默认)提示，yes-自动选择yes
PMT=">>>"
READ_TIMEOUT=30   # read timeout seconds

# Define Colors
RED='\e[41m'
BRB='\e[1;7;31;47m' # Blink Red bold
NC='\e[0m' # No color
BG='\e[7m' # Highlighting Background color
TC='\e[1m' # Highlighting Text color

function echo_greenr() { printf "\033[0;32;7m$@\033[0m"   ; }
function redr_line()   { printf "\033[0;31;7m$@\033[0m\n" ; }

item_index=0        # 记录菜单选项序号
item_line_count=3   # 每行显示菜单数量
ILEN=30             # 单个选项长度
MLEN=$(( (${ILEN}+1) * ${item_line_count}))   # 单行最大长度

function print_feed() {
    for i in $(seq 1 $item_line_count) ; do
        printf "+%${ILEN}s" | tr ' ' '-'
    done
    printf "+\n"
}

function menu_line() { echo -en "|  ${BRB} $@ $NC" ; tput hpa $MLEN ; echo "|" ; }

# 居中显示菜单选项
# 函数： center_line text max_len
function center_line() {
  local text="$1"
  local width="${2:-$MLEN}"  # 默认使用终端宽度

  local plain_text
  plain_text=$(echo -n "$text" | sed 's/\x1B\[[0-9;]*[JKmsu]//g')
  local text_len=${#plain_text}

  # 计算左右填充空格
  local total_padding=$((width - text_len))
  local left_padding=$((total_padding / 2))
  local right_padding=$((total_padding - left_padding))

  # 构建输出：| + 左空格 + 文本 + 右空格 + |
  printf "| %*s${BRB}%s${NC}" "$left_padding" "" "$text" ; tput hpa $width ; printf "|\n"
}

function menu_head() { center_line "$@" ; print_feed; }
# 一行可以有 item_line_count 个菜单选项
function menu_item() { let item_index=$item_index+1 ; n=$1 ; shift ; let rlen="$item_index * ($ILEN + 1)" ; echo -en "|  $BG ${n} $NC $@" ; tput hpa $rlen ; [[ "$item_index" == "$item_line_count" ]] && echo "|" && item_index=0 ; }
# 输出单行长菜单选项,长度有限制
function menu_iteml() { n=$1 ; shift ; echo -en "|  $BG ${n} $NC $@" ; tput hpa $MLEN ; echo "|" ; }
# 用于输入长信息(非菜单选项),不限制结尾长度
function menu_info() { n=$1 ; shift ; echo -e "|  $BG ${n} $NC $@" ; }
function menu_tail() { [[ "$item_index" != "0" ]] && echo "|" ; print_feed; item_index=0 ; }

# 日志记录
function output_msg() { LEVEL="$1" ; shift ; echo -e "$(date +'%Y年%m月%d日%H:%M:%S'):${LEVEL}: $@" ; }
function loginfo() { output_msg "INFO" $@  ; }
function logerr()  { output_msg "ERROR" $@ ; }

#### 检测当前终端支持色彩
function check_term() {
	# 指定 TERM ，避免对齐问题(已知某些rxvt-unicode终端版本存在对齐问题)
    if [[ "`tput colors`" -lt "256" ]] ; then
        export TERM=xterm
        export COLORTERM=truecolor
    fi
}

# 基础依赖命令检测与运行环境 #
function check_basic() {
    command -v curl &>/dev/null || sudo $pac_cmd_ins curl
    command -v fzf  &>/dev/null || sudo $pac_cmd_ins fzf

    # 运行环境检测 #
    [[ "$(id -u)" -ne 0 ]] && echo "Please run this script by root user!" && exit 1

    check_term

    if [ -d /sys/firmware/efi ]; then
        echo "BOOT mode: UEFI"
    else
        echo "BOOT mode: Legacy"
        echo "Please set BIOS to boot from UEFI mode"
        exit 1
    fi
}

function check_mounted() {
    if ! mountpoint -q /mnt ; then
        echo "not found /mnt, please mount /mnt first."
        return 1
    fi
}


################################################################
# 提示确认函数，如果使用 -y 参数默认为Y确认
function prompt() {
    msg="$@"
    if [ "$default_confirm" != "yes" ] ; then
        read -t $READ_TIMEOUT -r -n 1 -e  -p "$msg (y/`echo_greenr N`)" str_answer
        case "$str_answer" in
            y*|Y*)  echo "confirmed" ; return 0 ;;
            *)      echo "canceled" ; return 1 ;;
        esac
    fi
    return 0
}

######## FUNC:INSTALL ARCHER ###################

# 检查磁盘空间并交互设置分区大小
function prompt_sizes() {
    local disk="$1"
    # 获取磁盘总空间（单位：GB）
    local total_size_gb=$(lsblk -b -dn -o SIZE "$disk" | awk '{printf "%.0f", $1/1024/1024/1024}')
    echo "Disk: $disk , Total size: ${total_size_gb}G"

    # 预留1G给EFI，剩余空间用于LVM
    local total_gb=$((total_size_gb - 1))
    if [ "$total_gb" -le 0 ]; then
        echo "No more free space!"
        return 1
    fi
    local free_gb=$total_gb
    echo "Disk free space: ${total_gb}G"

    # 默认分区建议
    local default_root=30
    local default_home=20
    local default_opt=20

    # 交互输入分区大小
    read -p "set /(root) size(GB, default:${default_root}, free:${free_gb}): " root_size
    root_size=${root_size:-$default_root}
    free_gb=$((free_gb - root_size))
    read -p "set /home size(GB, default:${default_home}, free:${free_gb}): " home_size
    home_size=${home_size:-$default_home}
    free_gb=$((free_gb - home_size))
    read -p "set /opt size(GB, default:${default_opt}, free:${free_gb}): " opt_size
    opt_size=${opt_size:-$default_opt}

    # 校验总和
    local sum_size=$((root_size + home_size + opt_size))
    if [ "$sum_size" -gt "$total_gb" ]; then
        echo "Total size is not enough to use! total:[${total_gb}G], sum_size:[$sum_size]!"
        return 1
    fi

    # 导出变量供后续分区函数使用
    export DISK_ROOT_SIZE="${root_size}G"
    export DISK_HOME_SIZE="${home_size}G"
    export DISK_OPT_SIZE="${opt_size}G"
    echo "partation: root=${DISK_ROOT_SIZE}, home=${DISK_HOME_SIZE}, opt=${DISK_OPT_SIZE}"
    return 0
}

function partation_by_lvm() {
    DISK="$1"
    vg_name="sys"

    # 检测是否存在 sys VG组
    if vgs ${vg_name} &>/dev/null ; then
        echo "${vg_name} exist! you need delete it first!"
    else
        echo "Start partation by LVM: $DISK"

        # 创建新的 GPT 分区表
        parted -s $DISK mklabel gpt

        # 创建 EFI 分区 (1GB)
        parted -s $DISK mkpart ESP fat32 1MiB 1025MiB
        parted -s $DISK set 1 esp on

        # 创建 LVM 物理分区 (剩余空间)
        parted -s $DISK mkpart primary ext4 1025MiB 100%

        # 格式化 EFI 分区
        mkfs.fat -F32 ${DISK}1

        # 标记 LVM 分区
        pvcreate ${DISK}2

        # 创建卷组
        vgcreate $vg_name ${DISK}2

        # 检查磁盘空间大小,根据提示设置分区的空间大小,
        prompt_sizes "$DISK" || return 1

        # 创建逻辑卷
        lvcreate -L "$DISK_ROOT_SIZE" -n root $vg_name
        lvcreate -L "$DISK_HOME_SIZE" -n home $vg_name
        lvcreate -L "$DISK_OPT_SIZE" -n opt $vg_name

        # 格式化逻辑卷
        mkfs.ext4 /dev/$vg_name/root
        mkfs.ext4 /dev/$vg_name/home
        mkfs.ext4 /dev/$vg_name/opt
    fi
    # 挂载分区
    mount --mkdir /dev/$vg_name/root /mnt
    mount --mkdir ${DISK}1 /mnt/boot
    mount --mkdir /dev/$vg_name/home /mnt/home
    mount --mkdir /dev/$vg_name/opt /mnt/opt


    echo "Finish to partation and mount"
}

# 挂载分区功能，仅当已经完成分区后使用，可用于修复grub等其他用途或安装失败时重新挂载分区
function mount_partation() {
    # 列举所有可用磁盘
    disks=$(lsblk -d -n -p -o NAME,SIZE,TYPE | grep disk | awk '{print $1 " (" $2 ")"}'|$FZF)
    if [ -z "$disks" ]; then
        echo "no disk selected"
        return 1
    fi
    DISK=$(echo "$disks" | awk '{print $1}')

    vg_name="sys"
    # 检测 VG 组是否存在
    if ! vgs ${vg_name} &>/dev/null ; then
        echo "VG ${vg_name} not exist! you need create it first!"
        return 1
    fi
    echo "Mounting partitions for VG: $vg_name"
    # 挂载分区
    mount --mkdir /dev/$vg_name/root /mnt
    mount --mkdir ${DISK}1 /mnt/boot
    mount --mkdir /dev/$vg_name/home /mnt/home
    mount --mkdir /dev/$vg_name/opt /mnt/opt
    echo "Partitions mounted finished."
}

# 更新软件源(加速软件更新)#
function mirror_select() {
    mirror_file=/etc/pacman.d/mirrorlist

    if grep '^Server' $mirror_file &>/dev/null ; then
        echo -n "current mirror: "
        grep '^Server' $mirror_file
    fi

    prompt "change mirror ?"
    if [ "$?" == "0" ] ; then
        # 更换软件源国内清华源
        echo "updating mirrorlist..."
        pacman -Sy --noconfirm --needed pacman-mirrorlist

        if [ ! -f "$mirror_file" ] ; then
            mkdir -p $(dirname $mirror_file)
            echo 'Server = https://mirrors.tuna.tsinghua.edu.cn/archlinux/$repo/os/$arch' > $mirror_file
            return 0
        fi

        selected_mirror=$(grep 'Server' $mirror_file | grep 'https://' | sed 's/#//g' | $FZF)
        echo -e "selected mirror: ${RED}$selected_mirror${NC}"
        echo
        if [ -z "$selected_mirror" ] ; then
            echo "no mirror selected, exit."
            return 1
        fi

        cp $mirror_file ${mirror_file}.bak.$(date +%Y%m%d%H%M)
        sed -i "s/^Server/#Server/g" $mirror_file
        sed -i "s|#${selected_mirror}|${selected_mirror}|g" $mirror_file
        echo "mirrorlist updated done."
    fi
}

# 01. 磁盘分区
function disk_partation() {
    prompt "Start partation..." || return 1
    # 分区说明： GPT table + UEFI(bootloader) + LVM
    echo "DISK partation: GPT table + UEFI(bootloader) + LVM"
    # 列举所有可用磁盘
    disks=$(lsblk -d -n -p -o NAME,SIZE,TYPE | grep disk | awk '{print $1 " (" $2 ")"}'|$FZF)
    if [ -z "$disks" ]; then
        echo "no disk selected"
        return 1
    fi
    DISK=$(echo "$disks" | awk '{print $1}')

    read -p "Ready to partation: $DISK (y/N): " confirm
    if [ "$confirm" != "Y" ] && [ "$confirm" != "y" ]; then
        exit 1
    fi

    partation_by_lvm  $DISK
    if [ $? -ne 0 ]; then
        echo "Disk partation is failed!"
        return 1
    fi
    echo "Disk partation is done."
}

# 02. 基础安装 # 安装环境准备 #
function install_basic_packages() {
    # 检查磁盘分区是否已经挂载完成
    check_mounted || return 1

    # 安装基础软件包
    prompt "installing base packages..." || return 1
    
    # 更新软件源#
    mirror_select

    pacman -Syy --noconfirm --needed
    pacstrap -K /mnt base linux-lts linux-firmware linux-lts-headers vim sudo shadow lvm2 zsh networkmanager openssh  man-db man-pages texinfo  intel-ucode sof-firmware binutils git curl wget fzf plocate

}

# 03. 系统设置 #
function system_settings() {
    prompt "start to system settings..." || return 1

    read -s -p "root password: " new_password
    echo
    read -s -p "root password(confirm): " confirm_password
    echo
    [[ "$new_password" != "$confirm_password" ]] && echo "invalid password, please re-run the script." && exit 1

    ## 添加用户
    read -p "new username: " username
    read -s -p "new password: " user_new_password
    echo
    read -s -p "new password(confirm): " user_confirm_password
    echo

    [[ "$user_new_password" != "$user_confirm_password" ]] && echo "user[$username],invalid password, please re-run the script." && exit 1

    read -p "new hostname: " hostname
    echo "$hostname" > /mnt/etc/hostname

    arch-chroot /mnt bash -c "
        echo 'config locale zh_CN.UTF-8'
        echo 'en_US.UTF-8 UTF-8' > /etc/locale.gen
        echo 'zh_CN.UTF-8 UTF-8' >> /etc/locale.gen
        echo 'zh_CN.GB18030 GB18030' >> /etc/locale.gen
        echo 'zh_CN.GBK GBK' >> /etc/locale.gen
        locale-gen
        echo "export LANG=zh_CN.UTF-8" >> /etc/profile
        echo "export LC_ALL=zh_CN.UTF-8" >> /etc/profile

        echo 'change localtime to Asia/Shanghai'
        ln -sf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime
    "
    arch-chroot /mnt bash -c "
        echo 'add lvm2 to mkinitcpio.conf'
        sed -i 's/block filesystems/block lvm2 filesystems/g' /etc/mkinitcpio.conf

        echo 'generate initramfs'
        mkinitcpio -P

        echo 'install grub packages'
        pacman -Sy --noconfirm --needed grub-efi-x86_64 efibootmgr os-prober

        echo 'add lvm2 to grub'
        sed -i 's/part_gpt part_msdos/part_gpt lvm2 part_msdos/g' /etc/default/grub

        echo "enable os-prober for grub"
        sed -i 's/^.*GRUB_DISABLE_OS_PROBER=.*/GRUB_DISABLE_OS_PROBER=false/g' /etc/default/grub

        echo 'install grub to EFI'
        grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=arch_linux --recheck

        echo 'generate grub.cfg'
        grub-mkconfig -o /boot/grub/grub.cfg
    "

    # 设置 root 密码
    arch-chroot /mnt bash -c "echo 'root:$new_password' | chpasswd"

    # 创建用户并添加到wheel组
    arch-chroot /mnt bash -c "useradd --create-home -s /usr/bin/zsh -g users $username"
    arch-chroot /mnt bash -c "usermod -aG wheel $username"
    arch-chroot /mnt bash -c "echo '$username:$user_new_password' | chpasswd"
    arch-chroot /mnt bash -c "echo '$username ALL=(ALL) NOPASSWD: ALL' >/etc/sudoers.d/$username"

    # 开机自启动服务配置
    arch-chroot /mnt bash -c "systemctl enable NetworkManager"
    arch-chroot /mnt bash -c "systemctl enable sshd"

    # 生成fstab
    echo "generate /etc/fstab ..."
    genfstab -U -p /mnt >> /mnt/etc/fstab
 
    echo "finish configuration."

    echo "umount /mnt"
    umount -R /mnt


    echo "Now you can reboot your system."

}

function reinstall_grub() {
    # 自动检测已挂载的磁盘（假设/boot/efi或/boot已挂载，或用上次分区的DISK变量）
    check_mounted || return 1

    prompt "install grub..." || return 1

    if [ -d /sys/firmware/efi ]; then
        echo "UEFI mode:"
        arch-chroot /mnt bash -c "
            echo 'install grub packages'
            pacman -Sy --noconfirm --needed grub-efi-x86_64 efibootmgr os-prober

            echo 'add lvm2 to grub'
            sed -i 's/part_gpt part_msdos/part_gpt lvm2 part_msdos/g' /etc/default/grub

            echo "enable os-prober for grub"
            sed -i 's/^.*GRUB_DISABLE_OS_PROBER=.*/GRUB_DISABLE_OS_PROBER=false/g' /etc/default/grub

            echo 'install grub to EFI'
            grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=arch_linux --recheck

            echo 'generate grub.cfg'
            grub-mkconfig -o /boot/grub/grub.cfg

        "
        echo "finish reinstall grub."
    else
        echo "please set BIOS to boot from UEFI mode!"
    fi
}
###############################
#   显卡驱动
###############################
function download_nvidia_driver() {
    prompt "ready to download nvidia driver" || return 1
    command -v jq >/dev/null || sudo $pac_cmd_ins jq
    ! command -v jq >/dev/null && logerr "please install jq command first!" && return 1

    tmp_json=/tmp/tmp.nvidia.menu.json
    # 下载链接(以 NVIDIA GeForce GTX 1660 SUPER 显卡为例,驱动兼容大部分 GeForce系列)
    dn_url='https://gfwsl.geforce.cn/services_toolkit/services/com/nvidia/services/AjaxDriverService.php?func=DriverManualLookup&psid=112&pfid=910&osID=12&languageCode=2052&beta=null&isWHQL=0&dltype=-1&dch=0&upCRD=null&qnf=0&sort1=0&numberOfResults=10'
    
    # 检查文件最后修改时间是否为 15天内 , 是就不重复下载，否则就下载新的json文件
    old_json_file=`find $tmp_json  -type f -mtime -15 2>/dev/null`
    if [ "$old_json_file" = "" ] ; then
        ${CURL} -o $tmp_json $dn_url \
            -H 'authority: gfwsl.geforce.cn' \
            -H 'accept: */*' \
            -H 'accept-language: zh-CN,zh;q=0.9,en;q=0.8,zh-TW;q=0.7' \
            -H 'cache-control: no-cache' \
            -H 'dnt: 1' \
            -H 'origin: https://www.nvidia.cn' \
            -H 'pragma: no-cache' \
            -H 'referer: https://www.nvidia.cn/' \
            -H 'sec-ch-ua: "Google Chrome";v="111", "Not(A:Brand";v="8", "Chromium";v="111"' \
            -H 'sec-ch-ua-mobile: ?0' \
            -H 'sec-ch-ua-platform: "Linux"' \
            -H 'sec-fetch-dest: empty' \
            -H 'sec-fetch-mode: cors' \
            -H 'sec-fetch-site: cross-site' \
            -H 'user-agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36' \
            --compressed
        [[ "$?" != "0" ]] && logerr "failed to get nvidia driver package list!" && return 1
    fi
    file_url=`cat $tmp_json |jq |awk '/DownloadURL/{ print $2 }' | sed -n 's/[",]//g; 1p'`
    file_size=`cat $tmp_json |jq |awk '/DownloadURL/{ print $2 }' | sed -n 's/[",]//g; 2p'`
    loginfo "download url: $file_url , file size: $file_size MB"
    tmp_file="/tmp/`basename $file_url`"

    if [ -f "$tmp_file" ] ; then
        echo -e "${RED}Tips: ${tmp_file} is already downloaded !${NC}"
    fi
    prompt "Download nvidia driver package(file size: $file_size )" && ! ${CURL} -o $tmp_file -SL $file_url && logerr "Download nvidia driver failed!" && return 2

    loginfo "nvidia driver: $tmp_file ."
    sh ${tmp_file}
    loginfo "finish download_nvidia_driver."
}

# 自动安装Arch Linux
function auto_install_arch() {
    prompt "install arch linux" || return 1

    default_confirm="yes"
    disk_partation || return 1
    install_basic_packages || return 2
    system_settings || return 3
}

###############################
#    Desktop Install(桌面安装)
###############################
function install_i3wm() {
    prompt "install i3wm script: i3config" || return 1
    tmp_file=/tmp/i3config
    str_url="https://pushx.link/s/i3config"
    prompt "Download i3config script" && ! ${CURL} -o ${tmp_file} ${str_url} && logerr "Download i3config script failed!" && return 1

    mv $tmp_file /usr/local/bin/i3config
    i3config
}

function install_hyprland() {
    prompt "install hyprland script: wconfig" || return 1
    tmp_file=/tmp/wconfig
    str_url="https://pushx.link/s/wconfig"
    prompt "Download wconfig script" && ! ${CURL} -o ${tmp_file} ${str_url} && logerr "Download wconfig script failed!" && return 1
    mv $tmp_file /usr/local/bin/wconfig
    wconfig
}

function install_gnome() {
    prompt "install GNOME Desktop" || return 1

    sudo pacman -Syu --noconfirm --needed
    sudo pacman -Sy --noconfirm --needed gnome gnome-extra gdm gnome-tweaks
    sudo systemctl enable --now gdm.service
    systemctl is-enabled lightdm && sudo systemctl disable lightdm
    systemctl is-enabled sddm && sudo systemctl disable sddm
}

function install_kde() {
    prompt "install KDE Plasma Desktop" || return 1
    sudo pacman -Syu --noconfirm --needed
    sudo pacman -Sy --noconfirm --needed plasma kde-applications sddm
    sudo systemctl enable --now sddm.service
    systemctl is-enabled lightdm && sudo systemctl disable lightdm
    systemctl is-enabled gdm && sudo systemctl disable gdm
}

function install_Budgie() {
    prompt "install Budgie Desktop" || return 1

    sudo pacman -Syu --noconfirm --needed
    sudo pacman -Sy --noconfirm --needed budgie lightdm lightdm-gtk-greeter
    sudo pacman -Sy --noconfirm --needed arc-gtk-theme papirus-icon-theme network-manager-applet
    sudo systemctl enable --now lightdm.service
    systemctl is-enabled gdm && sudo systemctl disable gdm
    systemctl is-enabled sddm && sudo systemctl disable sddm
}

###############################################
#####  Display managers 安装函数 #####
###############################################

function show_desktop_menu() {

    menu_tail
    menu_head "Windows Manager(WM)"
    menu_item g GNOME
    menu_item k KDE Plasma
    menu_item b Budgie
    menu_tail
    menu_head "Tiling Windows Manager"
    menu_item i "i3WM(i3config)"
    menu_item w "Hyprland(wconfig)"
    menu_tail

    menu_item q Quit
    menu_tail
}


function start_install_desktop() {
    while true
    do
        show_desktop_menu
        read -r -n 1 -e  -p "`echo_greenr choice` ${PMT} " str_answer
        case "$str_answer" in

            g) install_gnome ;;
            k) install_kde ;;
            b) install_Budgie ;;
            i) install_i3wm ;;
            w) install_hyprland ;;
            q|"") return 0 ;;
            *) redr_line "no such choice [$str_answer]!" ;;
        esac
    done
}
######## MAIN:PROCESS ###################


# 主菜单显示 #
function show_menu_main() {
    menu_tail
    menu_head "$WELCOME"
    menu_item 0 auto install ArchLinux
    menu_tail
    
    menu_head "Install Manually"
    menu_item 1 disk partations
    menu_item 2 install basic packages
    menu_item 3 system settings
    menu_tail
    
    menu_item d Desktop Install
    menu_item s select mirror
    menu_item g grub2 reinstall
    menu_item m mount partation
    menu_item n install nvidia driver
    menu_tail

    menu_item r reboot system
    menu_item u Update ainstall
    menu_item q Quit
    menu_tail
}

function start_main(){
    while true
    do
        show_menu_main
        read -r -n 1 -e  -p "`echo_greenr choice` ${PMT} " str_answer
        case "$str_answer" in
            0) auto_install_arch       ;;
            1) disk_partation          ;;
            2) install_basic_packages  ;;
            3) system_settings         ;;

            d) start_install_desktop   ;;
            g) reinstall_grub          ;;
            n) download_nvidia_driver  ;;
            m) mount_partation         ;;
            q|"") return 0             ;;
            r) reboot                  ;;
            s) mirror_select           ;;
            u) update_repo ; exit 0    ;;

            *) redr_line "no such choice [$str_answer]!" ;;
        esac
    done
}



### MAIN:PROCESS ####

check_basic

# 主菜单
start_main

# 再见信息
echo
print_feed
center_line "$SEE_YOU"
print_feed

