#!/bin/bash
# Prerequisites check
if [ "$(id -u)" != "0" ]; then
   echo "This script must be run as root" 1>&2
   exit 1
fi

# Check if running on Debian
if [ ! -f /etc/debian_version ]; then
    echo "This script requires a Debian-based system"
    exit 1
fi

# Define variables - MODIFY THESE AS NEEDED
TARGET=/dev/mmcblk0          # Your SD card device
MOUNT_POINT=/mnt
IP_ADDRESS=192.168.0.2       # Your desired IP
GATEWAY=192.168.0.1         # Your network gateway
HOSTNAME="pve"              # Your desired hostname
TIMEZONE="Europe/London"    # Modify to your timezone

# Install required dependencies
apt update
apt install -y btrfs-progs dosfstools debootstrap qemu-user-static

# Clean the target device
echo "Cleaning target device ${TARGET}..."
dd if=/dev/zero of=${TARGET} bs=1M count=8

# Create partition table
echo "Creating partition table..."
(
echo g      # Create GPT partition table
echo n      # New partition
echo        # Default partition number (1)
echo        # Default first sector
echo +512M  # Boot partition size
echo t      # Change partition type
echo 1      # EFI System
echo n      # New partition
echo        # Default partition number (2)
echo        # Default first sector
echo        # Use remaining space
echo w      # Write changes
) | fdisk ${TARGET}

# Format partitions
echo "Formatting partitions..."
mkfs.fat -F32 ${TARGET}p1
mkfs.btrfs -f ${TARGET}p2

# Create and mount btrfs subvolumes
echo "Setting up btrfs subvolumes..."
mount -o compress=zstd ${TARGET}p2 ${MOUNT_POINT}
btrfs subvolume create ${MOUNT_POINT}/@
btrfs subvolume create ${MOUNT_POINT}/@home
btrfs subvolume create ${MOUNT_POINT}/@root
umount ${MOUNT_POINT}

# Mount all partitions
echo "Mounting filesystems..."
mount -o defaults,compress=zstd,subvol=@ ${TARGET}p2 ${MOUNT_POINT}
mkdir -p ${MOUNT_POINT}/home
mount -o defaults,compress=zstd,subvol=@home ${TARGET}p2 ${MOUNT_POINT}/home
mkdir -p ${MOUNT_POINT}/root
mount -o defaults,compress=zstd,subvol=@root ${TARGET}p2 ${MOUNT_POINT}/root
mkdir -p ${MOUNT_POINT}/boot/efi
mount ${TARGET}p1 ${MOUNT_POINT}/boot/efi

# Debootstrap
echo "Running first stage debootstrap..."
debootstrap --arch=arm64 --foreign --components=main,non-free,contrib \
    --include=locales bullseye ${MOUNT_POINT} \
    https://mirrors.tuna.tsinghua.edu.cn/debian

# Second stage
echo "Preparing for second stage debootstrap..."
cp /usr/bin/qemu-aarch64-static ${MOUNT_POINT}/usr/bin/
LANG=C.UTF-8 chroot ${MOUNT_POINT} /usr/bin/qemu-aarch64-static /bin/bash -c "/debootstrap/debootstrap --second-stage"

# Configure fstab
echo "Configuring fstab..."
cat > ${MOUNT_POINT}/etc/fstab << EOF
UUID=$(blkid -s UUID -o value ${TARGET}p2) / btrfs defaults,compress=zstd,subvol=@ 0 0
UUID=$(blkid -s UUID -o value ${TARGET}p2) /home btrfs defaults,compress=zstd,subvol=@home 0 0
UUID=$(blkid -s UUID -o value ${TARGET}p2) /root btrfs defaults,compress=zstd,subvol=@root 0 0
UUID=$(blkid -s UUID -o value ${TARGET}p1) /boot/efi vfat defaults 1 2
EOF

# Configure networking
echo "Configuring network..."
cat > ${MOUNT_POINT}/etc/network/interfaces << EOF
source /etc/network/interfaces.d/*
auto lo
iface lo inet loopback

auto eth0
iface eth0 inet static
    address ${IP_ADDRESS}/24
    gateway ${GATEWAY}
EOF

# Set hostname
echo ${HOSTNAME} > ${MOUNT_POINT}/etc/hostname
echo "${IP_ADDRESS} ${HOSTNAME} ${HOSTNAME}.local" >> ${MOUNT_POINT}/etc/hosts

# Enter chroot environment
echo "Entering chroot environment..."
for i in /dev /dev/pts /proc /sys /run; do
    mount -B $i ${MOUNT_POINT}$i
done

# Chroot operations
cat > ${MOUNT_POINT}/chroot-setup.sh << 'EOF'
#!/bin/bash

# Configure locale and timezone
sed -i -e 's/# en_US.UTF-8 UTF-8/en_US.UTF-8 UTF-8/' /etc/locale.gen
echo 'LANG="en_US.UTF-8"' > /etc/default/locale
dpkg-reconfigure -f noninteractive locales
update-locale LANG=en_US.UTF-8

# Install required packages
apt update
apt install -y grub-efi linux-image-arm64 openssh-server sudo haveged \
    zram-tools tmux bash-completion armbian-bsp-cli-rockpi-4b dialog \
    armbian-firmware linux-image-current-rockchip64 \
    linux-dtb-current-rockchip64 linux-u-boot-rockpi-4b-current \
    wpasupplicant btrfs-progs

# Install Proxmox
apt install -y proxmox-ve pve-kernel-helper ifupdown2 openvswitch-switch

# Configure bootloader
KERNEL=$(ls /boot/vmlinuz-* | sort -V | tail -n1 | sed 's/\/boot\/vmlinuz-//')
echo "Setting up boot for kernel ${KERNEL}..."

# Initialize Proxmox boot tool
proxmox-boot-tool init

# Clean up
apt clean
EOF

chmod +x ${MOUNT_POINT}/chroot-setup.sh
LANG=C.UTF-8 chroot ${MOUNT_POINT} /usr/bin/qemu-aarch64-static /bin/bash /chroot-setup.sh

# Cleanup and unmount
echo "Cleaning up..."
rm ${MOUNT_POINT}/usr/bin/qemu-aarch64-static
rm ${MOUNT_POINT}/chroot-setup.sh
sync
umount -R ${MOUNT_POINT}

echo "Installation complete. Please configure U-Boot as follows:"
cat << 'EOF'
In U-Boot console:
env default -f -a
setenv fdtover_addr_r 0x01fc0000
setenv fdtoverfile rockchip/overlay/rockchip-spi-jedec-nor.dtbo
saveenv
load mmc 0:1 $fdt_addr_r /dtb/$fdtfile
load mmc 0:1 $fdtover_addr_r /dtb/$fdtoverfile
load mmc 0:1 $kernel_addr_r /EFI/BOOT/BOOTAA64.EFI
fdt addr $fdt_addr_r
fdt resize 80000
fdt apply $fdtover_addr_r
bootefi $kernel_addr_r $fdt_addr_r
EOF