#!/bin/bash

# constants
VERSION='@VERSION@'     # updated at install time
BIOSBOOT_PARTITION_SIZE_KB=1024
IMAGE_SIZE_MARGIN_KB=0  # fs size estimation is enough pessimistic
MKSQUASHFS_OPTS="-b 1M -comp xz"
FAT_OVERHEAD_PERCENT=10
FS_OVERHEAD_PERCENT=18
LVM_OVERHEAD_PERCENT=4
ONE_GIGABYTE_KB=$((1024*1024))
MAX_WORK_IMAGE_OVERHEAD_KB=$ONE_GIGABYTE_KB
DEBUG=1
DD="dd status=none"
DBSTCK_DIR="/usr/share/debootstick"

# get functions now
. $DBSTCK_DIR/scripts/create-image/functions
# and have them available in chrooted scripts
export chrooted_functions="$(
    cat $DBSTCK_DIR/scripts/create-image/functions)"

# check options
root_password_request="NO_REQUEST"
root_password_on_first_boot=0
config_grub_on_serial_line=0
system_type="live"
kernel_bootargs=""
config_hostname=""
while [ $# != 0 ] 
do
    case "$1" in
        -h|--help)
            usage_and_exit 0
        ;;
        --help-os-support)
            describe_os_support
            exit 0
        ;;
        -v|--version)
            echo "debootstick $VERSION"
            exit 0
        ;;
        --kernel-package)
            kernel_package="$2"
            shift 2
        ;;
        --config-kernel-bootargs)
            kernel_bootargs="$2"
            shift 2
        ;;
        --config-root-password-ask)
            root_password_request="ASK"
            shift
        ;;
        --config-root-password-none)
            root_password_request="NO_PASSWORD"
            shift
        ;;
        --config-root-password-first-boot)
            root_password_on_first_boot=1
            shift
        ;;
        --config-hostname)
            config_hostname="$2"
            shift 2
        ;;
        --config-grub-on-serial-line)
            config_grub_on_serial_line=1
            shift
        ;;
        --system-type)
            system_type="$2"
            shift 2
        ;;
        *)
            break
        ;;
    esac
done

# we need 2 args
if [ -z "$2" ]
then
    usage_and_exit 1
fi

# let's verify system_type variable
case "$system_type" in
    'live'|'installer')
        ;;  # ok
    *)
        echo "--system-type option value must be either 'live' or 'installer'." >&2
        exit 1
esac

# ensure we are root
if [ $(id -u) -ne 0 ]; then
    echo "debootstick should be run as root. Trying sudo..."
    exec sudo "$0" "$@"
fi

# 2nd arg is a writable file path
if [ ! -w $(dirname "$2") ] 
then
    usage_and_exit
fi

# 1st arg is a directory
if [ ! -d "$1" ]
then
    usage_and_exit
fi

# this directory should contain a system 
# file hierarchy
check_fs_hierarchy "$1" || usage_and_exit

# get and verify target architecture
target_arch=$(check_target_arch "$1")

# if we are here, command line is ok :)
if [ "$root_password_request" = "ASK" ]
then
    while true
    do
        read -s -p "Enter embedded-os root password: " passwd1
        echo
        read -s -p "Enter embedded-os root password again: " passwd2
        echo
        if [ "$passwd1" = "$passwd2" ]
        then
            echo 'OK'
            root_password_request="root:$passwd1"
            break
        else
            echo 'Sorry, passwords do not match, please retry.'
        fi
    done
fi

ORIG_TREE="$(cd "$1"; pwd)"
STICK_OS_ID=$(uuidgen | tr -d '-')
FINAL_LVM_VG="DBSTCK-$STICK_OS_ID"
DRAFT_LVM_VG="DRAFT-$STICK_OS_ID"
DBSTCK_TMPDIR=$(mktemp -du --tmpdir tmp.dbstck.XXXXX.d)
final_image_path="$2"
final_image_abspath="$(abspath "$2")"
if [ "$DEBUG" = "1" ]
then
    CHROOTED_DEBUG="--debug"
fi

final_cleanup()
{
    return_code=$1
    if [ "$1" -gt 0 ]   # if error
    then
        rm -f $final_image_abspath
    fi
}

start_failsafe_mode --toplevel final_cleanup

failsafe mkdir -p $DBSTCK_TMPDIR
cd $DBSTCK_TMPDIR

# step0: generate a UEFI bootloader binary
echo -n "I: generating a UEFI bootloader binary... "
case $target_arch in
    "amd64")
        grub_arch="x86_64-efi"
        img_name="BOOTX64.efi"
        ;;
    "i386")
        grub_arch="i386-efi"
        img_name="BOOTIA32.efi"
        ;;
esac
build_uefi_binary_image $grub_arch $img_name
efi_image_size_bytes=$(stat -c "%s" $img_name)
efi_partition_size_kb=$(apply_overhead_percent \
            $((efi_image_size_bytes/1024)) $FAT_OVERHEAD_PERCENT)
echo done

# step1: compute a stick size large enough for our work 
# (i.e. not for the final minimized version)
echo -n "I: draft image - computing a size large enough... "
fs_size_estimation_kb=$(estimated_size_kb $ORIG_TREE)
draft_stick_size_kb=$((efi_partition_size_kb +
                    BIOSBOOT_PARTITION_SIZE_KB + 
                    fs_size_estimation_kb +
                    MAX_WORK_IMAGE_OVERHEAD_KB))
echo done

# step2: create draft image structure
echo -n "I: draft image - partitioning and formatting... "
create_formatted_image \
        draft \
        $draft_stick_size_kb \
        $efi_partition_size_kb \
        $DRAFT_LVM_VG
echo done

# step3: copy original tree to work image and modify it
echo -n "I: draft image - copying filesystem tree... "
cd $draft_lvroot_mountpoint
cp -a $ORIG_TREE/* .
echo done
mkdir -p opt/debootstick
cp -a $DBSTCK_DIR/scripts/live opt/debootstick/live
cp -a $DBSTCK_DIR/scripts/create-image/chrooted-customization-draft.sh .
# we will need internet connectivity for package 
# management in the chroot. Ensure we have a valid DNS setup there.
[ -f etc/resolv.conf ] || touch etc/resolv.conf
with mount -o bind /etc/resolv.conf $PWD/etc/resolv.conf; do
    # let's start the customization
    chroot . ./chrooted-customization-draft.sh $CHROOTED_DEBUG    \
                $draft_device $root_password_request        \
                final_lvm_vg=$FINAL_LVM_VG   \
                config_grub_on_serial_line=$config_grub_on_serial_line  \
                kernel_package="\"$kernel_package\""    \
                kernel_bootargs="\"$kernel_bootargs\"" \
                config_hostname="\"$config_hostname\"" \
                target_arch="$target_arch"
done
rm ./chrooted-customization-draft.sh
cd ..

# step4: finalyse filesystem setup 
finalize_fs $draft_lvroot_mountpoint

# step5: compute minimal size of final stick
echo -n "I: final image - computing minimal image size... "
cd $DBSTCK_TMPDIR
data_size_kb=$(estimated_size_kb $draft_lvroot_mountpoint)
final_lvmpart_size_kb=$(apply_overheads_percent $data_size_kb \
        $FS_OVERHEAD_PERCENT $LVM_OVERHEAD_PERCENT)
draft_lvmpart_size_kb=$(device_size_kb $draft_lvmpart_device)
# since the draft and final images are formatted the
# same way (only the size of the lvm part differs),
# (final_stick_size_kb - final_lvmpart_size_kb)
# should be equal to
# (draft_stick_size_kb - draft_lvmpart_size_kb)
# let's compute final_stick_size_kb accordingly:
final_stick_size_kb=$((
        final_lvmpart_size_kb +
        draft_stick_size_kb -
        draft_lvmpart_size_kb
))
echo done

# step6: copy work version to the final image (with minimal size)

# prepare a final image with minimal size
echo -n "I: final image - partitioning and formatting... "
create_formatted_image \
        final \
        $final_stick_size_kb \
        $efi_partition_size_kb \
        ${FINAL_LVM_VG} \
        $final_image_abspath
echo done
echo -n "I: final image - copying content from draft image... "
cp -a $draft_lvroot_mountpoint/* $final_lvroot_mountpoint/
echo done
release_image draft     # not needed anymore

# add the dbstck.conf file
cat > $final_lvroot_mountpoint/dbstck.conf << EOF
LVM_VG=$FINAL_LVM_VG
SYSTEM_TYPE=$system_type
ASK_ROOT_PASSWORD_ON_FIRST_BOOT=$root_password_on_first_boot
EOF

# step7: install BIOS bootloader

cd $final_lvroot_mountpoint
# since the size of the filesystem mounted there is minimized,
# creating new files may cause problems.
# so we will use the directory /tmp that we mount in memory.
with mount -t tmpfs none $final_lvroot_mountpoint/tmp; do
    cp -a $DBSTCK_DIR/scripts/create-image/chrooted-customization-final.sh tmp
    chroot . tmp/chrooted-customization-final.sh $final_device
done
cd ..

# step8: setup the EFI boot partition
echo -n "I: final image - setting up the EFI boot partition... "
mkdir -p $final_efipart_mountpoint/EFI/BOOT
mv $DBSTCK_TMPDIR/*.efi $final_efipart_mountpoint/EFI/BOOT/
echo done

# step9: clean up
echo -n "I: cleaning up... "
undo_all
echo done

chmod a+rw $final_image_abspath
stick_size=$(real_size_human_readable $final_image_abspath)
echo "I: $final_image_path ready (size: ${stick_size}). "

