#!/bin/sh

set -e
set -x

# shellcheck disable=SC1091
. /var/lib/rpi-sb-provisioner/manufacturing-data
# shellcheck disable=SC1091
. /var/lib/rpi-sb-provisioner/state-recording

DEBUG=

export PROVISIONER_FINISHED="NAKED-PROVISIONER-FINISHED"
export PROVISIONER_ABORTED="NAKED-PROVISIONER-ABORTED"
export PROVISIONER_STARTED="NAKED-PROVISIONER-STARTED"

# Source common helper functions
# shellcheck disable=SC1091
. "$(dirname "$0")/rpi-sb-common.sh"

die() {
    record_state "${TARGET_DEVICE_SERIAL}" "${PROVISIONER_ABORTED}" "${TARGET_USB_PATH}"
    # shellcheck disable=SC2086
    echo "$@" ${DEBUG}
    exit 1
}

log() {
    timestamp=$(date +"%Y-%m-%d %H:%M:%S.$(date +%N | cut -c1-3)")
    echo "[${timestamp}] $*" >> /var/log/rpi-sb-provisioner/"${TARGET_DEVICE_SERIAL}"/provisioner.log
    printf "[%s] %s\n" "${timestamp}" "$*"
}

read_config
compute_image_summary

TMP_DIR=""
CLEANUP_DONE=0

# check_file_is_expected ${path_to_file} ${expected_file_extension}
# Checks if a file exists, is not a directory, is not zero and has the right extension.
# If any of those checks fail, exit the script entirely and print a debug message
# If all checks succeed, supply the filepath via stdout
check_file_is_expected() {
    filepath="$1"
    ext="$2"

    if [ ! -e "${filepath}" ]; then
        die "Specified file does not exist: ${filepath}"
    fi

    if [ -d "${filepath}" ]; then
        die "Expected a file, got a directory for ${filepath}"
    fi

    if [ -z "${filepath}" ]; then
        die "Provided file is empty: ${filepath}"
    fi

    if [ -z "${ext}" ] || echo "${filepath}" | grep -q "${ext}"; then
        echo "${filepath}"
        return 0
    else
        die "Provided file is of the wrong extension, wanted ${ext}, provided ${filepath}"
    fi
}

check_command_exists() {
    command_to_test=$1
    if ! command -v "${command_to_test}" 1> /dev/null; then
        die "${command_to_test} could not be found"
    else
        echo "$command_to_test"
    fi
}

check_pidevice_storage_type() {
    case "${1}" in
        "sd")
            echo "mmcblk0"
            ;;
        "emmc")
            echo "mmcblk0"
            ;;
        "nvme")
            echo "nvme0n1"
            ;;
        *)
            die "Unexpected storage device type. Wanted sd, nvme or emmc, got '$1'"
            ;;
    esac
}

# Check if a customisation script exists and is executable
customisation_script_is_runnable() {
    SCRIPT_PATH="/etc/rpi-sb-provisioner/scripts/$1-$2.sh"
    [ -x "${SCRIPT_PATH}" ]
}

# TODO: Refactor these two functions to use the same logic, but with different consequences for failure.
timeout_nonfatal() {
    command="$*"
    set +e
    # shellcheck disable=SC2086
    timeout 10 ${command}
    command_exit_status=$?
    if [ ${command_exit_status} -eq 124 ]; then
        log "\"${command}\" failed, timed out."
    elif [ ${command_exit_status} -ne 0 ]; then
        log "\"${command}\" failed, exit status: ${command_exit_status}"
    else
        log "\"$command\" succeeded."
    fi
    set -e
    return ${command_exit_status}
}

timeout_fatal() {
    command="$*"
    set +e
    # shellcheck disable=SC2086
    timeout 120 ${command}
    command_exit_status=$?
    if [ ${command_exit_status} -eq 124 ]; then
        record_state "${TARGET_DEVICE_SERIAL}" "${PROVISIONER_ABORTED}" "${TARGET_USB_PATH}"
        die "\"${command}\" failed, timed out."
    elif [ ${command_exit_status} -ne 0 ]; then
        record_state "${TARGET_DEVICE_SERIAL}" "${PROVISIONER_ABORTED}" "${TARGET_USB_PATH}"
        die "\"$command\" failed, exit status: ${command_exit_status}"
    else
        log "\"$command\" succeeded."
    fi
    set -e
}

cleanup() {
    # Capture the exit status that triggered the trap BEFORE any other
    # command runs, otherwise $? is clobbered by the guard/assignment below
    # and a genuine failure is reported as success.
    return_value=$?

    # Guard against multiple invocations (signal + EXIT trap)
    [ "$CLEANUP_DONE" -eq 1 ] && return
    CLEANUP_DONE=1

    # Disable errexit so cleanup runs to completion even if umount/sync fail
    set +e

    # Unmount any mounted partitions from image customisation
    [ -d "${TMP_DIR}/rpi-boot-img-mount" ] && { umount "${TMP_DIR}"/rpi-boot-img-mount 2>/dev/null; sync; }
    [ -d "${TMP_DIR}/rpi-rootfs-img-mount" ] && { umount "${TMP_DIR}"/rpi-rootfs-img-mount 2>/dev/null; sync; }

    # Detach any loop devices associated with the modified image copy
    if [ -n "${TMP_DIR}" ] && [ -f "${TMP_DIR}/gold-master-modified.img" ]; then
        unmount_image "${TMP_DIR}/gold-master-modified.img" 2>/dev/null
    fi

    if [ -d "${TMP_DIR}" ]; then
        rm -rf "${TMP_DIR}"
    fi

    if [ -n "${DELETE_PRIVATE_TMPDIR}" ]; then
        announce_start "Deleting customised intermediates"
        # shellcheck disable=SC2086
        rm -rf "${RPI_SB_WORKDIR}" ${DEBUG}
        DELETE_PRIVATE_TMPDIR=
        announce_stop "Deleting customised intermediates"
    fi

    cleanup_orphans

    if [ "${return_value}" -ne 0 ]; then
        run_provision_failed_hook "naked-provisioner" "provisioning"
    fi

    exit ${return_value}
}
trap cleanup EXIT INT TERM

# Start the provisioner phase


# These tools are used to modify the supplied images, and deal with mounting and unmounting the images.
check_command_exists losetup
check_command_exists mknod
check_command_exists lsblk
check_command_exists cut
check_command_exists findmnt

# Fastboot is used as a transfer mechanism to get images and metadata to and from the Raspberry Pi device
check_command_exists fastboot

check_command_exists blockdev

check_command_exists grep

check_command_exists systemd-notify

setup_fastboot_and_id_vars "$1"

record_state "${TARGET_DEVICE_SERIAL}" "${PROVISIONER_STARTED}" "${TARGET_USB_PATH}"

# Run provision-started hook (e.g. LED control on programming rigs)
run_customisation_script "naked-provisioner" "provision-started" "${FASTBOOT_DEVICE_SPECIFIER}" "${TARGET_DEVICE_SERIAL}" "${RPI_DEVICE_STORAGE_TYPE}"

TMP_DIR=$(make_temp_dir)
RPI_DEVICE_STORAGE_TYPE="$(check_pidevice_storage_type "${RPI_DEVICE_STORAGE_TYPE}")"
DELETE_PRIVATE_TMPDIR=
announce_start "Finding the cache directory"
if [ -z "${RPI_SB_WORKDIR}" ]; then
    RPI_SB_WORKDIR=$(make_temp_dir "rpi-sb-provisioner.XXX")
    announce_stop "Finding the cache directory: Created ${RPI_SB_WORKDIR} (none configured)"
    DELETE_PRIVATE_TMPDIR="true"
elif [ ! -d "${RPI_SB_WORKDIR}" ]; then
    RPI_SB_WORKDIR=$(make_temp_dir "rpi-sb-provisioner.XXX")
    announce_stop "Finding the cache directory: Created ${RPI_SB_WORKDIR} (configured path isn't a directory)"
    DELETE_PRIVATE_TMPDIR="true"
else
    # Deliberately do nothing
    announce_stop "Finding the cache directory: Using specified name"
fi


systemd-notify --ready --status="Provisioning started"

announce_start "Writing OS images"

# Determine which image to flash. If bootfs-mounted or rootfs-mounted customisation
# scripts are present and executable, we need to mount the image, run the scripts,
# and flash the modified copy. Otherwise, flash the gold master directly.
FLASH_IMAGE="${GOLD_MASTER_OS_FILE}"

if customisation_script_is_runnable "naked-provisioner" "bootfs-mounted" || \
   customisation_script_is_runnable "naked-provisioner" "rootfs-mounted"; then

    announce_start "OS Image Customisation"

    announce_start "Copying gold master image for customisation"
    cp --reflink=auto "${GOLD_MASTER_OS_FILE}" "${TMP_DIR}"/gold-master-modified.img
    announce_stop "Copying gold master image for customisation"

    announce_start "OS Image Mounting"

    # Mount the image as a series of partitions via loop device
    cnt=0
    until ensure_next_loopdev && LOOP_DEV="$(losetup --show --find --partscan "${TMP_DIR}/gold-master-modified.img")"; do
        if [ $cnt -lt 5 ]; then
            cnt=$((cnt + 1))
            log "Error in losetup.  Retrying..."
            sleep 5
        else
            log "ERROR: losetup failed; exiting"
            die "Failed to set up loop device for image customisation"
        fi
    done

    ensure_loopdev_partitions "$LOOP_DEV"
    BOOT_DEV="${LOOP_DEV}"p1
    ROOT_DEV="${LOOP_DEV}"p2

    NEED_ROOTFS_MOUNT=0
    if customisation_script_is_runnable "naked-provisioner" "rootfs-mounted"; then
        NEED_ROOTFS_MOUNT=1
    fi

    # shellcheck disable=SC2086
    mkdir -p "${TMP_DIR}"/rpi-boot-img-mount ${DEBUG}
    # shellcheck disable=SC2086
    mkdir -p "${TMP_DIR}"/rpi-rootfs-img-mount ${DEBUG}

    # OS Images are, by convention, packed as a MBR whole-disk file,
    # containing two partitions: A FAT boot partition, which contains the kernel, command line,
    # and supporting boot infrastructure for the Raspberry Pi Device.
    # And in partition 2, the OS rootfs itself.
    # shellcheck disable=SC2086
    mount -t vfat "${BOOT_DEV}" "${TMP_DIR}"/rpi-boot-img-mount ${DEBUG}
    if [ "${NEED_ROOTFS_MOUNT}" -eq 1 ]; then
        # shellcheck disable=SC2086
        mount -t ext4 "${ROOT_DEV}" "${TMP_DIR}"/rpi-rootfs-img-mount ${DEBUG}
    fi

    announce_stop "OS Image Mounting"

    # Naked mode leaves the customer's image as-is; warn rather than modify
    # if lock_device_private_key=1 is missing, because rpi-verity-verifier
    # (if installed in the image) will abort at boot with KEY_NOT_LOCKED.
    ensure_lock_device_private_key "${TMP_DIR}/rpi-boot-img-mount/config.txt" warn

    # Run customisation script for bootfs-mounted stage
    run_customisation_script "naked-provisioner" "bootfs-mounted" "${TMP_DIR}/rpi-boot-img-mount" "${TMP_DIR}/rpi-rootfs-img-mount"

    # Run customisation script for rootfs-mounted stage
    run_customisation_script "naked-provisioner" "rootfs-mounted" "${TMP_DIR}/rpi-boot-img-mount" "${TMP_DIR}/rpi-rootfs-img-mount"

    announce_start "OS Image Unmounting"
    umount "${TMP_DIR}"/rpi-boot-img-mount
    if [ "${NEED_ROOTFS_MOUNT}" -eq 1 ]; then
        umount "${TMP_DIR}"/rpi-rootfs-img-mount
    fi
    sync; sync; sync;
    losetup -d "${LOOP_DEV}"
    announce_stop "OS Image Unmounting"

    FLASH_IMAGE="${TMP_DIR}/gold-master-modified.img"
    announce_stop "OS Image Customisation"
fi

announce_start "Erase Device Storage"
fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" erase "${RPI_DEVICE_STORAGE_TYPE}"
sleep 3
announce_stop "Erase Device Storage"

# Re-check the fastboot devices specifier, as it may take a while for a device to gain IP connectivity
setup_fastboot_and_id_vars "${FASTBOOT_DEVICE_SPECIFIER}"

# Prefer the TCP data-plane specifier when the daemon advertises split
# mode (-i usb+tcp); fall back to whatever the control plane is using.
FLASH_SPECIFIER="${FASTBOOT_TCP_FLASH_SPECIFIER:-${FASTBOOT_DEVICE_SPECIFIER}}"
fastboot -s "${FLASH_SPECIFIER}" flash "${RPI_DEVICE_STORAGE_TYPE}" "${FLASH_IMAGE}"

# If we customised the image, delete the modified copy immediately after flash.
# The bootfs-mounted/rootfs-mounted scripts may have injected per-device material
# (keys, certificates, unique configs), so the modified image must not linger on disk.
if [ "${FLASH_IMAGE}" != "${GOLD_MASTER_OS_FILE}" ]; then
    rm -f "${FLASH_IMAGE}"
    log "Deleted per-device customised image"
fi
announce_stop "Writing OS images"

announce_start "Set LED status"
fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" oem led PWR 0
announce_stop "Set LED status"

metadata_gather

# Run post-flash customisation script
run_customisation_script "naked-provisioner" "post-flash" "${FASTBOOT_DEVICE_SPECIFIER}" "${TARGET_DEVICE_SERIAL}" "${RPI_DEVICE_STORAGE_TYPE}"
log "Post-flash customization script completed"

record_state "${TARGET_DEVICE_SERIAL}" "${PROVISIONER_FINISHED}" "${TARGET_USB_PATH}"

log "Provisioning completed. Remove the device from this machine."

# Indicate successful completion to systemd
# This is used when the script is run as a systemd service
# The special exit code 0 indicates success to systemd
# Additionally, we can use systemd-notify if available to indicate completion
systemd-notify --status="Provisioning completed successfully" STOPPING=1

true
