#!/bin/bash
source $(dirname $BASH_SOURCE)/vtcommon

DEVICE=""
TARGET=""
SHUT=""
USE_LUKS=""
NO_BOOT_PARTITION=""

# Parse arguments
while [ $# -gt 0 ]; do
    case "$1" in
        --luks)
            USE_LUKS=yes
            ;;
        --no-boot-partition)
            NO_BOOT_PARTITION=yes
            ;;
        shut)
            SHUT=yes
            ;;
        *)
            if [ "$DEVICE" == "" ]; then
                DEVICE="$1"
                TARGET="/dev/$1"
            fi
            ;;
    esac
    shift
done


# setting SHUT environment variable will disable all useless data and prompt
# this is only recommended for a wrapper around the script

set -e # exit on error

root_device(){
    lsblk --raw -d -no PKNAME $(findmnt / --raw -no SOURCE)
}

valid_devices(){
    lsblk -d --raw -no NAME | grep -v mtd | grep -v $(root_device) | grep -v boot
}

valid_devices_ex(){
    lsblk -d --raw -no NAME,SIZE | grep -v mtd | grep -v $(root_device) | grep -v boot
}

# List valid devices if no argument
if [ "$DEVICE" == "" ];then
    if [ "$SHUT" == "" ];then
        if [ "$(valid_devices_ex)" == "" ];then
            echo "no possible installation devices"
            exit 0
        fi

        echo "Usage: vtinstall [--luks] <device> [shut]"
        echo ""
        echo "Options:"
        echo "  --luks    Enable LUKS encryption on the root partition"
        echo "  --no-boot-partition    Do not create a separate /boot partition"
        echo "  shut      Suppress prompts and verbose output"
        echo ""
        echo "Examples:"
        echo "  vtinstall mmcblk0                 # Standard installation"
        echo "  vtinstall --luks mmcblk0          # Installation with LUKS encryption"
        echo "  vtinstall mmcblk0 shut            # Silent installation"
        echo "  vtinstall --luks mmcblk0 shut     # Silent LUKS installation"
        echo ""
        echo "Possible installation devices:"
        valid_devices_ex
        echo ""
        exit 0
    fi

    valid_devices
    # filter out block devices which might not be a good place to install a system
    exit 0
fi

# Validating the device
if [ ! -e "$TARGET" ]; then
    echo "$TARGET doesn't exist"
    exit 0
fi

if [ "$(valid_devices | grep "$DEVICE")" == "" ];then
    echo "The device is not a valid device."
    echo "It might be a partition."
    exit 0
fi



# Begin installation
need_root

if [ "${SHUT}" == "" ];then
    echo "The selected device $TARGET is going to be wiped."
    echo "This operation CANNOT BE UNDONE."
    echo "If you have any sensitive data present on this device"
    echo "please back it up before proceeding."
    echo ""
    echo "Would you like to continue ?"
    need_confirmation 
fi

echo "[info] unmounting all partitions of target."
if [ "$(lsblk --raw -no MOUNTPOINTS $TARGET | tr '\n' ' ' | tr -d ' ')" != "" ];then
    umount $(lsblk --raw -no MOUNTPOINTS $TARGET | tr '\n' ' ') #sed didn't work but tr seems to also be in the unix core util
fi

echo "[info] wiping partition table."
sgdisk -Z "$TARGET"
partprobe "$TARGET"
sgdisk -C -e -G "$TARGET"
partprobe "$TARGET"
echo "[info] creating a new one."
cgpt create "$TARGET"

echo "[info] creating the kernel partitions."
cgpt add -i 1 -t kernel -b 8192 -s 262144 -l KernelA -S 1 -T 2 -P 10 "$TARGET"
cgpt add -i 2 -t kernel -b 270336 -s 262144 -l KernelB -S 0 -T 2 -P 5 "$TARGET"
if [ "$NO_BOOT_PARTITION" == "yes" ]; then
    BOOT_PART_NUM=""
    ROOT_PART_NUM="3"
    echo "[info] creating the other partitions (single root)."
    fdisk --wipe-partitions always "$TARGET" << EOF > /dev/null #only report issue
n
3


p
w
EOF
else
    BOOT_PART_NUM="3"
    ROOT_PART_NUM="4"
    echo "[info] Creating the other partitions (boot + root)."
    fdisk --wipe-partitions always "$TARGET" << EOF > /dev/null #only report issue
n
3

+1024M
t
3
20
n
4


t
4
20
x
A
3
p
r
p
w
EOF
fi

partprobe $TARGET

lsblk $TARGET

echo "[info] creating/mounting the file system."

#get partition like /dev/sda or /dev/mmcblk1p
TARGET_PART="/dev/$(lsblk --raw -no NAME $TARGET | sed -n '2 p' | sed 's/[0-9]*$//')"

if [ "$NO_BOOT_PARTITION" == "" ]; then
    mkfs -t ext4 -O ^has_journal -m 0 -L bootemmc -F ${TARGET_PART}${BOOT_PART_NUM}
fi

if [ "$USE_LUKS" == "yes" ]; then
    echo "[info] Setting up LUKS encryption on root partition."
    
    if [ "${SHUT}" == "" ];then
        echo "LUKS encryption will be set up on the root partition."
        echo "You will be prompted to enter a passphrase for the encryption."
        echo "Please choose a strong passphrase, and be sure to remember it!"
        echo ""
    fi
    
    cryptsetup luksFormat ${TARGET_PART}${ROOT_PART_NUM}
    cryptsetup open --type luks ${TARGET_PART}${ROOT_PART_NUM} encrypted
    mkfs -t btrfs -m single -L rootemmc /dev/mapper/encrypted
    mount -o ssd,compress-force=zstd,noatime,nodiratime /dev/mapper/encrypted /mnt

else
    mkfs -t btrfs -m single -L rootemmc -f ${TARGET_PART}${ROOT_PART_NUM}
    mount -o ssd,compress-force=zstd,noatime,nodiratime ${TARGET_PART}${ROOT_PART_NUM} /mnt


fi

partprobe $TARGET

if [ "$NO_BOOT_PARTITION" == "" ]; then
    mkdir -p /mnt/boot
    mount ${TARGET_PART}${BOOT_PART_NUM} /mnt/boot
fi

echo "[info] Copying the file system over."
rsync --info=name -axADHSX --no-inc-recursive --delete /boot/ /mnt/boot || true #sometimes files like to vanish
rsync --info=name -axADHSX --no-inc-recursive --delete --exclude='/swap/*' / /mnt || true

if [ "$USE_LUKS" == "yes" ]; then
    
    # Configure crypttab for LUKS before initramfs generation
    echo "[info] Configuring crypttab for LUKS."
    
    # Get the UUID of the encrypted partition
    ENCRYPTED_UUID=$(blkid -s UUID -o value ${TARGET_PART}${ROOT_PART_NUM})
    
    # Create/update crypttab
    echo "encrypted UUID=${ENCRYPTED_UUID} none luks" > /mnt/etc/crypttab
    
    # Set up chroot environment for initramfs generation
    echo "[info] Setting up a chroot environment."
    mount --bind /dev /mnt/dev
    mount --bind /proc /mnt/proc  
    mount --bind /sys /mnt/sys
    mount --bind /run /mnt/run
    
    get_kver
    
    echo "[info] Building the kernel partition."
    # clean initramfs to force rebuild by vtbuild
    rm  /mnt/boot/initrd.img-${kver}
    chroot /mnt /bin/bash -c "/usr/bin/vtbuild"
    
    # Clean up chroot mounts
    umount /mnt/run /mnt/sys /mnt/proc /mnt/dev

    
else
    echo "[info] Building the kernel partition."
    echo "CUSTOM_ROOT_DEVICE=\"${TARGET_PART}${ROOT_PART_NUM}\" NO_FLASH=\"yes\" CUSTOM_PREFIX=\"install-\" $(dirname $BASH_SOURCE)/vtbuild"
    CUSTOM_ROOT_DEVICE="${TARGET_PART}${ROOT_PART_NUM}" NO_FLASH="yes" CUSTOM_PREFIX="install-" $(dirname $BASH_SOURCE)/vtbuild
    
    get_kver


    rm /mnt/boot/vmlinux.kpart-initrd-* #flashing other kernel version from previous install will guarantee unbootable system
    mv /boot/install-vmlinux.kpart-initrd-${kver} /mnt/boot/vmlinux.kpart-initrd-${kver} -v

    echo "[info] Flashing the kernel partition."
    dd if=/mnt/boot/vmlinux.kpart-initrd-${kver} of=${TARGET_PART}1 bs=1024k status=progress
fi


echo "[info] Synchronizing the device."
sync

echo "[info] Adjusting the fstab."
if [ "$NO_BOOT_PARTITION" == "yes" ]; then
    echo "LABEL=rootemmc / btrfs defaults,ssd,compress-force=zstd,noatime,nodiratime 0 1" > /mnt/etc/fstab
else
    sed -i 's,bootpart,bootemmc,g;s,rootpart,rootemmc,g' /mnt/etc/fstab
fi

echo "[info] Clearing existing ssh key, machine id and random seed"
# remove the generated ssh keys so that fresh ones are generated on
# first boot for each installed image
rm -f /mnt/etc/ssh/*key*
# activate the one shot service to recreate them on first boot
mkdir -p /mnt/etc/systemd/system/multi-user.target.wants
( cd /mnt/etc/systemd/system/multi-user.target.wants ;  ln -s ../regenerate-ssh-host-keys.service . )

# delete random-seed and machine-id according to https://systemd.io/BUILDING_IMAGES/
# so that they get created unique per machine on first boot
# inspired by: https://github.com/armbian/build/pull/3774
echo "uninitialized" > /mnt/etc/machine-id
rm -f /mnt/var/lib/systemd/random-seed /mnt/var/lib/dbus/machine-id

if [ "$USE_LUKS" == "yes" ]; then
    # Update fstab to use the encrypted device
    sed -i 's,LABEL=rootemmc,/dev/mapper/encrypted,g' /mnt/etc/fstab  
fi

echo "[info] Umount the partitions."
if [ "$NO_BOOT_PARTITION" == "" ]; then
    umount /mnt/boot
fi
umount /mnt

if [ "$USE_LUKS" == "yes" ]; then
  echo "[info] Closing the LUKS device."
  cryptsetup close encrypted
fi

if [ "$USE_LUKS" == "yes" ] && [ "${SHUT}" == "" ]; then
    echo ""
    echo "[info] LUKS installation completed successfully!"
    echo "The root partition is now encrypted and the system is configured to"
    echo "unlock it at boot time. You will be prompted for your passphrase"
    echo "during boot. Make sure to remember your passphrase!"
    echo ""
fi
