#!/bin/bash

DONT_UMOUNT=false

[ -z "${MESSAGE_FILE}" ] && MESSAGE_FILE="/dev/stderr"
echo -n >"${MESSAGE_FILE}"

function log() {
  echo "${1}" >> "${MESSAGE_FILE}"
}

function install_grub() {
  if [ -d "/sys/firmware/efi" ]; then
    mkdir -p /target/efi
    mount -o rw "${efi_partition}" /target/efi
    chroot /target grub-install --force --removable --efi-directory=/efi --bootloader-id=GRUB "${target_disk}" || {
      log "Fail when installing GRUB on ${target_partition} with EFI, exit code: ${?}"
      exit 4
    }
    # Fallback for some VMs
    grub-install --force --bootloader-id=GRUB --boot-directory=/target/boot/ --efi-directory=/target/efi
  else
    chroot /target grub-install --force "${target_disk}" || {
      log "Fail when installing GRUB on ${target_partition} with Legacy BIOS, exit code: ${?}"
      exit 4
    }
  fi
  
  chroot /target update-grub || {
      log "Fail when updating GRUB on ${target_partition}, exit code: ${?}"
      exit 5
  }
}
 
function install_system() {
  log "Ensuring partitions are unmounted before installation..."
  umount -l "${efi_partition}"    >> "${MESSAGE_FILE}" 2>&1
  umount -l "${target_partition}" >> "${MESSAGE_FILE}" 2>&1
  umount -l "${home_partition}"   >> "${MESSAGE_FILE}" 2>&1
  umount -l "${target_disk}"      >> "${MESSAGE_FILE}" 2>&1
  umount -l /target > /dev/null 2>&1
  umount -l /target/home > /dev/null 2>&1

  log "Mounting /target..."
  mkdir -p "/target"
  mount -o rw "${target_partition}" /target || {
    log "Fail when trying to mount ${target_partition} in read-write mode in /target, exit code: ${?}"
    exit 1
  }
  
  [ -b "/${home_partition}" ] && {
    log "Mounting /target/home..."
    rm -rf "/target/home"
    mkdir -p "/target/home"
    mount -o rw "${home_partition}" /target/home || {
      log "Fail when trying to mount ${home_partition} in read-write mode in /target/home, exit code: ${?}"
      exit 2
    }
  }
 
  log "Extracting the system..."
  [ -n "${squashfs_file}" ] && {
    unsquashfs -f -d /target "${squashfs_file}" || {
      log "Fail when extracting ${squashfs_file} to /target, exit code: ${?}"
      exit 3
    }
  }
  
  log "Generating fstab"
  mount | grep "${target_partition}" | sed 's| on /target | / |;s|type ||; s|(.*| rw 0 1|'          > /target/etc/fstab
  mount | grep "${home_partition}"   | sed 's| on /target/home | /home |;s|type ||; s|(.*| rw 0 1|' >> /target/etc/fstab
 
  log "Installing GRUB"
  mount -o bind /proc /target/proc
  mount -o bind /sys  /target/sys
  mount -o bind /dev  /target/dev
  mkdir -p /target/tmp
  chmod 1777 /target/tmp

  install_grub
  
  # HOTFIX: update-grub taking wrong UUID of root partition
  #(
    cd /target/boot/grub
    wrong_uuid=$(grep -om1 "root=UUID=.*" "grub.cfg" | cut -d ' ' -f1 | cut -d '=' -f3)
    correct_uuid=$(blkid -o value -s UUID "${target_partition}")
    
    sed -i "s|${wrong_uuid}|${correct_uuid}|g" "grub.cfg"
    
    install_grub > /dev/null 2>&1
  #)
 
  log "Setting up language..."
  echo "LANG=${language}.UTF-8"   >  /target/etc/default/locale
  echo "LC_ALL=${language}.UTF-8" >> /target/etc/default/locale
  
  log "Setting up keyboard..."
  echo "KEYMAP=${keyboard_map}" > /target/etc/vconsole.conf
 
  log "Setting up timezone..."
  ln -sf "/usr/share/zoneinfo/${timezone}" /target/etc/localtime
  echo "${timezone}" > /target/etc/timezone
 
  log "Setting up hostname..."
  echo "${hostname}" > /target/etc/hostname
 
  log "Setting up user..."
  chroot /target useradd -m -s /bin/bash "${user_name}"
  chroot /target usermod -aG sudo "${user_name}"
  echo -e "${user_password}\n${user_password}" | chroot /target passwd "${user_name}"
 
  log "Running post install"
  [ -f "/target/usr/bin/post-install" ] && {
    chroot /target /usr/bin/post-install "${@}"
  }

  log "Install finished"
  
  exit 0
}

function cleanup(){
  [ "${DONT_UMOUNT}" = "true" ] && return ;

  log "Cleaning up..."
  umount -l /target/proc
  umount -l /target/sys
  umount -l /target/dev
  
  [ -d "/sys/firmware/efi" ] && {
    umount -l /target/efi && {
      rm -rf /target/efi
    }
  }
  
  [ -b "/${home_partition}" ] && {
    umount -l /target/home
  }

  umount -l "${target_partition}" >> "${MESSAGE_FILE}" 2>&1
  umount -l "${home_partition}"   >> "${MESSAGE_FILE}" 2>&1
  umount -l "${efi_partition}"    >> "${MESSAGE_FILE}" 2>&1
}

trap cleanup EXIT

for arg in "${@}"; do
  if echo "${arg}" | grep -q "^--target-partition="; then
    target_partition=$(echo "${arg}" | sed 's|^--target-partition=||')
    shift
  elif echo "${arg}" | grep -q "^--home-partition="; then
    home_partition=$(echo "${arg}" | sed 's|^--home-partition=||')
    shift
  elif echo "${arg}" | grep -q "^--user-name="; then
    user_name=$(echo "${arg}" | sed 's|^--user-name=||')
    shift
  elif echo "${arg}" | grep -q "^--user-password="; then
    user_password=$(echo "${arg}" | sed 's|^--user-password=||')
    shift
  elif echo "${arg}" | grep -q "^--squashfs="; then
    squashfs_file=$(echo "${arg}" | sed 's|^--squashfs=||')
    shift
  elif echo "${arg}" | grep -q "^--hostname="; then
    hostname=$(echo "${arg}" | sed 's|^--hostname=||')
    shift
  elif echo "${arg}" | grep -q "^--timezone="; then
    timezone=$(echo "${arg}" | sed 's|^--timezone=||')
    shift
  elif echo "${arg}" | grep -q "^--keyboard-map="; then
    keyboard_map=$(echo "${arg}" | sed 's|^--keyboard-map=||')
    shift
  elif echo "${arg}" | grep -q "^--target-disk="; then
    target_disk=$(echo "${arg}" | sed 's|^--target-disk=||')
    shift
  elif echo "${arg}" | grep -q "^--efi-partition="; then
    efi_partition=$(echo "${arg}" | sed 's|^--efi-partition=||')
    shift
  elif echo "${arg}" | grep -q "^--language="; then
    language=$(echo "${arg}" | sed 's|^--language=||')
    shift
  elif echo "${arg}" | grep -q "^--skip-umount$"; then
    DONT_UMOUNT=true
    shift
  fi
done

[ -z "${target_partition}" ] && { echo "Missing target partition";  exit 42; }
[ -z "${target_disk}" ]      && { echo "Missing GRUB target disk";  exit 42; }

[ -d "/sys/firmware/efi" ] && {
  [ -z "${efi_partition}" ]  && { echo "Missing GRUB target disk";  exit 42; }
}

[ -z "${language}" ]         && { echo "Missing language";          exit 42; }
[ -z "${user_name}" ]        && { echo "Missing user name";         exit 42; }
[ -z "${user_password}" ]    && { echo "Missing password";          exit 42; }
[ -z "${squashfs_file}" ]    && { echo "Missing system image file"; exit 42; }
[ -z "${hostname}" ]         && { echo "Missing hostname";          exit 42; }
[ -z "${timezone}" ]         && { echo "Missing time zone";         exit 42; }
[ -z "${keyboard_map}" ]     && { echo "Missing keyboard map";      exit 42; }

log "Setting up partition tables"
partx -u  "/dev/$(lsblk -no PKNAME ${target_partition})"
partprobe "/dev/$(lsblk -no PKNAME ${target_partition})"

partx -u  "/dev/$(lsblk -no PKNAME ${efi_partition})"
partprobe "/dev/$(lsblk -no PKNAME ${efi_partition})"

[ -b "/${home_partition}" ] && {
  partx -u  "/dev/$(lsblk -no PKNAME ${home_partition})"
  partprobe "/dev/$(lsblk -no PKNAME ${home_partition})"
}

install_system "${@}"