#!/bin/bash

srcdev="$1"
destdev="$2"
mediatype="$3"
bootloader="$4"

FORCE="${FORCE:-no}"
ORIG_DESTDEV="${destdev}"

declare -A partition_sizes
declare -A partition_filesystems
declare -A partition_mounts
declare -A partition_fstabparams
partition_filesystems["rootfs"]=ext4
partition_filesystems["log"]=ext4
partition_filesystems["settings"]=ext4
partition_mounts["rootfs"]="/"
partition_mounts["log"]="/var/log"
partition_mounts["settings"]="/opt/stpc/var/lib/etc"
partition_fstabparams["rootfs"]="ro,noatime,errors=remount-ro 0 1"
partition_fstabparams["log"]="rw,noatime,nofail,errors=remount-ro 0 0"
partition_fstabparams["settings"]="rw,noatime,nofail,errors=remount-ro 0 0"

boot_partition="rootfs"
partition_order=("rootfs" "log" "settings")

tmproot="/tmp/rootfs"
imgroot="/tmp/imgroot"
imgloop="loop0"
imgfilename="stpc-image.img"
onpanic=("heartbeat" "heartbeat" "heartbeat" "heartbeat")
onsuccess=("heartbeat" "mmc0" "none" "mmc1")


function setled () {
    local base="/sys/class/leds/beaglebone:green:usr"
    local ledindex="$1"
    local ledvalue="$2"

    if [ "${ledvalue}" = "1" ]; then
        echo 255 > "${base}${ledindex}/brightness"
    elif [ "${ledvalue}" = "0" ]; then
        echo 0 > "${base}${ledindex}/brightness"
    fi
}

function settrigger () {
    local base="/sys/class/leds/beaglebone:green:usr"
    local ledindex="$1"
    local ledtrigger="$2"

    echo "${ledtrigger}" > "${base}${ledindex}/trigger"
    sleep 1
}

function setupgpio () {
    local base="/sys/class/gpio"
    local gpioindex="$1"
    local gpiodirection="$2"

    if ! [ -f "${base}/gpio${gpioindex}" ]; then
        echo ${gpioindex} > ${base}/export
        sleep 1
        echo ${gpiodirection} > ${base}/gpio${gpioindex}/direction
        sleep 1
    fi
}

function setgpio () {
    local base="/sys/class/gpio"
    local gpioindex="$1"
    local gpiovalue="$2"

    echo ${gpiovalue} > ${base}/gpio${gpioindex}/value
}

function die () {
    echo -e "\033[31merror: $1\033[0m"

    local modeindex=0
    if [ -e /sys/class/leds/beaglebone\:green\:usr0/trigger ] ; then
        for mode in "${onpanic[@]}"; do
            settrigger ${modeindex} "${mode}"
            let modeindex=modeindex+1
        done
    fi
    exit 1
}

function info () {
    echo -e "\033[33m$1\033[0m"
}

function usage () {
    echo "flasher srcdev destdev mediatype"
    echo "    srcdev     - source device such as /dev/mmcblk1"
    echo "    destdev    - destination device such as /dev/mmcblk0"
    echo "               -   if using imagefile, this must be a mountable"
    echo "               -   partition to stick the image file on"
    echo "               -   such as /dev/mmcblk0p1"
    echo "    mediatype  - media type of destdev (disk or imagefile)"
    echo "    bootloader - which bootloader to use (uSD or eMMC)"
    echo
    echo "other options that may be passed in are"
    echo "FORCE=yes               - force answer yes"
    echo "MAKE_FLASHER=yes        - make the image a flasher image"
    echo "__IS_INIT_FLASHER__=yes - is an init flasher at runtime (examine the code to see what this does)"
}

function setupsizes () {
    if [ "${mediatype}" = "imagefile" ]; then
        partition_sizes["rootfs"]=2G
        partition_sizes["log"]=500M
        partition_sizes["settings"]=64M
        partition_method="imagefile"
        imagefile_size=3G
    elif [ "${mediatype}" = "disk" ]; then
        partition_sizes["rootfs"]=2G
        partition_sizes["log"]=1G
        partition_sizes["settings"]=REST
        partition_method="disk"
    else
        usage
        die "unknown media type"
    fi
}

function setupsystem () {
    # setup a temporary file system that should be writable just in case
    mount -t tmpfs tmpfs /tmp || true

    if [ ! -e "/sys/class/leds/beaglebone:green:usr0/trigger" ] ; then
        # attempt to probe the leds
        modprobe leds_gpio || true
        sleep 1
    fi
}

function eyecandy () {
    local states=(\
        "0001" "0011" "0010" "0110" "0100" "1100" "1000"\
        "1100" "0100" "0110" "0010" "0011"\
    )

    setupgpio 51 out
    setupgpio 50 out

    # check if we have leds to cycle through
    if [ -e /sys/class/leds/beaglebone\:green\:usr0/trigger ] ; then
        settrigger 0 none
        settrigger 1 none
        settrigger 2 none
        settrigger 3 none

        while true; do
            for state in "${states[@]}"; do
                led0="${state:3:1}"
                led1="${state:2:1}"
                led2="${state:1:1}"
                led3="${state:0:1}"

                setled 0 "${led0}"
                setled 1 "${led1}"
                setled 2 "${led2}"
                setled 3 "${led3}"
                setgpio 51 "${led0}"
                setgpio 50 "${led1}"
                sleep 0.1
            done
         done
    fi
}

function flushcache () {
    sync
    blockdev --flushbufs "${ORIG_DESTDEV}"
}

function getbytes_from_gigabytes () {
    value="$1"
    value=`echo $(getbytes_from_megabytes ${value})*1024 | bc`
    echo "${value}"
}

function getbytes_from_megabytes () {
    value="$1"
    value=`echo $(getbytes_from_kilobytes ${value})*1024 | bc`
    echo "${value}"
}

function getbytes_from_kilobytes () {
    value="$1"
    value=`echo ${value}*1024 | bc`
    echo "${value}"
}

function getbytes () {
    value="$1"

    if [[ "${value}" =~ ^[0-9]+(\.[0-9]+)?[GMk]$ ]]; then
        let i=${#value}-1
        type="${value:${i}:1}"
        value="${value:0:${i}}"
    elif [[ "${value}" =~ ^[0-9]+$ ]]; then
        type=""
        value="${value}"
    else
        echo "UNKNOWN TYPE"
        return
    fi

    if [ "${type}" = "G" ]; then
        getbytes_from_gigabytes "${value}"
    elif [ "${type}" = "M" ]; then
        getbytes_from_megabytes "${value}"
    elif [ "${type}" = "k" ]; then
        getbytes_from_kilobytes "${value}"
    elif [ "${type}" = "" ]; then
        echo "${value}"
    else
        echo "UNKNOWN TYPE"
        return
    fi
}

function bytes_to_kilobytes () {
    value="$1"
    value=`echo ${value}/1024 | bc`
    echo "${value}"
}

function bytes_to_megabytes () {
    value="$1"
    value=`echo $(bytes_to_kilobytes ${value})/1024 | bc`
    echo "${value}"
}

function bytes_to_gigabytes () {
    value="$1"
    value=`echo $(bytes_to_megabytes ${value})/1024 | bc`
    echo "${value}"
}

function floor () {
    value="$1"
    value=`echo "${value}/1" | bc`
    echo "${value}"
}

function checkparams () {
    if [ "${srcdev}" = "" ]; then
        usage
        die "must define srcdev"
    fi

    if [ "${destdev}" = "" ]; then
        usage
        die "must define destdev"
    fi

    [ "`getrootdev`" = "" ] && die "cannot find rootdev in configuration"
}

function askokay () {
    if [ "${FORCE}" = "yes" ]; then
        return
    fi

    read -p "Do you really want to continue with ${destdev}? (y/n) " result
    if ! [ "${result}" = "y" ]; then
        die "stopping on user command"
    fi
}

function find_dev_by_mount () {
    local mnt="$1"
    local partindex=1

    for partition in "${partition_order[@]}"; do
        if [ "${partition_mounts[${partition}]}" = "${mnt}" ]; then
            echo "${destdev}p${partindex}"
            return
        fi
        let partindex=partindex+1
    done
    echo ""
}

function getrootdev () {
    find_dev_by_mount "/"
}

function sfdisk_format () {
    local boot
    local size
    start=8192

    for partition in "${partition_order[@]}"; do
        if [ "${partition}" = "${boot_partition}" ]; then
            boot="*"
        else
            boot=""
        fi

        size="${partition_sizes[${partition}]}"
        if [ "${size}" = "REST" ]; then
            size=""
        fi

        # this works because it tells sfdisk to start searching
        # at ${start} sectors so as partitions are added, sfdisk
        # will find the end of the last partition and stick the next
        # partition there.
        echo "${start},${size},,${boot}"
    done
}

function cleanup () {
    info "unmounting '${tmproot}' and its subcomponents"
    umount -R "${tmproot}" || true
    umount -R "${imgroot}" || true
    kpartx -d "/dev/${imgloop}" || true
    sleep 1
}

function flash_bootloader () {
    local destination="$1"

    info "flashing '${bootloader}/u-boot.img' to '${destination}'"
    dd if=/opt/stpc/lib/firmware/${bootloader}/u-boot.img of=${destination} count=2 seek=1 conv=notrunc bs=384k \
        || die "failed to flash u-boot.img"

    info "flashing '${bootloader}/MLO' to '${destination}'"
    dd if=/opt/stpc/lib/firmware/${bootloader}/MLO of=${destination} count=1 seek=1 conv=notrunc bs=128k \
        || die "failed to flash MLO"
}

function flash_bootloader_disk () {
    flash_bootloader "${destdev}"
}

function flash_bootloader_imagefile () {
    flash_bootloader "${imgroot}/${imgfilename}"
}

function partition_disk () {
    info "cleaning '${destdev}'"
    dd if=/dev/zero of="${destdev}" bs=1M count=108 || die "failed to clean '${destdev}'"
    sync
    dd if="${destdev}" of=/dev/zero bs=1M count=108 || die "failed to sync cleaned '${destdev}'"
    sync
    flushcache
    partprobe
    sleep 1

    info "partitioning '${destdev}'"
    sfdisk "${destdev}" <<< "`sfdisk_format`" \
        || die "failed to partition '${destdev}'"
    flushcache
    partprobe
    sleep 1
}

function partition_imagefile () {
    info "partitioning '${imgroot}/${imgfilename}'"
    sfdisk "${imgroot}/${imgfilename}" <<< "`sfdisk_format`" \
        || die "failed to partition '${imgroot}/${imgfilename}'"
    flushcache

    info "mounting fake disk"
    kpartx -a "${imgroot}/${imgfilename}" || die "failed to mount imagefile"
    destdev="/dev/mapper/${imgloop}"
    flushcache
    sleep 1
}

function create_filesystems () {
    local partition_index=1
    local fs
    for partition in "${partition_order[@]}"; do
        fs="${partition_filesystems[${partition}]}"
        info "creating filesystem '${fs}' for partition '${partition}'"
        mkfs.${fs} -F ${destdev}p${partition_index} \
            || die "failed to make '${partition}' filesystem '${fs}'"
        let partition_index=partition_index+1
    done
    flushcache
}

function mount_partitions () {
    local partition_index=1
    local mnt
    for partition in "${partition_order[@]}"; do
        mnt="${partition_mounts[${partition}]}"
        info "mounting partition '${partition}' as '${mnt}'"
        mkdir -p "${tmproot}/${mnt}" || die "failed to create mount point '${tmproot}/${mnt}'"
        mount "${destdev}p${partition_index}" "${tmproot}/${mnt}" \
            || die "failed to mount '${partition}' to '${tmproot}/${mnt}'"
        let partition_index=partition_index+1
    done
}

function sync_rootfs () {
    info "syncing rootfs to '${tmproot}'"
    rsync -aAx /* "${tmproot}" \
        --exclude={/dev/*,/proc/*,/sys/*,/tmp/*,/run/*,/mnt/*,/media/*,/lost+found,/lib/modules/*,/uEnv.txt}

    info "creating ssh key regeneration file"
    touch "${tmproot}/etc/ssh/ssh.regenerate"

    info "syncing kernel modules to '${tmproot}'"
    mkdir -p "${tmproot}/lib/modules/`activekernel`"
    rsync -aAx /lib/modules/`activekernel`/* "${tmproot}/lib/modules/`activekernel`/"
    flushcache
}

function generate_fstab () {
    local partindex=1
    local uuid
    local mnt
    local fs
    local fstabparams

    echo "# /etc/fstab: static file system information."
    echo "#"

    for partition in "${partition_order[@]}"; do
        uuid=`/sbin/blkid -c /dev/null -s UUID -o value ${destdev}p${partindex}`
        mnt="${partition_mounts[${partition}]}"
        fs="${partition_filesystems[${partition}]}"
        fstabparams="${partition_fstabparams[${partition}]}"
        echo "UUID=${uuid} ${mnt} ${fs} ${fstabparams}"
        let partindex=partindex+1
    done

    echo "tmpfs /tmp tmpfs rw,nodev,nosuid,size=64M 0 0"
    echo "debugfs /sys/kernel/debug debugfs defaults 0 0"
}

function patch_new_system () {
    # setup uuid
    local rootdev=`getrootdev`
    local root_uuid=`/sbin/blkid -c /dev/null -s UUID -o value ${rootdev}`
    info "setup uEnv.txt uuid=${root_uuid}"
    sed -i -e 's:uuid=:#uuid=:g' "${tmproot}/boot/uEnv.txt"
    echo "uuid=${root_uuid}" >> "${tmproot}/boot/uEnv.txt"

    # patch other stuff in uEnv.txt
    info "reconfigure uEnv.txt"
    sed -i -e 's,^enable_uboot_cape_universal=,#enable_uboot_cape_universal=,g' "${tmproot}/boot/uEnv.txt"
    sed -i -e 's,^uboot_overlay_pru=,#uboot_overlay_pru=,g' "${tmproot}/boot/uEnv.txt"
    sed -i -e 's,^cmdline=,#cmdline=,g' "${tmproot}/boot/uEnv.txt"

    if [ "${MAKE_FLASHER}" = "yes" ]; then
        echo "cmdline=init=/opt/stpc/bin/init-flasher" >> "${tmproot}/boot/uEnv.txt"
    else
        echo "cmdline=coherent_pool=1M net.ifnames=0 quiet" >> "${tmproot}/boot/uEnv.txt"
    fi

    if [ "${MAKE_FLASHER}" = "yes" ]; then
        # disable stpcd service so that it does not start and
        # enable the poweroff service
        chroot "${tmproot}" /bin/bash <<< "
        systemctl disable stpcd.service
        systemctl enable poweroff-on-boot.service
        "
    else
        # enable stpcd service so that it does not start and
        # disable the poweroff service
        chroot "${tmproot}" /bin/bash <<< "
        systemctl enable stpcd.service
        systemctl disable poweroff-on-boot.service
        "
    fi

    # generate fstab
    info "generate fstab"
    fstab=`generate_fstab`
    echo "${fstab}" > "${tmproot}/etc/fstab"
    cat "${tmproot}/etc/fstab"

    # ensure /opt/stpc/var/lib/etc has dPG.dosis owner.group
    chmod -Rv dPG.dosis "${tmproot}/opt/stpc/var/lib/etc"

    # syncronize everything
    flushcache
    dd if=/dev/mmcblk1 of=/dev/null count=100000
}

function check_filesystem () {
    local partindex=1
    for partition in "${partition_order[@]}"; do
        info "filesystem check ${destdev}p${partindex}"
        fsck -a "${srcdev}p${partindex}"
        let partindex=partindex+1
    done
}

function mount_imagefile_destdev () {
    info "mounting '${destdev}' to '${imgroot}'"
    mkdir -p "${imgroot}" || die "failed to create '${imgroot}'"
    mount "${destdev}" "${imgroot}" || die "failed to mount '${imgroot}'"
}

function setup_imagefile () {
    imgsz=$(floor `getbytes ${imagefile_size}`)
    if [ "${imgsz}" = "" ]; then
        die "bad image size ${imagefile_size}, programmer error"
    fi

    imgsz=`bytes_to_megabytes "${imgsz}"`

    info "creating empty image '${imgroot}/${imgfilename}'"
    dd if=/dev/zero of="${imgroot}/${imgfilename}" bs=1M count=${imgsz} status=progress \
        || die "failed to create '${imgroot}/${imgfilename}'"
}

function configure_partition_method () {
    if [ "${partition_method}" = "imagefile" ]; then
        info "partition method imagefile"
        mount_imagefile_destdev
        setup_imagefile
        partition_imagefile
        flash_bootloader_imagefile
    elif [ "${partition_method}" = "disk" ]; then
        info "partition method disk"
        partition_disk
        flash_bootloader_disk
    else
        die "unknown partition method"
    fi
}

function cleanupeyecandy () {
    if ! [ "${eyecandypid}" = "" ]; then
        kill ${eyecandypid} || true
    fi
}

function finished () {
    local modeindex=0
    if [ -e /sys/class/leds/beaglebone\:green\:usr0/trigger ] ; then
        for mode in "${onsuccess[@]}"; do
            settrigger ${modeindex} "${mode}"
            let modeindex=modeindex+1
        done
    fi

    cleanupeyecandy || true
    if [ "${__IS_INIT_FLASHER__}" = "yes" ]; then
        exec /sbin/init
    fi
}

checkparams
setupsizes
askokay
setupsystem
cleanup

eyecandy &
eyecandypid=$!

check_filesystem
configure_partition_method
create_filesystems
mount_partitions
sync_rootfs
patch_new_system
finished
