#!/usr/bin/env bash

# === Motivation ===
# Ensure the system can boot with encrypted root after installation.

# === Problem Statement ===
# Post-install steps must align initramfs, bootloader, and encryption metadata.

# === Scope ===
# In scope: required configuration updates inside the installed system.
# Out of scope: package selection and user account management.

# === Concepts ===
# Initramfs: early boot image that unlocks encrypted storage.
# Bootloader config: entries that point to the encrypted root.
# Dracut modules: crypt (LUKS unlock), resume (hibernation support if swap file used).
# Crypttab: /etc/crypttab maps LUKS UUIDs to device names for initramfs.
# GRUB cryptodisk: GRUB_ENABLE_CRYPTODISK=y in /etc/default/grub enables LUKS unlock.
# Kernel parameters: rd.luks.uuid=<UUID> tells dracut which LUKS container to unlock.

# === Decisions ===
# Keep post-install steps explicit and minimal, focused on boot viability.
# Default to GRUB because it is widely supported on Void without systemd dependencies.
# Regenerate initramfs for all installed kernels to avoid boot drift.
# Use UUID-based references in /etc/crypttab for stability across device naming changes.
# Enable dracut crypt module explicitly in /etc/dracut.conf.d/10-crypt.conf.
# Set GRUB_ENABLE_CRYPTODISK=y in /etc/default/grub to allow GRUB to unlock LUKS.
# Add rd.luks.uuid=<UUID> to GRUB_CMDLINE_LINUX for initramfs LUKS unlock.
# Run grub-install to embed cryptodisk support, then grub-mkconfig to update menu entries.

# === Alternatives Considered ===
# Skipping post-install updates rejected because it risks unbootable systems.
# systemd-boot and rEFInd rejected as defaults due to availability and scope constraints.
# EFISTUB rejected as default because it increases manual UEFI entry management.

# === Constraints ===
# Steps must run in the target system context.

# === Open Questions ===
# Should we generate a rescue initramfs in addition to the default one?
# Should we verify GRUB can unlock LUKS2 before rebooting, or trust the configuration?
# How do we handle future kernel updates - should we document the dracut reconfiguration process?
# Should swap file activation be configured in this phase, or deferred to first boot?

# === Success Criteria ===
# After reboot, the system prompts for decryption and boots successfully.

postinstall_bind_mounts() {
  mkdir -p "$MOUNT_ROOT/dev" "$MOUNT_ROOT/proc" "$MOUNT_ROOT/sys" "$MOUNT_ROOT/run"
  mount --rbind /dev "$MOUNT_ROOT/dev"
  mount --make-rslave "$MOUNT_ROOT/dev"
  mount -t proc /proc "$MOUNT_ROOT/proc"
  mount --rbind /sys "$MOUNT_ROOT/sys"
  mount --make-rslave "$MOUNT_ROOT/sys"
  mount --rbind /run "$MOUNT_ROOT/run"
  mount --make-rslave "$MOUNT_ROOT/run"
}

postinstall_unbind_mounts() {
  umount -R "$MOUNT_ROOT/dev" 2>/dev/null || true
  umount -R "$MOUNT_ROOT/proc" 2>/dev/null || true
  umount -R "$MOUNT_ROOT/sys" 2>/dev/null || true
  umount -R "$MOUNT_ROOT/run" 2>/dev/null || true
}

postinstall_run() {
  : "${ROOT_PART:?Root partition is required}"
  : "${CRYPT_NAME:?Crypt mapping name is required}"
  : "${MOUNT_ROOT:?Mount root is required}"
  : "${FS_TYPE:?Filesystem type is required}"
  : "${SWAP_SIZE:?Swap size is required}"
  : "${HOSTNAME:?Hostname is required}"
  local luks_uuid
  local root_uuid

  luks_uuid="$(cryptsetup luksUUID "$ROOT_PART")"
  root_uuid="$(blkid -s UUID -o value "/dev/mapper/${CRYPT_NAME}")"

  log_info "Post-install: configuring target system in chroot"
  postinstall_bind_mounts

  LUKS_UUID="$luks_uuid" ROOT_UUID="$root_uuid" FS_TYPE="$FS_TYPE" SWAP_SIZE="$SWAP_SIZE" HOSTNAME="$HOSTNAME" \
    chroot "$MOUNT_ROOT" /bin/bash -s <<'EOS'
set -euo pipefail

# Ensure required tools are available inside the target system.
for cmd in grub-install grub-mkconfig xbps-reconfigure fallocate mkswap; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "Missing required command inside chroot: $cmd" >&2
    exit 1
  fi
done

# Write the hostname explicitly to match wrapper configuration.
printf '%s\n' "$HOSTNAME" > /etc/hostname

# Ensure crypttab exists with the correct UUID mapping.
mkdir -p /etc
printf 'cryptroot UUID=%s none luks\n' "$LUKS_UUID" > /etc/crypttab

# Ensure dracut loads crypt and optional resume modules.
mkdir -p /etc/dracut.conf.d
if [[ "$SWAP_SIZE" != "0" ]]; then
  printf 'add_dracutmodules+=" crypt resume "\n' > /etc/dracut.conf.d/10-crypt.conf
else
  printf 'add_dracutmodules+=" crypt "\n' > /etc/dracut.conf.d/10-crypt.conf
fi

# Prepare a swap file inside the encrypted root if enabled.
if [[ "$SWAP_SIZE" != "0" ]]; then
  mkdir -p /swap
  if [[ "$FS_TYPE" == "btrfs" ]]; then
    if command -v chattr >/dev/null 2>&1; then
      chattr +C /swap
    fi
    if command -v btrfs >/dev/null 2>&1; then
      btrfs property set /swap compression none || true
    fi
  fi
  fallocate -l "$SWAP_SIZE" /swap/swapfile
  chmod 600 /swap/swapfile
  mkswap /swap/swapfile
  if ! grep -q '^/swap/swapfile' /etc/fstab; then
    printf '/swap/swapfile none swap defaults 0 0\n' >> /etc/fstab
  fi
fi

# Ensure btrfs subvolumes are reflected in /etc/fstab.
if [[ "$FS_TYPE" == "btrfs" ]]; then
  touch /etc/fstab

  update_fstab_entry() {
    local mount_point="$1"
    local options="$2"
    local line="UUID=$ROOT_UUID $mount_point btrfs $options 0 0"

    if awk -v mp="$mount_point" '$2==mp {found=1} END {exit found?0:1}' /etc/fstab; then
      awk -v mp="$mount_point" -v line="$line" 'BEGIN{OFS=" "} $2==mp {$0=line} {print}' /etc/fstab > /etc/fstab.tmp
      mv /etc/fstab.tmp /etc/fstab
    else
      printf '%s\n' "$line" >> /etc/fstab
    fi
  }

  update_fstab_entry / "defaults,subvol=@"
  update_fstab_entry /home "defaults,subvol=@home"
  update_fstab_entry /var "defaults,subvol=@var"
  update_fstab_entry /var/log "defaults,subvol=@log"
  update_fstab_entry /.snapshots "defaults,subvol=@snapshots"
  update_fstab_entry /swap "defaults,subvol=@swap"
fi

# Ensure GRUB unlocks the encrypted root and passes the LUKS UUID to initramfs.
mkdir -p /etc/default
touch /etc/default/grub
if grep -q '^GRUB_ENABLE_CRYPTODISK=' /etc/default/grub; then
  sed -i 's/^GRUB_ENABLE_CRYPTODISK=.*/GRUB_ENABLE_CRYPTODISK=y/' /etc/default/grub
else
  printf 'GRUB_ENABLE_CRYPTODISK=y\n' >> /etc/default/grub
fi

if grep -q '^GRUB_CMDLINE_LINUX=' /etc/default/grub; then
  if ! grep -q "rd.luks.uuid=$LUKS_UUID" /etc/default/grub; then
    sed -i "s/^GRUB_CMDLINE_LINUX=\"/GRUB_CMDLINE_LINUX=\"rd.luks.uuid=$LUKS_UUID /" /etc/default/grub
  fi
else
  printf 'GRUB_CMDLINE_LINUX="rd.luks.uuid=%s"\n' "$LUKS_UUID" >> /etc/default/grub
fi

# Regenerate initramfs and GRUB configuration.
xbps-reconfigure -fa

grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=Void --recheck
grub-mkconfig -o /boot/grub/grub.cfg
EOS

  postinstall_unbind_mounts
  log_info "Post-install complete."
}