#!/bin/bash
set -euo pipefail

apt update && sudo apt install -y \
  debootstrap gdisk zfsutils-linux zfs-initramfs zfs-zed \
  dosfstools curl wget \
  grub-efi-amd64 shim-signed efibootmgr \
  locales console-setup keyboard-configuration \
  sudo

### === Uživatelské proměnné === ###
DISK1="/dev/disk/by-id/ata-QEMU_HARDDISK_QM00005"                # např. /dev/sda nebo /dev/nvme0n1
DISK2="/dev/disk/by-id/ata-QEMU_HARDDISK_QM00007"                # např. /dev/sdb nebo /dev/nvme1n1
HOSTNAME="ubntzfstest"
ZFS_POOL="rpool"
ZFS_BOOTFS="ubuntu-noble"
ZBM_VERSION="2.4.0"
MOUNTPOINT="/mnt"
UBUNTU_SUITE="noble"
TIMEZONE="Europe/Prague"
LOCALE="cs_CZ.UTF-8"
KEYMAP="cz"
DEBIAN_MIRROR="http://cz.archive.ubuntu.com/ubuntu"
### ============================ ###

ZFS_OPTS_POOL="-o ashift=12 -o autotrim=on"
ZFS_OPTS_DATASETS="-O acltype=posixacl -O xattr=sa -O compression=zstd -O relatime=on -O normalization=formD -O dnodesize=auto -O atime=off -O mountpoint=none -O canmount=off"

echo "== Čištění disků =="
for DISK in "$DISK1" "$DISK2"; do
  sgdisk --zap-all "$DISK"
  dd if=/dev/zero of="$DISK" bs=512 count=2048
  wipefs -a "$DISK"
done

echo "== Vytváření GPT tabulky a oddílů =="
for DISK in "$DISK1" "$DISK2"; do
  sgdisk -n1:1M:+1G   -t1:EF00 -c1:"EFI System Partition" "$DISK"
  sgdisk -n2:0:0      -t2:BF00 -c2:"ZFS root" "$DISK"
done

PART1_DISK1=$(ls ${DISK1}* | grep -E "${DISK1}p?1$")
PART1_DISK2=$(ls ${DISK2}* | grep -E "${DISK2}p?1$")
PART2_DISK1=$(ls ${DISK1}* | grep -E "${DISK1}p?2$")
PART2_DISK2=$(ls ${DISK2}* | grep -E "${DISK2}p?2$")

echo "== Formátování EFI oddílů =="
mkfs.vfat -F32 -n EFI "$PART1_DISK1"
mkfs.vfat -F32 -n EFI "$PART1_DISK2"

echo "== Vytváření ZFS poolu ($ZFS_POOL) =="
zpool create -f $ZFS_OPTS_POOL $ZFS_OPTS_DATASETS \
  -R "$MOUNTPOINT" \
  "$ZFS_POOL" mirror "$PART2_DISK1" "$PART2_DISK2"

zfs create "$ZFS_POOL/ROOT"
zfs create -o mountpoint=/ "$ZFS_POOL/ROOT/$ZFS_BOOTFS"
zfs create -o mountpoint=/home "$ZFS_POOL/home"
zfs create -o mountpoint=/var "$ZFS_POOL/var"
zfs create -o mountpoint=/var/log "$ZFS_POOL/var/log"
zfs create -o mountpoint=/var/tmp "$ZFS_POOL/var/tmp"
zfs create -o mountpoint=/var/cache "$ZFS_POOL/var/cache"
zfs create -o mountpoint=/var/lib "$ZFS_POOL/var/lib"
zfs create -o mountpoint=/srv "$ZFS_POOL/srv"
zfs create -o mountpoint=/tmp "$ZFS_POOL/tmp"

zfs set canmount=noauto "$ZFS_POOL/ROOT/$ZFS_BOOTFS"
zfs mount "$ZFS_POOL/ROOT/$ZFS_BOOTFS"

mkdir -p "$MOUNTPOINT"/{boot/efi,home,var/log,var/tmp,var/cache,var/lib,srv,tmp}

echo "== Montování EFI oddílů =="
mount "$PART1_DISK1" "$MOUNTPOINT/boot/efi"

echo "== Bootstrap Ubuntu base systému =="
debootstrap --components=main,universe --arch=amd64 "$UBUNTU_SUITE" "$MOUNTPOINT" "$DEBIAN_MIRROR"

echo "== Příprava chroot prostředí =="
mount --rbind /dev "$MOUNTPOINT/dev"
mount --rbind /proc "$MOUNTPOINT/proc"
mount --rbind /sys "$MOUNTPOINT/sys"
cp /etc/resolv.conf "$MOUNTPOINT/etc/"

cat <<EOF > "$MOUNTPOINT/root/post-chroot.sh"
#!/bin/bash
set -e

export HOSTNAME="$HOSTNAME"
export ZFS_POOL="$ZFS_POOL"
export ZFS_BOOTFS="$ZFS_BOOTFS"
export TIMEZONE="$TIMEZONE"
export LOCALE="$LOCALE"
export KEYMAP="$KEYMAP"
export ZBM_VERSION="$ZBM_VERSION"

echo "\$HOSTNAME" > /etc/hostname
ln -sf /usr/share/zoneinfo/\$TIMEZONE /etc/localtime

apt update
apt install -y locales keyboard-configuration console-setup

sed -i "s/# \$LOCALE UTF-8/\$LOCALE UTF-8/" /etc/locale.gen
locale-gen
update-locale LANG=\$LOCALE
echo "KEYMAP=\$KEYMAP" > /etc/vconsole.conf

apt install -y linux-image-generic zfs-initramfs zfsutils-linux zfs-zed \
               dosfstools grub-efi-amd64 shim-signed efibootmgr curl

zfs set canmount=noauto \$ZFS_POOL/ROOT/\$ZFS_BOOTFS
zfs mount \$ZFS_POOL/ROOT/\$ZFS_BOOTFS
zfs set mountpoint=/ \$ZFS_POOL/ROOT/\$ZFS_BOOTFS
zpool set bootfs=\$ZFS_POOL/ROOT/\$ZFS_BOOTFS \$ZFS_POOL

useradd -m -s /bin/bash ubuntu
echo "ubuntu:ubuntu" | chpasswd
adduser ubuntu sudo
echo "⚠️ Výchozí heslo pro uživatele 'ubuntu' je 'ubuntu'. Změň ho po restartu."

echo "== Instalace GRUB na oba disky =="
grub-install --target=x86_64-efi --efi-directory=/boot/efi \
             --removable --recheck --bootloader-id=ubuntu
update-grub

echo "== Instalace ZFSBootMenu (volitelný fallback) =="
mkdir -p /boot/efi/EFI/ZBM
curl -L https://github.com/zbm-dev/zfsbootmenu/releases/download/v\$ZBM_VERSION/zfsbootmenu-efi-\$ZBM_VERSION-x86_64.EFI -o /boot/efi/EFI/ZBM/ZBM.EFI

cat <<GRUBENTRY >> /etc/grub.d/40_custom

menuentry "ZFSBootMenu" {
  insmod part_gpt
  insmod fat
  insmod chain
  search --no-floppy --label EFI --set=root
  chainloader /EFI/ZBM/ZBM.EFI
}
GRUBENTRY

update-grub

echo "== Hotovo. Nyní spusť: exit && umount a export poolu"
EOF

chmod +x "$MOUNTPOINT/root/post-chroot.sh"

echo
echo "=========================="
echo "✅ Chroot a konfigurace:"
echo "sudo chroot $MOUNTPOINT /bin/bash"
echo "/root/post-chroot.sh"
echo
echo "Poté ukonči chroot a spusť:"
echo "umount -l $MOUNTPOINT/dev $MOUNTPOINT/proc $MOUNTPOINT/sys $MOUNTPOINT/boot/efi"
echo "zpool export $ZFS_POOL"
echo "=========================="