#! /usr/bin/env bash
#
# Copyright 2019 Northern.tech AS
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

# Default that can be overridden by providing this method in a
# configuration file passed with '--config'
function platform_package() {
  true
}

function trap_exit() {
  echo "mender-convert-package has finished. Cleaning up..."
  sudo umount -f work/boot > /dev/null
  sudo umount -f work/rootfs > /dev/null
}

function trap_term() {
  true
}

trap trap_term INT TERM
trap trap_exit EXIT

echo "Running $(basename $0): $@"

source modules/bootstrap.sh
source modules/disk.sh

# The mender_convert_config is always used and provides all the defaults
declare -a configs=("configs/mender_convert_config")

while (( "$#" )); do
  case "$1" in
    -o | --overlay)
      overlays+="${2}"
      shift 2
      ;;
    -c | --config)
      configs+=("${2}")
      shift 2
      ;;
    -d | --disk-image)
      disk_image="${2}"
      shift 2
      ;;
    *)
      log_fatal "Sorry but the provided option is not supported: $1"
      ;;
  esac
done

# Note the use of %q formatting here. This is a bash feature to add
# proper quoting to the strings so that spaces and special characters
# will be treated properly.  Primarily for supporting spaces in
# pathnames and avoid splitting those into multiple parameters.
source modules/config.sh $(printf "%q " "${configs[@]}")

PARTED="/usr/bin/parted"
if [ ! -f ${PARTED} ]; then
  PARTED="/sbin/parted"
fi

output_dir=work

boot_part=$(disk_boot_part)
root_part=$(disk_root_part)

# Final output
mkdir -p deploy

# Create mount points
mkdir -p work/boot
mkdir -p work/rootfs

sudo mount ${boot_part} work/boot
sudo mount ${root_part} work/rootfs

# Convert to 512 blocks
disk_image_total_sectors=$(disk_mb_to_sectors ${MENDER_STORAGE_TOTAL_SIZE_MB})
boot_part_sectors=$(disk_mb_to_sectors ${MENDER_BOOT_PART_SIZE_MB})
data_part_sectors=$(disk_mb_to_sectors ${MENDER_DATA_PART_SIZE_MB})
alignment_sectors=$((${MENDER_PARTITION_ALIGNMENT} / 512))

if [ "${MENDER_GRUB_EFI_INTEGRATION}" == "y" ]; then
  boot_part_start=${alignment_sectors}
  overhead_sectors=$(( ${alignment_sectors} * 4 ))
else
  # The two first blocks are typically reserved for U-boot environment.
  boot_part_start=$(( ${alignment_sectors} * 3))
  overhead_sectors=$(( ${alignment_sectors} * 6 ))
fi

# Validate boot-gap.bin size, ensuring it will fit between MBR and boot_part_start
if [ "${MENDER_COPY_BOOT_GAP}" == "y" ]; then
  # Add one block for MBR
  boot_gap_sectors=$(( $(stat --printf="%b" work/boot-gap.bin) + 1))
  if [ ${boot_gap_sectors} -ge ${boot_part_start} ]; then
    log_warn "The work/boot-gap.bin file will overwrite the boot partition"
    log_fatal "Please increase MENDER_PARTITION_ALIGNMENT (2x will typically solve this))"
  fi
fi

# Validate boot part size
actual_boot_size=$(du --apparent-size -s --block-size=512 ${boot_part} | cut -f 1)
if [ ${actual_boot_size} -gt ${boot_part_sectors} ]; then
  log_warn "The allocated boot part size $(disk_sectors_to_mb ${boot_part_sectors}) MiB is too small."
  log_warn "The actual boot part size is $(disk_sectors_to_mb ${actual_boot_size}) MiB"
  log_warn "Will adjust MENDER_BOOT_PART_SIZE_MB automatically"
  log_warn "Considered adjusting the configuration file to avoid this message"
  boot_part_sectors=${actual_boot_size}
fi

# Make sure that the boot part sector size is aligned to
# MENDER_PARTITION_ALIGNMENT, it is possible that the boot part is extracted
# as-is from the input image, and the input image might not have used the same
# alignment as our configuration
boot_part_sectors=$(disk_align_sectors ${boot_part_sectors} ${MENDER_PARTITION_ALIGNMENT} )

# Calculate rootfs size
rootfs_part_sectors=$(((${disk_image_total_sectors} - ${data_part_sectors} - \
  ${boot_part_sectors} - ${overhead_sectors}) / 2))

# Make sure rootfs size is aligned to MENDER_PARTITION_ALIGNMENT
rootfs_part_sectors=$(disk_align_sectors ${rootfs_part_sectors} ${MENDER_PARTITION_ALIGNMENT} )

device_type=$(cat work/rootfs/data/mender/device_type | sed 's/[^=]*=//')
artifact_name=$(cat work/rootfs/etc/mender/artifact_info | sed 's/[^=]*=//')

image_name="${device_type}-${artifact_name}"

actual_rootfs_size=$(sudo du -s --block-size=512 work/rootfs | cut -f 1)

# KiB -> 512 sectors
image_rootfs_size_sectors=$((${IMAGE_ROOTFS_SIZE} * 2))

if [ "${IMAGE_ROOTFS_SIZE}" -eq -1 ]; then
  rootfs_image_sectors="${rootfs_part_sectors}"
elif [ "${image_rootfs_size_sectors}" -lt "${actual_rootfs_size}" ]; then
  rootfs_image_sectors="${actual_rootfs_size}"
else
  rootfs_image_sectors="${image_rootfs_size_sectors}"
fi

rootfs_image_sectors=$((${rootfs_image_sectors} + ${IMAGE_ROOTFS_EXTRA_SPACE} * 2))

rootfs_image_sectors_overhead=$(awk -v r1="$actual_rootfs_size" "BEGIN{printf \"%.0f\", r1 * ${IMAGE_OVERHEAD_FACTOR}}")
if [ "${rootfs_image_sectors_overhead}" -gt "${rootfs_image_sectors}" ]; then
  rootfs_image_sectors="${rootfs_image_sectors_overhead}"
fi

if [ ${rootfs_image_sectors} -gt ${rootfs_part_sectors} ]; then
  log_warn "The calculated rootfs partition size $(disk_sectors_to_mb ${rootfs_part_sectors}) MiB is too small."
  log_warn "The actual rootfs image size is $(disk_sectors_to_mb ${rootfs_image_sectors}) MiB"
  log_fatal "You can try adjusting the MENDER_STORAGE_TOTAL_SIZE_MB variable to increase available space, or modify one of the variables IMAGE_ROOTFS_SIZE, IMAGE_ROOTFS_EXTRA_SPACE or IMAGE_OVERHEAD_FACTOR to reduce the size of the root filesystem."
fi

log_info "Rootfs filesystem size will be $(disk_sectors_to_mb ${rootfs_image_sectors}) MiB"

# Extract file-system type from rootfs
if file ${root_part} | grep -q ext4; then
  image_fs_type="ext4"
elif file ${root_part} | grep -q XFS; then
  image_fs_type="xfs"
else
  log_warn "$(file work/${root_part})"
  log_fatal "Could not determinate root file-system type. Aborting..."
fi

disk_create_file_system_from_folder "work/rootfs/data/" "work/data.img" \
  "${data_part_sectors}" "${image_fs_type}"

# Clear this area as it will be contained in the data.img
sudo rm -rf work/rootfs/data/*

disk_create_file_system_from_folder "work/rootfs/" "work/rootfs.img" \
  "${rootfs_image_sectors}" "${image_fs_type}"

log_info "Copying root filesystem image to deploy directory"
run_and_log_cmd "cp --sparse=always work/rootfs.img deploy/${image_name}.${image_fs_type}"

mender_create_artifact "${device_type}" "${artifact_name}"

log_info "Creating Mender compatible disk-image"

sdimg_path=deploy/${image_name}.sdimg

log_info "Total disk size: $(disk_sectors_to_mb ${disk_image_total_sectors}) MiB"
log_info "  Boot partition    $(disk_sectors_to_mb ${boot_part_sectors}) MiB"
log_info "  RootFS partitions $(disk_sectors_to_mb ${rootfs_part_sectors}) MiB x 2"
log_info "  Data partition    $(disk_sectors_to_mb ${data_part_sectors}) MiB"

# Initialize sdcard image file
run_and_log_cmd \
  "dd if=/dev/zero of=${sdimg_path} bs=512 count=0 seek=${disk_image_total_sectors} status=none"

# boot_part_start, is defined at the beginning of this file
boot_part_end=$(( ${boot_part_start} + ${boot_part_sectors} - 1 ))

rootfsa_start=$(disk_align_sectors ${boot_part_end} ${MENDER_PARTITION_ALIGNMENT} )
rootfsa_end=$(( ${rootfsa_start} + ${rootfs_part_sectors} - 1 ))

rootfsb_start=$(disk_align_sectors ${rootfsa_end} ${MENDER_PARTITION_ALIGNMENT} )
rootfsb_end=$(( ${rootfsb_start} + ${rootfs_part_sectors} - 1 ))

data_start=$(disk_align_sectors ${rootfsb_end} ${MENDER_PARTITION_ALIGNMENT} )
data_end=$(( ${data_start} + ${data_part_sectors} - 1 ))

# Create partition table. TODO: GPT support
run_and_log_cmd "${PARTED} -s ${sdimg_path} mklabel msdos"
run_and_log_cmd "${PARTED} -s ${sdimg_path} unit s mkpart primary fat32 ${boot_part_start} ${boot_part_end}"
run_and_log_cmd "${PARTED} -s ${sdimg_path} set 1 boot on"
run_and_log_cmd "${PARTED} -s ${sdimg_path} -- unit s mkpart primary ext2 ${rootfsa_start} ${rootfsa_end}"
run_and_log_cmd "${PARTED} -s ${sdimg_path} -- unit s mkpart primary ext2 ${rootfsb_start} ${rootfsb_end}"
run_and_log_cmd "${PARTED} -s ${sdimg_path} -- unit s mkpart primary ext2 ${data_start} ${data_end}"
run_and_log_cmd "${PARTED} -s ${sdimg_path} print"

# Write boot-gap
if [ "${MENDER_COPY_BOOT_GAP}" == "y" ]; then
  log_info "Writing boot gap of size: ${boot_part_sectors} (sectors)"
  disk_write_at_offset "${output_dir}/boot-gap.bin" "${sdimg_path}" "1"
fi

# Burn Partitions
disk_write_at_offset "${boot_part}" "${sdimg_path}" "${boot_part_start}"
disk_write_at_offset "${output_dir}/rootfs.img" "${sdimg_path}" "${rootfsa_start}"
disk_write_at_offset "${output_dir}/rootfs.img" "${sdimg_path}" "${rootfsb_start}"
disk_write_at_offset "${output_dir}/data.img" "${sdimg_path}" "${data_start}"

log_info "Performing platform specific package operations (if any)"
platform_package

# Create bmap index
if [ "${MENDER_USE_BMAP}" == "y" ]; then
  BMAP_TOOL="/usr/bin/bmaptool"
  if [ ! -e "${BMAP_TOOL}" ]; then
    log_error "You have enabled the MENDER_USE_BMAP option, but we could not find the required 'bmaptool'"
    log_fatal "You can install 'bmaptool' with: apt-get install bmap-tools (on Debian based distributions)"
  fi
  run_and_log_cmd "${BMAP_TOOL} create ${sdimg_path} > ${sdimg_path}.bmap"
fi

case "${MENDER_COMPRESS_DISK_IMAGE}" in
  gzip)
    log_info "Compressing ${sdimg_path}.gz"
    run_and_log_cmd "pigz --best --force ${sdimg_path}"
    ;;
  lzma)
    log_info "Compressing ${sdimg_path}.xz"
    run_and_log_cmd "pxz --best --force ${sdimg_path}"
    ;;
  none)
    :
    ;;
  *)
    log_fatal "Unknown MENDER_COMPRESS_DISK_IMAGE value: ${MENDER_COMPRESS_DISK_IMAGE}"
    ;;
esac

log_info "Conversion has completed! \o/"

##################### Create configuration file for tests ######################

if [ "${MENDER_GRUB_EFI_INTEGRATION}" == "y" ]; then
  boot_part_mountpoint="/boot/efi"

  # This is the name of the DISTRO_FEATURES in Yocto
  distro_feature="mender-grub"
else
  boot_part_mountpoint="/uboot"

  # This is the name of the DISTRO_FEATURES in Yocto
  distro_feature="mender-uboot"
fi
distro_feature="${distro_feature} mender-convert"

cat <<- EOF > deploy/${image_name}.cfg
MENDER_BOOT_PART="${MENDER_STORAGE_DEVICE}${MENDER_BOOT_PART_INDEX}"
MENDER_ROOTFS_PART_A="${MENDER_STORAGE_DEVICE}${MENDER_ROOTFS_PART_A_INDEX}"
MENDER_ROOTFS_PART_B="${MENDER_STORAGE_DEVICE}${MENDER_ROOTFS_PART_B_INDEX}"
MENDER_BOOT_PART_MOUNT_LOCATION="${boot_part_mountpoint}"
MENDER_BOOT_PART_SIZE_MB="$(disk_sectors_to_mb ${boot_part_sectors})"
MENDER_DATA_PART_SIZE_MB="${MENDER_DATA_PART_SIZE_MB}"
MENDER_DEVICE_TYPE="${device_type}"
MENDER_PARTITIONING_OVERHEAD_KB="$(( (${overhead_sectors} * 512) / 1024 ))"
MENDER_PARTITION_ALIGNMENT="${MENDER_PARTITION_ALIGNMENT}"
MENDER_STORAGE_DEVICE="${MENDER_STORAGE_DEVICE}"
MENDER_STORAGE_TOTAL_SIZE_MB="${MENDER_STORAGE_TOTAL_SIZE_MB}"
MENDER_UBOOT_ENV_STORAGE_DEVICE_OFFSET="12582912"
MENDER_ARTIFACT_NAME="${artifact_name}"
DISTRO_FEATURES="${distro_feature}"
DEPLOY_DIR_IMAGE="${PWD}/deploy"
MENDER_MACHINE="${device_type}"
EOF

# Something that the tests expect to be defined (originally from Yocto)
cat <<- EOF >> deploy/${image_name}.cfg
IMAGE_FSTYPES="${image_fs_type} mender sdimg"
ARTIFACTIMG_FSTYPE="${image_fs_type}"
EOF