#!/usr/bin/env bash
#=====================================================================
#  File:   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_repack_$$"
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_PART_NAME="bootfs"

# Runtime Variables / Control flag
OCS_ZIP=""
LOOP_DEVICE=""
KPARTX_MAP=""
DOWNLOADED_TEMPLATE=""
FULL_SELF_CLI="$BASH_SOURCE$(printf " %q" "$@")"

DEBUG_MODE=false
VERBOSE_MODE=

#---------------------------#
#  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. 
   -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 
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 "${KPARTX_MAP}" ]]; then
        echo "Removing kpartx map ${KPARTX_MAP} ..."
        kpartx ${VERBOSE_MODE} -d "${LOOP_DEVICE}" || echo "Warning: kpartx -d failed"
        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
            ;;
        -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. Create temp directory    #
#---------------------------#
mkdir ${VERBOSE_MODE} -p "${MOUNT_POINT}"
echo "Create Temp directory : ${MOUNT_POINT}"

#---------------------------------------#
#  2. Deat with "bootfs" in 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}"

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

# Search device mapper id which the partition "name" = "bootfs"
KPARTX_MAP="/dev/mapper/$(lsblk -ln -o NAME,LABEL ${LOOP_DEVICE}  \
    | awk -v bootfs_part_name="$BOOTFS_PART_NAME" '$2 == bootfs_part_name {print $1}' \
    2>/dev/null || true)"

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

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

#---------------------------#
#  3. 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"

echo "Rsyncing extracted files to mounted partition ..."
rsync -a ${VERBOSE_MODE} "${TMP_UNZIP}/." "${MOUNT_POINT}/" || die "rsync failed"
#cp ${VERBOSE_MODE} -a "${TMP_UNZIP}/." "${MOUNT_POINT}/" || die "copy failed"

#--------------------------------#
#  4. Write repack info to file  #
#--------------------------------#
INFO_FILE="${MOUNT_POINT}/repack-raw_info.txt"
{
    echo "=== Repack Info  ==="
    echo "Date                  : $(date '+%Y-%m-%d %H:%M:%S')"
    echo "Original OCS Zip      : ${OCS_ZIP}"
    echo "Bootable Raw Template : ${BOOTABLE_TEMPLATE_RAWDISK}"
    echo "Loop Device           : ${LOOP_DEVICE}"
    echo "Bootfs Mapper         : ${KPARTX_MAP}"
    echo "Full Self Command     : ${FULL_SELF_CLI}"
    echo "=== End of info ==="
} > "${INFO_FILE}"

echo "Write Repack Info into: ${INFO_FILE}"

#---------------------------------------#
#  5. Umount then rename raw 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"

# Generate its output name
OCS_BASENAME="$(basename "${OCS_ZIP}" .zip)"   # Remove suffix '.zip'
FINAL_IMG="${OCS_BASENAME}_u-bootable.img"

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

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/mmcdevice 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
