#!/bin/bash

set -o nounset -o errexit
readonly DEFAULT_DISK_SIZE="2G"
readonly IMAGE="image.img"
readonly MIRROR='https://fastly.mirror.pkgbuild.com/$repo/os/$arch'

function init() {
  readonly ORIG_PWD="${PWD}"
  readonly OUTPUT="${PWD}/output"
  mkdir -p "${OUTPUT}" "${PWD}/tmp"
  local tmpdir
  tmpdir="$(mktemp --directory --tmpdir="${PWD}/tmp")"
  readonly TMPDIR="${tmpdir}"
  if [ -n "${SUDO_UID:-}" ] && [[ -n "${SUDO_GID:-}" ]]; then
    chown "${SUDO_UID}:${SUDO_GID}" "${OUTPUT}" "${TMPDIR}"
  fi
  cd "${TMPDIR}"

  readonly MOUNT="${PWD}/mount"
  mkdir "${MOUNT}"
}

function cleanup() {
  set +o errexit
  if [ -n "${MOUNT:-}" ] && mountpoint -q "${MOUNT}"; then
    umount --recursive "${MOUNT}" || exit 1
  fi
  if [ -n "${LOOPDEV:-}" ]; then
    losetup -d "${LOOPDEV}"
  fi
  if [ -n "${TMPDIR:-}" ]; then
    rm -rf "${TMPDIR}"
  fi
}
trap cleanup EXIT

function setup_disk() {
  truncate -s "${DEFAULT_DISK_SIZE}" "${IMAGE}"
  sgdisk --align-end \
    --clear \
    --new 0:0:+1M --typecode=0:ef02 --change-name=0:'BIOS boot partition' \
    --new 0:0:+300M --typecode=0:ef00 --change-name=0:'EFI system partition' \
    --new 0:0:0 --typecode=0:8304 --change-name=0:'Parch Linux root' \
    "${IMAGE}"

  LOOPDEV=$(losetup --find --partscan --show "${IMAGE}")
  wait_until_settled "${LOOPDEV}"
  mkfs.fat -F 32 -S 4096 "${LOOPDEV}p2"
  mkfs.btrfs "${LOOPDEV}p3"
  mount -o compress=zstd:1 "${LOOPDEV}p3" "${MOUNT}"
  mount --mkdir "${LOOPDEV}p2" "${MOUNT}/efi"
}

function bootstrap() {
  cat <<EOF >pacman.conf
[options]
Architecture = auto
SigLevel = DatabaseOptional

[world]
SigLevel = Never
Server = https://github.com/parchlinux/world/raw/refs/heads/main/\$arch/

[core]
Include = mirrorlist

[extra]
Include = mirrorlist
EOF
  echo "Server = ${MIRROR}" >mirrorlist

  pacstrap -c -C pacman.conf -M "${MOUNT}" base linux grub openssh sudo btrfs-progs dosfstools efibootmgr qemu-guest-agent parchlinux-keyring parch-branding lsb-release
  gpgconf --homedir "${MOUNT}/etc/pacman.d/gnupg" --kill gpg-agent
  cp -a --no-preserve=ownership "${ORIG_PWD}/branding/." "${MOUNT}/"
  install -Dm 644 "${ORIG_PWD}/branding/usr/lib/os-release_parch" "${MOUNT}/usr/local/share/parch-branding/os-release"
  install -Dm 644 "${ORIG_PWD}/branding/etc/issue_parch" "${MOUNT}/usr/local/share/parch-branding/issue"
  install -Dm 644 "${ORIG_PWD}/branding/etc/lsb-release_parch" "${MOUNT}/usr/local/share/parch-branding/lsb-release"
  install -Dm 644 "${ORIG_PWD}/branding/usr/share/pixmaps/parch-logo.svg" "${MOUNT}/usr/local/share/parch-branding/pixmaps/parch-logo.svg"
  install -Dm 644 "${ORIG_PWD}/branding/usr/share/pixmaps/parch-logo-dark.svg" "${MOUNT}/usr/local/share/parch-branding/pixmaps/parch-logo-dark.svg"
}

function image_cleanup() {
  rm -rf "${MOUNT}/etc/pacman.d/gnupg/"

  printf '\nMODULES=(virtio_pci virtio_blk virtio_scsi btrfs sr_mod)\n' >>"${MOUNT}/etc/mkinitcpio.conf"
  arch-chroot "${MOUNT}" /usr/bin/mkinitcpio -P

  sync -f "${MOUNT}/etc/os-release"
  fstrim --verbose "${MOUNT}"
  fstrim --verbose "${MOUNT}/efi"
}

function wait_until_settled() {
  udevadm settle
  blockdev --flushbufs --rereadpt "${1}"
  until test -e "${1}p3"; do
    echo "${1}p3 doesn't exist yet..."
    sleep 1
  done
}

function mount_image() {
  LOOPDEV=$(losetup --find --partscan --show "${1:-${IMAGE}}")
  wait_until_settled "${LOOPDEV}"
  mount -o compress=zstd:1 "${LOOPDEV}p3" "${MOUNT}"
  mount --bind "/var/cache/pacman/pkg" "${MOUNT}/var/cache/pacman/pkg"
}

function unmount_image() {
  umount --recursive "${MOUNT}"
  losetup -d "${LOOPDEV}"
  LOOPDEV=""
}

function mv_to_output() {
  sha256sum "${1}" >"${1}.SHA256"
  if [ -n "${SUDO_UID:-}" ]; then
    chown "${SUDO_UID}:${SUDO_GID}" "${1}"{,.SHA256}
  fi
  mv "${1}"{,.SHA256} "${OUTPUT}/"
}

function create_image() {
  local tmp_image
  tmp_image="$(basename "$(mktemp -u)")"
  cp -a "${IMAGE}" "${tmp_image}"
  if [ -n "${DISK_SIZE}" ]; then
    truncate -s "${DISK_SIZE}" "${tmp_image}"
    sgdisk --align-end --delete 3 "${tmp_image}"
    sgdisk --align-end --move-second-header \
      --new 0:0:0 --typecode=0:8304 --change-name=0:'Parch Linux root' \
      "${tmp_image}"
  fi
  mount_image "${tmp_image}"
  if [ -n "${DISK_SIZE}" ]; then
    btrfs filesystem resize max "${MOUNT}"
  fi

  if [ 0 -lt "${#PACKAGES[@]}" ]; then
    arch-chroot "${MOUNT}" /usr/bin/pacman -S --noconfirm "${PACKAGES[@]}"
  fi
  if [ 0 -lt "${#SERVICES[@]}" ]; then
    arch-chroot "${MOUNT}" /usr/bin/systemctl enable "${SERVICES[@]}"
  fi
  "${2}"
  image_cleanup
  unmount_image
  "${3}" "${tmp_image}" "${1}"
  mv_to_output "${1}"
}

function main() {
  if [ "$(id -u)" -ne 0 ]; then
    echo "root is required"
    exit 1
  fi
  init

  if [ -n "${2:-}" ] && [ ! -f "${ORIG_PWD}/images/${2}.sh" ]; then
    echo "Unknown image variant '${2}'. Available variants are: $(find "${ORIG_PWD}/images" -maxdepth 1 -type f -name '*.sh' ! -name base.sh -printf '%f\n' | sed 's/\.sh$//' | tr '\n' ' ')"
    exit 1
  fi

  setup_disk
  bootstrap
  source "${ORIG_PWD}/images/base.sh"
  pre
  unmount_image

  local build_version
  if [ -z "${1:-}" ]; then
    build_version="$(date +%Y%m%d).0"
    echo "WARNING: BUILD_VERSION wasn't set!"
    echo "Falling back to $build_version"
  else
    build_version="${1}"
  fi

  for image in "${ORIG_PWD}/images/"*.sh; do
    if [ "$(basename "${image}")" = "base.sh" ]; then
      continue
    fi
    if [ -n "${2:-}" ] && [ "$(basename "${image}" .sh)" != "${2}" ]; then
      continue
    fi
    source "${image}"
    create_image "${IMAGE_NAME}" pre post
  done
}
main "$@"