#!/usr/bin/env bash
set -euo pipefail

# ============
# Global variables
# ============
DISK=""
DISK_EFI=""
DISK_ROOT=""

LFS=/mnt/lfs
LFS_USER=lfs
LFS_USER_PASS=lfs

# ============
# Helper: Ask for value if global is empty
# ============
require_var() {
  local varname="$1"
  local prompt="$2"

  if [[ -z "${!varname}" ]]; then
    read -rp "$prompt: " value
    printf -v "$varname" '%s' "$value"
  fi
}

# ============
# Disk & Mount Setup
# ============

format_disk() {
  echo "WARNING: THIS WILL DESTROY ALL DATA ON THE TARGET DISK."

  # If DISK was not passed OR global DISK is empty, prompt for it
  if [[ $# -eq 1 ]]; then
    DISK="$1"
  fi

  require_var DISK "Enter target disk (e.g., /dev/sda or /dev/nvme0n1)"

  if [[ ! -b "$DISK" ]]; then
    echo "Error: $DISK is not a block device"
    exit 1
  fi

  # Decide partition naming scheme
  if [[ "$DISK" =~ (nvme|mmcblk) ]]; then
    DISK_BIOS="${DISK}p1"
    DISK_EFI="${DISK}p2"
    DISK_ROOT="${DISK}p3"
  else
    DISK_BIOS="${DISK}1"
    DISK_EFI="${DISK}2"
    DISK_ROOT="${DISK}3"
  fi

  echo "About to partition disk: $DISK"
  echo "  BIOS partition : $DISK_BIOS"
  echo "  EFI partition  : $DISK_EFI"
  echo "  ROOT partition : $DISK_ROOT"
  echo "Press Ctrl+C to abort, or Enter to continue."
  read -r _

  parted -s "$DISK" mklabel gpt

  parted -s "$DISK" mkpart primary 1MiB 3MiB
  parted -s "$DISK" set 1 bios_grub on

  parted -s "$DISK" mkpart ESP fat32 3MiB 515MiB
  parted -s "$DISK" set 2 esp on

  parted -s "$DISK" mkpart primary ext4 515MiB 100%

  partprobe "$DISK" || true
  sleep 2

  mkfs.fat -F32 "$DISK_EFI"
  mkfs.ext4 -F "$DISK_ROOT"

  echo "Done."
  echo "Created:"
  echo "  BIOS partition : $DISK_BIOS (bios_grub)"
  echo "  EFI partition  : $DISK_EFI (FAT32)"
  echo "  ROOT partition : $DISK_ROOT (ext4)"
}

mount_disk() {
  # If globals are empty, prompt the user
  require_var DISK_ROOT "Enter root partition device (e.g., /dev/sda3)"
  require_var DISK_EFI  "Enter EFI partition device (e.g., /dev/sda2)"

  mkdir -pv $LFS
  mount $DISK_ROOT $LFS

  mkdir -pv $LFS/boot/efi
  mount $DISK_EFI $LFS/boot/efi

  mkdir -v $LFS/sources
  chmod -v a+wt $LFS/sources 

  mkdir -pv $LFS/{bin,etc,lib,sbin,usr,var,tools}

  case $(uname -m) in
    x86_64) mkdir -pv $LFS/lib64 ;;
  esac
}

# ============
# Create User
# ============

create_user() {
  groupadd $LFS_USER
  useradd -s /bin/bash -g $LFS_USER -m -k /dev/null $LFS_USER
  
  # Set password for the LFS user
  echo "$LFS_USER:$LFS_USER_PASS" | chpasswd

  # Set ownership of LFS directories
  chown -v "$LFS_USER" "$LFS"/{sources,bin,etc,lib,sbin,usr,var,tools}
  case $(uname -m) in
    x86_64) chown -v "$LFS_USER" "$LFS"/lib64 ;;
  esac

  # Set users .bash_profile
  cat > /home/$LFS_USER/.bash_profile << "EOF"
exec env -i HOME=$HOME TERM=$TERM PS1='\u:\w\$ ' /bin/bash
EOF
    # Set users .bashrc
    cat > /home/$LFS_USER/.bashrc << "EOF"
set +h
umask 022
LFS=/mnt/lfs
LC_ALL=POSIX
LFS_TGT=$(uname -m)-lfs-linux-gnu
PATH=/usr/bin
if [ ! -L /bin ]; then PATH=/bin:$PATH; fi
PATH=$LFS/tools/bin:$PATH
CONFIG_SITE=$LFS/usr/share/config.site
export LFS LC_ALL LFS_TGT PATH CONFIG_SITE
EOF

}

everything() {
  # Pass through args to format_disk
  format_disk "$@"
  mount_disk "$@"
  create_user "$@"
}

menu() {
  echo "1) Everything"
  echo "2) Format Disk Only"
  echo "3) Mount Disk Only"
  echo "4) Create User"
  echo -n "Choice: "
  read -r CHOICE

  case "$CHOICE" in
    1) everything "$@";;
    2) format_disk "$@";;
    3) mount_disk "$@";;
    4) create_user "$@";;
    *) echo "Invalid choice";;
  esac
}

menu "$@"
