#!/bin/bash
export LANG=C
rootfs=""
efifs=""
special_disk="$1"
pause_exit(){
    echo -e "\033[32;1mPRESS ANY KEY TO EXIT\033[;0m"
    read -n 1
    exit $1
}
echo -e "\033[31;1mDetecting rootfs:\033[;0m"
export live_media=$(lsblk -J | grep "/.../live/" | grep -v loop| sed "s/.*name\":\"//g" | sed "s/\".*//g" | sed "s/[0-9]//g")
[[ "${live_media}"  == "" ]] && live_media=null
echo -e "\033[31;1mLive media:\033[;;0m/dev/${live_media}"
for part in $(ls /sys/class/block | grep "[0-9]$"| grep -v "${live_media}" | grep -v fd | grep -v ram | grep -v loop | grep -v sr)
do
    set +e
    umount -lf /dev/$part &>/dev/null || true
    echo -e "\033[31;1mLooking for:\033[;0m/dev/$part"
    tmpdir=$(mktemp -d)
    fsck -yf /dev/$part || true

    if timeout 10 mount -o defaults,ro /dev/$part "$tmpdir" &>/dev/null
    then
        if [ -f "$tmpdir/etc/default/grub" ]
        then
            export distro=$(cat $tmpdir/etc/os-release | grep "^NAME" | sed "s/.*=//g")
            export rootfs=$part
            echo -e "        Found rootfs: /dev/$rootfs (${distro})"
        fi
        if fdisk -l /dev/$(echo $part | sed "s/[0-9]//g") | grep ": gpt" &>/dev/null ; then
            if find "$tmpdir/EFI/" -name "grubx64.efi" &>/dev/null ; then
                export efivar="$(ls $tmpdir/EFI/)"
                export efifs=$part
                echo -e "        Found efi: /dev/$efifs (${efivar[@]})"
            fi
        fi
        if [ -f "$tmpdir/Windows/System32/ntoskrnl.exe" ] ; then
            echo -e "        Found Windows: /dev/$disk"
        fi
    fi
    umount -lf -R $tmpdir &>/dev/null
    set -e
done


if [ "$rootfs" == "" ] ; then
    echo "Error: Failed to detect rootfs."
    pause_exit 1
fi
echo -e "\033[31;1mReinstalling grub:\033[;0m"
rmdir $tmpdir || true
mount /dev/$rootfs /mnt
if [ -d "/sys/firmware/efi" ] && [ ! "$efifs" == "" ]; then
        mkdir -p /mnt/boot/efi
        mount /dev/$efifs /mnt/boot/efi
fi
for i in dev sys proc
do
    mount --bind /$i /mnt/$i
done

if [ -d "/sys/firmware/efi" ] && [ ! "$efifs" == "" ]; then
    mount --bind /sys/firmware/efi/efivars /mnt/sys/firmware/efi/efivars
    chroot /mnt grub-install --target=x86_64-efi /dev/$(echo $rootfs | sed "s/[0-9]//g") || pause_exit 1
else
    chroot /mnt grub-install --target=i386-pc /dev/$(echo $rootfs | sed "s/[0-9]//g") || pause_exit 1
fi
sync
for kernel in $(ls /mnt/lib/modules)
do
    echo -e "\033[31;1mDepmod: \033[;0m$kernel"
    if [ -d /mnt/lib/modules/$kernel/kernel ] ; then
        chroot /mnt depmod -a $kernel || pause_exit 1
    else
        rm -rf /mnt/lib/modules/$kernel
    fi
done
echo -e "\033[31;1mGenerating initramfs:\033[;0m"
if chroot /mnt which update-initramfs &>/dev/null ; then
    chroot /mnt update-initramfs -u -k all  || pause_exit 1
else
    echo "\033[33;1minitramfs-tools not found\033[;0m"
fi
echo -e "\033[31;1mGenerating grub.cfg:\033[;0m"
chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg
if [ -d "/sys/firmware/efi" ] && [ ! "$efifs" == "" ] ; then
    umount /mnt/boot/efi
fi
echo -e "\033[31;1mClearing filesystem:\033[;0m"
find /mnt/var/log -type f | xargs rm -f 2>/dev/null || true
rm -rf /mnt/home/*/.cache/
echo -e "\033[31;1mFix executables:\033[;0m"
for dir in bin usr/bin sbin usr/sbin ; do
    chmod +x -R /mnt/$dir/* 2>/dev/null || true
done
echo -e "\033[31;1mRefresh caches:\033[;0m"
chroot /mnt glib-compile-schemas /usr/share/glib-2.0/schemas/ || true
chroot /mnt update-mime-database /usr/share/mime || true
chroot /mnt update-ca-certificates --fresh || true
chroot /mnt fc-cache -fv || true
umount -lf -R /mnt
pause_exit 0
