#!/bin/sh
#
# Init Script for Klipper Mod - running in Flashforge Host System
#
set -ex # abort on errors, print scripts steps

# the mod is installed here
mod_dir="/data/.klipper_mod"
chroot_dir="$mod_dir/chroot"

# log files for init process
log_file="/root/mod_init.log"
log_file_debug="/root/mod_init_debug.log"

# marker files allow to intervene in init process
# marker files in mod root home
mod_marker_skip="$chroot_dir/root/.reboot_stock_once"
mod_marker_remove="$chroot_dir/root/.reboot_and_remove"
# marker files on usb
usb_marker_skip="klipper_mod_skip"
usb_marker_remove="klipper_mod_remove"
usb_marker_debug="klipper_mod_debug"
# debug script is started from usb if debug marker is found
usb_debug_script="mod_debug.sh"

continue_flashforge_boot() {
    # continue with booth of the flashforge system
    # add boot delay, because otherwise in case of warm reboot from klipper mod
    # the flashforge sofware might fail to initialize (probably related to MCU state)
    sleep 5
    exit 0
}

uninstall() {
    rm -rf $mod_dir || true
    rm -f /root/mod_init* || true
    rm -f "$0" || true # remove this script
}

mount_data() {
    # mount data - this would otherwise be mounted later by flashforge
    if [ ! "`mount | grep "mmcblk0p7"`" ]; then
        echo "Mount /data partition"
        fsck -y /dev/mmcblk0p7 || true
        mount /dev/mmcblk0p7 /data;
    fi
}

# check for marker files on usb
check_usb_marker() {
    # wait for usb devices to appear
    sleep 2

    # scan for mountable drives with marker files
    for sd in sda sdb; do
        for i in 1 2 3 4 ""; do
            if [ -e /dev/$sd$i ]; then
                mount -t vfat -o ro,codepage=437,iocharset=utf8 /dev/$sd$i /mnt || continue
                # check for mod marker usb
                usb_mod_marker=$(ls -1 /mnt/ | grep "^klipper_mod" | head -n1)
                # check for flashforge update on usb
                usb_update_marker=$(ls -1 /mnt/ | grep "^Adventurer.*tgz" | head -n1)
                [ -n "$usb_update_marker" ] || usb_update_marker=$(ls -1 /mnt/ | grep "^flashforge_init.sh" | head -n1)
                umount /dev/$sd$i || true;

                if ([ -n "$usb_mod_marker" ] || [ -n "$usb_update_marker" ]); then
                    usb_drive="$sd$i"
                    echo "Found marker file."
                    break 2;
                fi
            fi
        done
    done

    # check for debug marker. copy init logs and start debug script if detected
    if [ "$usb_mod_marker" = "$usb_marker_debug" ]; then
        echo "Found debug marker $usb_mod_marker on $usb_drive.";
        mount -t vfat -o codepage=437,iocharset=utf8 /dev/$usb_drive /mnt || true
        sync
        echo "Copy init log files."
        cp -f /root/mod_init* /mnt/ || true
        sync
        if [ -f /mnt/$usb_debug_script ]; then
            echo "Execute debug script."
            /bin/sh /mnt/$usb_debug_script || true
        fi
        umount /dev/$usb_drive || true;
    fi

    # bail, if we found a flashforge update
    if [ -n "$usb_update_marker" ]; then
        echo "Skip Klipper Mod init due to flashforge update '$usb_update_marker' found on $usb_drive."
        continue_flashforge_boot
    fi

    # allow to boot into stock linux, if skip marker is present
    if [ "$usb_mod_marker" = "$usb_marker_skip" ]; then
        echo "Skip Klipper Mod init due to skip marker '$usb_mod_marker' found on $usb_drive."
        continue_flashforge_boot
    # remove mod, if remove marker is present
    elif [ "$usb_mod_marker" = "$usb_marker_remove" ]; then
        echo "Uninstall Klipper Mod due to remove maker '$usb_mod_marker' found on $usb_drive."
        uninstall
        continue_flashforge_boot
    fi
}

# check for marker files in mod dir
check_mod_marker() {

    # check if mod is installed on /data, else skip init.
    if [ ! -d $chroot_dir ]; then
        echo "Klipper Mod not found on /data partition. Skip init.";
        exit 0;
    fi

    # allow to boot into stock linux, if skip marker is present
    if [ -f $mod_marker_skip ]; then
        echo "Skip Klipper Mod init once due to skip maker $mod_marker_skip."
        rm -f $mod_marker_skip
        sync
        continue_flashforge_boot
    fi

    # remove mod, if remove marker is present
    if [ -f $mod_marker_remove ]; then
        echo "Uninstall Klipper Mod due to remove maker $mod_marker_remove."
        uninstall
        continue_flashforge_boot
    fi
}

prepare_chroot() {
    # we need to mount /dev/pts, here because it doenst exists at this init stage for some reason
    # check if it is there, because this function is also used during install
    if [ ! -d /dev/pts ]; then
        mkdir /dev/pts
        mount -o gid=5,mode=620,ptmxmode=0666 -t devpts devpts /dev/pts
    fi

    # mount bind virtual filesystems + tmp
    for dir in /dev /dev/pts /proc /sys /run /tmp; do mount -o bind $dir $chroot_dir$dir; done
    # mount bind data
    mount -o bind /data $chroot_dir/mnt/data
    # mount bind original root, remount ro
    mount -o bind / $chroot_dir/mnt/orig_root
    mount -o bind,remount,ro $chroot_dir/mnt/orig_root
}

teardown_chroot() {
    umount -f $chroot_dir/mnt/data || true
    umount -f $chroot_dir/mnt/orig_root || true
    for dir in /dev /dev/pts /proc /sys /run /tmp; do umount -f $chroot_dir$dir || true; done
}

start_klipper_mod() {
    echo "Start Klipper Mod Init" | tee /dev/kmsg
    check_usb_marker
    mount_data
    check_mod_marker
    prepare_chroot
    chroot $chroot_dir /etc/init/init_chroot.sh
    # now we have hijacked the init process and let it go :-)
    kill -9 $PPID
    sync # make sure our logs are saved
} > $log_file 2> $log_file_debug

start() {
    # rotate init logs
    mv -f $log_file $log_file.old || true
    mv -f $log_file_debug $log_file_debug.old || true
    start_klipper_mod
}

stop() {
    printf "Stop Klipper Mod"
    chroot $chroot_dir /etc/init/stop_chroot.sh
    teardown_chroot
}

setup() {
    # called once from install script
    prepare_chroot
    # flashforge installer mounts flash drive directly to /mnt
    # mount bind it inside the chroot for user provided overlay
    mount -o bind /mnt $chroot_dir/media
    chroot $chroot_dir /etc/init/setup_chroot.sh
}

case "$1" in
    start)
        start
        ;;
    stop)
        stop
        ;;
    # this is used by installer
    setup)
        setup
        ;;
    # expose some functions for testing
    prepare_chroot)
        prepare_chroot
        ;;
    teardown_chroot)
        teardown_chroot
        ;;
    mount_data)
        mount_data
        ;;
    check_usb_marker)
        check_usb_marker
        ;;
    check_mod_marker)
        check_mod_marker
        ;;
    uninstall)
        uninstall
        ;;
    *)
        echo "Usage: $0 {start|stop}"
        exit 1
esac

exit $?
