#!/usr/bin/env bash
type emergency_shell > /dev/null 2>&1 || source /lib/dracut-lib.sh
source /etc/openslx

export PATH="/usr/local/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/lib"


# hardcode dnbd device path
declare -rg _dnbd3_dev="/dev/dnbd0"

# all outputs are redirected to stderr, since this functions should
# only echo the path to the unpacked container to stdout.
container_unpack_raw() {
    local in_device="$1"
    local out_device="$(losetup -f)"
    echo "raw: Preparing to access image at '$in_device' via '$out_device'"
    if ! modprobe "loop"; then
        echo "raw: Failed to load kernel module: loop"
    fi
    if ! losetup "$out_device" "$in_device" --partscan; then
        echo "raw: Failed to attach '$in_device' to '$out_device'."
        return 1
    fi
    read_only_device="$out_device"
    return 0
}

container_unpack_xloop() {
    local in_device="$1"
    local out_device="$(xlosetup -f)"
    echo "xloop: Preparing to access image at '$in_device' via '$out_device'"
    for kmod in xloop xloop_file_fmt_qcow xloop_file_fmt_raw; do
        if ! modprobe "${kmod}"; then
            echo "xloop: Failed to load kernel module: $kmod"
        fi
    done
    if ! xlosetup -r -t QCOW "$out_device" "$in_device" --partscan; then
        echo "xloop: Failed to attach '$in_device' to '$out_device'."
        return 1
    fi
    read_only_device="$out_device"
    return 0
}
# endregion

# region connect dnbd3 image
# Determine path to dnbd3 image: either on the kcl or via config file
declare -r KCL_DNBD3_IMAGE="$(getarg slx.stage4.path=)"
if [ -n "$KCL_DNBD3_IMAGE" ]; then
    SLX_DNBD3_IMAGE="$KCL_DNBD3_IMAGE"
    echo "SLX_DNBD3_IMAGE='$SLX_DNBD3_IMAGE'" >> /etc/openslx
fi
if [ -z "$SLX_DNBD3_IMAGE" ]; then
    emergency_shell "Failed to determine which DNBD3 image to use." \
        "It was neither specified on kernel command line nor in the" \
        "configuration file."
fi
declare -r KCL_DNBD3_RID="$(getarg slx.stage4.rid=)"
if [ -n "$KCL_DNBD3_RID" ]; then
    # specified on the KCL?
    SLX_DNBD3_RID="$KCL_DNBD3_RID"
    echo "SLX_DNBD3_RID='$SLX_DNBD3_RID'" >> /etc/openslx
fi
unset _dnbd3_client_additional_args
declare -a _dnbd3_client_additional_args
if [ -n "$SLX_DNBD3_RID" ]; then
    _dnbd3_client_additional_args=("--rid" "$SLX_DNBD3_RID")
fi

if ! modprobe dnbd3; then
    echo "Failed to load kernel module: dnbd3"
fi

for try in {1..5} FINAL; do
    if [ "$try" = "FINAL" ]; then
        emergency_shell "Failed to connect '${SLX_DNBD3_IMAGE}'" \
            "(revision: ${SLX_DNBD3_RID:-0})" \
            "from one of '$SLX_DNBD3_SERVERS' to '$_dnbd3_dev'." \
            "Check if the image exists on one of the servers" \
            "and if any is reachable from this client."
    fi
    echo "dnbd3-connect: Try ($try/5), trying hosts '$SLX_DNBD3_SERVERS'"
    if dnbd3-client \
            --host "$SLX_DNBD3_SERVERS" \
            --image "$SLX_DNBD3_IMAGE" \
            --device "$_dnbd3_dev" \
            "${_dnbd3_client_additional_args[@]}"; then
        echo "dnbd3-connect: Connected"
        break
    fi
    sleep 1
done

# endregion
# region unpack dnbd3 image
if [[ $SLX_QCOW_HANDLER = xmount ]]; then
    emergency_shell "xmount support is unmaintained, broken, and has been removed. Consider using xloop"
    exit 1
fi
if ! [[ $SLX_QCOW_HANDLER =~ ^(kernel|xloop|raw)?$ ]]; then
    emergency_shell "Unsupported image handler: $SLX_QCOW_HANDLER" \
        "Use either 'raw' or 'xloop'."
fi
if [ -z "$SLX_QCOW_HANDLER" ]; then
    SLX_QCOW_HANDLER="xloop"
    echo "SLX_QCOW_HANDLER='$SLX_QCOW_HANDLER'" >> /etc/openslx
fi
if [[ $SLX_QCOW_HANDLER =~ ^(kernel|xloop)$ ]]; then
    container_unpack_xloop "$_dnbd3_dev"
fi
if [ -z "$read_only_device" ]; then
    container_unpack_raw "$_dnbd3_dev"
fi

# Fail fast if unpacking dnbd3 image failed.
if [ -z "$read_only_device" ]; then
    emergency_shell "Failed to unpack the qcow2 image!"
fi

# endregion
# region find system partition within image
if [[ "$SLX_SYSTEM_PARTITION_IDENTIFIER" =~ ^\+[0-9]+$ ]]; then
    # Partition number, e.g. +2 for second partition
    #
    num="${SLX_SYSTEM_PARTITION_IDENTIFIER#+}"
    echo "Got partition number $num from configuration, using that"
    if [ -b "${read_only_device}p${num}" ]; then
        read_only_partition="${read_only_device}p${num}"
    elif [ -b "${read_only_device}${num}" ]; then
        read_only_partition="${read_only_device}${num}"
    else
        emergency_shell "Failed to find partition $num from '$read_only_device'"
    fi
else
    # Find requested root partition by MBRID, GPT label, or GPT type - default to SLX_SYS label
    #
    identifier="${SLX_SYSTEM_PARTITION_IDENTIFIER:-SLX_SYS}"
    echo "Trying to find partition with ID or label '$identifier' from '$read_only_device'"
    declare -a parts
    parts=( "${read_only_device}"?* )
    if [ -b "${parts[0]}" ]; then
        # There is at least one partition on the device, scan
        mapfile -t parts < <( slx-tools dev_find_partitions \
            "$read_only_device" "${identifier}" )
        if [ "${#parts[@]}" = 1 ]; then
            # One match, perfect
            read_only_partition="${parts[0]}"
            echo "Found read-only partition by identifier: $read_only_partition"
        elif [ "${#parts[@]}" = 0 ]; then
            # Nothing found, scan all partitions. This will include the device itself.
            parts=( "$read_only_device"* )
        fi
        # If we found 2 or more matching partitions, they'll be probed below
    else
        # No partitions on device - check if it's a file-system
        parts=( "$read_only_device" )
    fi
fi

if [ -z "$read_only_partition" ]; then
    # Do a scan
    echo "Warning: Don't know which partition to use, trying all candidates (${parts[*]})" >&2
    p="/tmp/mounttest"
    mkdir -p "$p"
    for part in "${parts[@]}"; do
        if mount -v -o ro "$part" "$p"; then
            # See if it looks like a system partition
            if [ -x "$p/sbin/init" ] || [ -x "$p/lib/systemd/systemd" ]; then
                umount -lf "$p"
                echo "Found init on $part, will try to boot off of it."
                read_only_partition="$part"
                break
            fi
            umount -lf "$p"
        fi
    done
    echo "Warning: ***** Please fix this by either specifying SLX_SYSTEM_PARTITION_IDENTIFIER or labeling your root partition SLX_SYS *****" >&2
fi
# endregion

if [ -z "$read_only_partition" ]; then
    echo "Error: Failed to find bootable partition"
    exit 1
fi
if ! [ -b "$read_only_partition" ]; then
    echo "Error: Bootable partition $read_only_partition does not exist, or is not a block device"
    exit 1
fi

echo "Using read-only partition: $read_only_partition"
echo "$read_only_partition" > "/.read_only_device"

exit 0