#!/usr/bin/env bash
#=====================================================================
#  File:   ocs-live-gen-ubrd (a.k.a: ocs-live-u-bootable-rawdisk)
#  Author: Ceasar Sun
#  Co-Translator: Jouni Järvinen (@rautamiekka)
#  Desc:   Merge OCS zip and U-boot enabled bootable raw image as a new U-boot enabled OCS live raw disk.
#          The output can then be used with the 'dd' command to create a U-boot enabled bootable microSD for RISC-V64 machine.
#          E.g. :DC-ROMA L2A,... 
#=====================================================================

set -euo pipefail   # Strict mode: exit on error, error on unset variable

#------------------------------------#
#  Variables and default values      #
#------------------------------------#
# Download U-boot enabled template via URL

DEFAULT_BOOTABLE_RAW_TEMPLATE="http://free.nchc.org.tw/clonezilla-riscv/DC-ROMA-L2A/bootable-img.template/bootable.ocs-latest.img"
DEFAULT_BOOTABLE_RAW_CHECKSUM="http://free.nchc.org.tw/clonezilla-riscv/DC-ROMA-L2A/bootable-img.template/MD5SUM.txt"
#DEFAULT_BOOTABLE_RAW_TEMPLATE="http://localhost:8080/clonezilla-riscv/DC-ROMA-L2A/bootable-img.template/bootable.ocs-latest.img"
#DEFAULT_BOOTABLE_RAW_CHECKSUM="http://localhost:8080/clonezilla-riscv/DC-ROMA-L2A/bootable-img.template/MD5SUM.txt"

# Temp directory (PID as name to avoid conflict)
TMPDIR_ROOT="${PWD}/ocs_gen-ubrd_$$"
MOUNT_POINT="${TMPDIR_ROOT}/mnt"

# TEXT for searching partition name in bootable template image
# "rootfs" is referred to in the DC-ROMA L2A startup procedure :
#   -  https://bianbu-linux.spacemit.com/en/device/boot/
BOOTFS_PARTLABEL="bootfs"
ENV_PARTLABEL="env"

# Runtime Variables / Control flag
OCS_ZIP=""
LOOP_DEVICE=""
BOOTFS_KPARTX_MAP=""
ENV_KPARTX_MAP=""
DOWNLOADED_TEMPLATE=""
OUTPUT_RAWDISK=""
FULL_SELF_CLI="$BASH_SOURCE$(printf " %q" "$@")"
SOFT_ITEAMS="vmlinuz initrd.img dtbs repack-raw_info.txt"

DEBUG_MODE=false
VERBOSE_MODE=
THIS_VERSION="1.0.2-1"
OCS_U_BOOT_ENV_VERSION=""

ANS_IF_OVERWRITE=""

#---------------------------#
#  Functions                #
#---------------------------#

# 1.  Super privilege check ----------
require_root() {
    #$EUID reliable only in bash; sh may use $UID
    if [[ "${EUID:-$(id -u)}" -ne 0 ]]; then
        cat >&2 <<'EOF'
=====================================================================
Warning： Root or sudo privilege required ; e.g：
  sudo $0 clonezilla-live-3.2.2-5-riscv64.zip
=====================================================================
EOF
        exit 1
    fi
}

# 2. Usage / die ----------
usage() {
    cat <<EOF

Usage: $(basename $0) [OPTION]... OCS_ZIP_FILE ...
Merge OCS zip and U-boot enabled bootable raw image as a new U-boot enabled OCS live raw disk.

Mandatory arguments to long options are mandatory for short options too.
   OCS_ZIP_FILE                     Required, OCS zip file
   -bt|--bootable_template=FILE     OCS bootable template raw image。
                                    Script will use '${DEFAULT_BOOTABLE_RAW_TEMPLATE}' as default template if none is assigned. 
   -o|--output=FILE                 File name of output。
   -d|--debug                       Debug mode : to keep TMPDIR_ROOT
   -v|--verbose                     Verbose mode
   -h|--help                        Show help

Example:
 ~$ sudo $(basename $0) clonezilla-live-3.2.2-5-riscv64.zip	
 ~$ sudo $(basename $0) my_repack_ocs.zip --bootable_template /path/to/boot.raw --debug
 ~$ sudo $(basename $0) my_repack_ocs.zip -d -v -bt /path/to/boot.raw -o /path/to/my-output.img
EOF
}

die() {
    echo "Error: $*" >&2
    exit 1
}

# 3. cleanup
# Cleanup：umount, remove loop device、kpartx mapper ,temp dir and downloaded files
cleanup() {

    STATUS=$?
    if [[ $STATUS -eq 0 ]]; then
      # echo "Normal exit, skip cleanup"
      return 0 
    fi

    echo "=== Cleanup start ==="
    # 1. Umoun raw partition
    if mountpoint -q "${MOUNT_POINT}" 2>/dev/null; then
        echo "Unmounting ${MOUNT_POINT} ..."
        umount ${VERBOSE_MODE} -l "${MOUNT_POINT}" || echo "Warning: umount failed"
    fi

    # 2. Delete kpartx mapper
    if [[ -n "${BOOTFS_KPARTX_MAP}" ]]; then
        echo "Removing kpartx map ${BOOTFS_KPARTX_MAP} ..."
        kpartx ${VERBOSE_MODE} -d "${LOOP_DEVICE}" || echo "Warning: kpartx -d failed"
        BOOTFS_KPARTX_MAP=""
    fi

    # 3. Release loop device
    if [[ -n "${LOOP_DEVICE}" ]]; then
        echo "Detaching loop device ${LOOP_DEVICE} ..."
        losetup ${VERBOSE_MODE} -d "${LOOP_DEVICE}" || echo "Warning: losetup -d failed"
        LOOP_DEVICE=""
    fi

    # 4. Delete temp directory
    # Keep or remove temp dir depending on --debug
    if $DEBUG_MODE; then
        echo "Debug mode : to keep temp dir : ${TMPDIR_ROOT}"
    else
        if [[ -d "${TMPDIR_ROOT}" ]]; then
            echo "Removing temporary directory ${TMPDIR_ROOT} ..."
            rm ${VERBOSE_MODE} -rf "${TMPDIR_ROOT}"
        fi
    fi

    echo "=== Cleanup finished ==="
}

# trap: ensure cleanup runs on script exit, Ctrl+C, or termination signal
trap cleanup EXIT INT TERM

#---------------------------#
#  Argument parsing         #
#---------------------------#
if [[ $# -eq 0 ]]; then
    usage
fi

while [[ $# -gt 0 ]]; do
    case "$1" in
        -bt|--bootable_template)
            BOOTABLE_TEMPLATE_RAWDISK="${2:-}"
            [[ -z "${BOOTABLE_TEMPLATE_RAWDISK}" ]] && die "Missing argument for --bootable_template"
            [[ ! -e "${BOOTABLE_TEMPLATE_RAWDISK}" ]] && die "Template file : '${BOOTABLE_TEMPLATE_RAWDISK}' doesn't existed !"
            shift 2
            ;;
        -o|--output)
            OUTPUT_RAWDISK="${2:-}"
            [[ -z "${OUTPUT_RAWDISK}" ]] && die "Missing argument for --output"
            shift 2
            ;;
        -d|--debug)
            DEBUG_MODE=true
            shift
            ;;
        -v|--verbose)
            VERBOSE_MODE="-v"
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        --*)  # Unknown long option
            die "Unknown argument: $1"
            ;;
        *)   # default required arg (first non-option is OCS zip)
            if [[ -z "${OCS_ZIP}" ]]; then
                OCS_ZIP="$1"
                [[ ! -e "${OCS_ZIP}" ]] && die "OCS zip file : '${OCS_ZIP}' doesn't existed !"
            else
                die "Unknown argument: $1"
            fi
            shift
            ;;
    esac
done

require_root

# Required parameter check
[[ -z "${OCS_ZIP:-}" ]] && die "Inout : <ocs zip file> is required !"

#-------------------------------------------#
#  Download or check bootable template      #
#-------------------------------------------#
if [[ -z "${BOOTABLE_TEMPLATE_RAWDISK:-}" ]]; then
    echo  "The '--bootable_template' parameter wasn't specified, using default template : $DEFAULT_BOOTABLE_RAW_TEMPLATE ..."
    # Create temp subdir for downloaded file
    DOWNLOAD_DIR="${TMPDIR_ROOT}/download"
    mkdir ${VERBOSE_MODE} -p "${DOWNLOAD_DIR}"
    DOWNLOADED_TEMPLATE="${DOWNLOAD_DIR}/bootable_template.img"

    # Prefer curl, fallback to wget
    if command -v curl >/dev/null 2>&1; then
        echo "Use 'curl' to download  ${DEFAULT_BOOTABLE_RAW_TEMPLATE}"
        curl ${VERBOSE_MODE} -L -o "${DOWNLOADED_TEMPLATE}" "${DEFAULT_BOOTABLE_RAW_TEMPLATE}" \
            || die "curl download failed"
        curl ${VERBOSE_MODE} -L -o "${DOWNLOAD_DIR}/checksum.txt" "${DEFAULT_BOOTABLE_RAW_CHECKSUM}" \
            || die "curl download checksum failed"
    elif command -v wget >/dev/null 2>&1; then
        echo "Use 'wget' to download  ${DEFAULT_BOOTABLE_RAW_TEMPLATE}"
        wget ${VERBOSE_MODE} -O "${DOWNLOADED_TEMPLATE}" "${DEFAULT_BOOTABLE_RAW_TEMPLATE}" \
            || die "wget download failed"
        wget ${VERBOSE_MODE} -O "${DOWNLOAD_DIR}/checksum.txt" "${DEFAULT_BOOTABLE_RAW_TEMPLATE}" \
            || die "wget download checksum failed"
    else
        die "No curl or wget installed, download file failed"
    fi

    # Verify download success and file is readable
    [[ -f "${DOWNLOADED_TEMPLATE}" && -r "${DOWNLOADED_TEMPLATE}" ]] \
        || die "The download URL is not availabe or unreadable !"

    # Verify chechsum : bootable_template.img -vs-  ${DEFAULT_BOOTABLE_RAW_TEMPLATE}
    download_checksum="$(md5sum ${DOWNLOADED_TEMPLATE} | awk '{print $1}')"
    echo "'$(basename ${DOWNLOADED_TEMPLATE})' verify checksum '${download_checksum}'... "
    [[ "x${download_checksum}" != "x$(grep $(basename ${DEFAULT_BOOTABLE_RAW_TEMPLATE}) ${DOWNLOAD_DIR}/checksum.txt | awk '{print $1}')" ]] \
        &&  die "Checksum : ERROR !" \
        || echo "Checksum : PASS !"

    BOOTABLE_TEMPLATE_RAWDISK="${DOWNLOADED_TEMPLATE}"
else
    # If user provides file, skip here
    :
fi

#---------------------------#
#  Pre-check                #
#---------------------------#
for f in "${OCS_ZIP}" "${BOOTABLE_TEMPLATE_RAWDISK}"; do
    [[ -f "${f}" ]] || die "File: ${f} doesn't exist !"
    [[ -r "${f}" ]] || die "File: ${f} is not readable !"
done
for cmd in kpartx losetup mount umount unzip basename dirname; do
    command -v "${cmd}" >/dev/null 2>&1 || die "Command or package doesn't existed: ${cmd}"
done

#---------------------------#
#  1. Prepare : working directory and OCS zip file    #
#---------------------------#

# Ask overwrite answer if output file existed
OCS_BASENAME="$(basename "${OCS_ZIP}" .zip)"   # Remove suffix '.zip'
[[ -z "${OUTPUT_RAWDISK}" ]] && FINAL_IMG="${OCS_BASENAME}_u-bootable.img" || FINAL_IMG="${OUTPUT_RAWDISK}"

# If target exists, ask before overwrite
if [[ -e "${FINAL_IMG}" ]]; then
    read -p "Output file : ${FINAL_IMG} exists, overwrite (y/N)？" ANS_IF_OVERWRITE
    # Rename if $ans_if_overwrite not ^[Yy]$ ]]
    [[ "${ANS_IF_OVERWRITE}" =~ ^[Yy]$ ]] || mv ${VERBOSE_MODE} "${FINAL_IMG}" "${FINAL_IMG}_$$"
fi

mkdir ${VERBOSE_MODE} -p "${MOUNT_POINT}"
echo "Create Temp directory : ${MOUNT_POINT}"

#  1.1 Unzip OCS zip then rsync contents   #
echo "Extracting OCS zip ..."
TMP_UNZIP="${TMPDIR_ROOT}/unzipped"
mkdir ${VERBOSE_MODE} -p "${TMP_UNZIP}"
unzip $( [[ $VERBOSE_MODE ]] || echo "-q" ) "${OCS_ZIP}" -d "${TMP_UNZIP}" || die "unzip failed"

#---------------------------------------#
#  2. Deal with U-bootable raw device   #
#---------------------------------------#
TMP_RAW="${TMPDIR_ROOT}/raw_image.img"
echo "Copying bootable raw image to temporary file ..."
cp -a ${VERBOSE_MODE} "${BOOTABLE_TEMPLATE_RAWDISK}" "${TMP_RAW}"

#---------------------------------------#
#  2.1 Create loop and mapper device    #
#---------------------------------------#
echo "Setting up loop device ..."
LOOP_DEVICE=$(losetup ${VERBOSE_MODE} -f --show "${TMP_RAW}")
[[ -z "${LOOP_DEVICE}" ]] && die "Create loop device failed ..."

echo "Creating partition mappings with kpartx ..."
kpartx ${VERBOSE_MODE} -a -s "${LOOP_DEVICE}"
sleep 1   # wait for device mapper

#----------------------------------------------#
#  2.2 Find "env", "bootfs" then mount ,sync from ZIP  #
#----------------------------------------------#
# Search device mapper id which the partition "name" = "bootfs"
BOOTFS_KPARTX_MAP="/dev/mapper/$(lsblk -ln -o NAME,TYPE,PARTLABEL,PARTUUID ${LOOP_DEVICE}  \
    | awk -v this_partlabel="$BOOTFS_PARTLABEL" '$3 == this_partlabel {print $1}' \
    2>/dev/null || true)"

# Search device mapper id which the partition "name" = "env"
ENV_KPARTX_MAP="/dev/mapper/$(lsblk -ln -o NAME,TYPE,PARTLABEL,PARTUUID ${LOOP_DEVICE}  \
    | awk -v this_partlabel="$ENV_PARTLABEL" '$3 == this_partlabel {print $1}' \
    2>/dev/null || true)"

[[ -b "${BOOTFS_KPARTX_MAP}" && -b "${ENV_KPARTX_MAP}" ]] && echo "Find available partition pf 'env' and 'bootfs': ''${ENV_KPARTX_MAP}' / '${BOOTFS_KPARTX_MAP}'  .." ||  die "Unable to find partition "name" = bootfs (say : /dev/mapper/*p5)"

echo "Mounting partition ${BOOTFS_KPARTX_MAP} to ${MOUNT_POINT} ..."
mount ${VERBOSE_MODE} "${BOOTFS_KPARTX_MAP}" "${MOUNT_POINT}"

#  Save u-Boot env of raw imag as file  #
dd if=${ENV_KPARTX_MAP} bs=1 skip=4 | tr '\0' '\n' | sed -e '/^$/d' -e  '/^\o377*$/d' > "${TMP_UNZIP}"/live/ocs-u-boot-env-exported.txt
OCS_U_BOOT_ENV_VERSION="$(grep -E '^ocs_u_boot_env_version=' "${TMP_UNZIP}/live/ocs-u-boot-env-exported.txt" | cut -d'=' -f2-)"

#----------------------------------------------#
# 2.3 Check soft links and save repack info
#----------------------------------------------#
for _item in $SOFT_ITEAMS ; do 
    echo "Check soft link :'${_item}' in '${MOUNT_POINT}/' ..."
    [[ -L "${MOUNT_POINT}/${_item}" && \
        "$(readlink -f "${MOUNT_POINT}/${_item}")" == "$(readlink -f "${MOUNT_POINT}/live/${_item}")" ]] \
        || ( rm -rf ${VERBOSE_MODE} "${MOUNT_POINT}/${_item}" ; ln -sf ${VERBOSE_MODE} "${MOUNT_POINT}/live/${_item}" "${MOUNT_POINT}/${_item}")
done

#  Write repack info to file  #
INFO_FILE="${TMP_UNZIP}/live/repack-raw_info.txt"
{
    echo "=== Repack Info  ==="
    echo " * Date                  : $(date '+%Y-%m-%d %H:%M:%S')"
    echo " * OCS-live-ubrd Version : ${THIS_VERSION}"
    echo " * U-Boot ENV Version    : ${OCS_U_BOOT_ENV_VERSION}"
    echo " * Original OCS Zip      : $(md5sum ${OCS_ZIP})"
    echo " * Template Image        : $(md5sum ${BOOTABLE_TEMPLATE_RAWDISK})"
    echo " * Loop Device           : ${LOOP_DEVICE}"
    echo " * Bootfs / ENV Mapper   : ${BOOTFS_KPARTX_MAP} / ${ENV_KPARTX_MAP} "
    echo " * Full Self Command     : ${FULL_SELF_CLI}"
    echo "=== End of info ==="
} > "${INFO_FILE}"
echo "Write Repack Info into: ${INFO_FILE}"
cat "${TMP_UNZIP}/live/repack-raw_info.txt"

# To sync tmp-zip dir to "${MOUNT_POINT}
echo "Syncing extracted files to mounted partition ..."
rsync -a ${VERBOSE_MODE} "${TMP_UNZIP}/." "${MOUNT_POINT}/" || die "rsync failed"

#---------------------------------------#
#  3. Umount , remove loop device       #
#---------------------------------------#
echo "Syncing filesystem ..."
sync

echo "Unmounting partition ..."
umount ${VERBOSE_MODE} "${MOUNT_POINT}" || die "umount failed"

# Delete kpartx mapper（same with function : cleanup ）
kpartx ${VERBOSE_MODE} -d "${LOOP_DEVICE}" || echo "Warning: kpartx -d failed"

# Release loop device (same with function : cleanup ）
losetup ${VERBOSE_MODE} -d "${LOOP_DEVICE}" || echo "Warning: losetup -d failed"

#---------------------------------------#
#  4. Rename output ,show dd command    #
#---------------------------------------#
echo "Renaming temporary raw image to ${FINAL_IMG} ..."
mv ${VERBOSE_MODE} "${TMP_RAW}" "${FINAL_IMG}"

echo "=== Done ==="
echo "Output : ${FINAL_IMG}"
echo "Then, you can use the following command to create a bootable OCS disk with U-boot/SPI enabled :"
echo " $ sudo dd if=${FINAL_IMG} of=/dev/mmcblk0 bs=1024M conv=fsync"

# If debug mode, show temp dir location before exit (same with cleanup)
if $DEBUG_MODE; then
    echo "Debug mode: Keep TMPDIR_ROOT = ${TMPDIR_ROOT}"
else 
    if [[ -d "${TMPDIR_ROOT}" ]]; then
        echo "Removing temporary directory ${TMPDIR_ROOT} ..."
        rm ${VERBOSE_MODE} -rf "${TMPDIR_ROOT}"
    fi
fi

exit 0