#!/bin/sh
# Called from /etc/fstab to ensure we have something writable mounted
# at /cfg, /etc, /home, /root, and /var.
#
# /mnt/cfg is used to store device modifications made to /etc, /home,
# /root.
#
# /mnt/var is used as the backing store for /var.  If /mnt/var isn't
# available, /mnt/cfg will still be be used to persist /var/lib.
#
# If a disk partiion with the corresponding label ("cfg"/"var") is
# available, that is used to persist changes to the aforementioned
# directories. Otherwise fall back to a tmpfs based RAM disk. This
# effectively brings up the system with the default configuration, but
# obviously any subsequent configuration is ephemeral.
. /etc/os-release

set -e

nm=$(basename "$0")
err=0
mmc=""
opt="-k"

# External button or bootloader changed kernel command line
check_factory()
{
    if [ -f /mnt/cfg/infix/.reset ]; then
	return 0;
    fi

    if grep -q 'finit.cond=factory-reset' /proc/cmdline; then
	return 0;
    fi

    # Add to your br2-external to extend factory-reset check
    if [ ! -x /usr/libexec/infix/check-factory ]; then
	return 1;
    fi

    /usr/libexec/infix/check-factory
}

factory_reset()
{
    find /sys/class/leds/ -type l -exec sh -c 'echo 100 > $0/brightness' {} \;
    print_start "Resetting to factory defaults."
    logger $opt -p user.notice -t "$nm" "Resetting to factory defaults."

    rm -rf /mnt/cfg/* /mnt/var/*
    sync

    logger $opt -p user.notice -t "$nm" "Factory reset complete."
    print_end 0 "Factory reset complete."
    print_restore
}

is_mmc()
{
    [ -n "$mmc" ] && return $mmc

    # Check if primary or secondary partition (our rootfs) is on MMC
    for label in primary secondary; do
        devname=$(find_partition_by_label "$label" 2>/dev/null)
        if [ -n "$devname" ]; then
            case "$devname" in
                mmcblk*)
		    mmc=0
                    return 0
                    ;;
            esac
        fi
    done

    mmc=1
    return 1
}

wait_mmc()
{
    # Try up to 50 times with 0.2s sleep = 10 second timeout
    for _ in $(seq 50); do
        if ls /dev/mmcblk* >/dev/null 2>&1; then
            logger $opt -p user.notice -t "$nm" "MMC device available after delay"
            return 0
        fi
        sleep .2
    done

    logger $opt -p user.warn -t "$nm" "Timeout waiting for MMC device"
    return 1
}

# This early on we don't have the luxury of /dev/disk/by-label/$1
find_partition_by_label()
{
    label="$1"

    for diskpath in /sys/class/block/*; do
	devname=$(basename "$diskpath")

	# Skip partitions, only check whole disks
	[ -f "$diskpath/partition" ] && continue

	# Skip ram, loop, and other virtual devices
	case "$devname" in
	    ram*|loop*|nullb*|dm-*) continue ;;
	esac

	disk="/dev/$devname"

	#
	# 1. Try GPT/MBR partition label using sgdisk
	#
	result=$(sgdisk -p "$disk" 2>/dev/null | awk -v label="$label" -v devname="$devname" '
	    /^ *[0-9]/ {
		if ($7 == label) {
		    if (devname ~ /^(mmcblk|nvme|loop)/)
			print devname "p" $1;
		    else
			print devname $1;
		    exit 0;
		}
	    }
	')

	if [ -n "$result" ]; then
	    echo "$result"
	    return 0
	fi

	#
	# 2. Fallback: Check if the whole disk is an ext4/ext2/ext3 filesystem
	#

	# Check for ext4/ext2/ext3 magic number (0xEF53) at offset 1080 (1024+56).
	magic_number=$(dd if="$disk" bs=1 skip=1080 count=2 2>/dev/null | od -t x2 -A n | tr -d ' \n')

	# Check for both Little-Endian ('53ef') and Big-Endian ('ef53') interpretations of 0xEF53.
	# This supports bi-endian architectures like MIPS that may run in BE mode,
	# as well as the LE mode which is standard for RISC-V and ext filesystems.
	if [ "$magic_number" = "ef53" ] || [ "$magic_number" = "53ef" ]; then

	    # Read the volume label from offset 1144 (1024+120)
	    fslabel=$(dd if="$disk" bs=1 skip=1144 count=16 2>/dev/null | tr -d '\000')
	    logger $opt -p user.notice -t "$nm" "Found label $fslabel on disk $disk ..."

	    if [ "$fslabel" = "$label" ]; then
		echo "$devname"
		return 0
	    fi
	fi
    done

    return 1
}

status()
{
    GREEN="$(printf '\033[1;32m')"
    RED="$(printf '\033[1;31m')"
    YELLOW="$(printf '\033[1;33m')"
    BOLD="$(printf '\033[1m')"
    RESET="$(printf '\033[0m')"

    case $1 in
        0) color=$GREEN; text=' OK ' ;;
        1) color=$RED;   text='FAIL' ;;
        2) color=$YELLOW;text='WARN' ;;
        *) color=$YELLOW;text=' ⋯  ' ;;
    esac

    printf '%s[%s%s%s%s]%s ' "$BOLD" "$color" "$text" "$RESET" "$BOLD" "$RESET"
}

print_start()
{
    printf '\r\033[K%s%s' "$(status 3)" "$*" > /dev/console
}

print_end()
{
    rc=$1; shift
    if [ $# -gt 0 ]; then
        printf '\r\033[K%s%s\n' "$(status "$rc")" "$*" > /dev/console
    else
        printf '\r%s\n' "$(status "$rc")" > /dev/console
    fi
}

# Restore Finit's original progress message so its [ OK ] appears correctly
print_restore()
{
    print_start "Mounting filesystems from /etc/fstab"
}

# Helper to log resize error and create failure marker (stage 1)
resize_err()
{
    print_end 1
    logger $opt -p user.err -t "$nm" "$1"
    echo "1" > /mnt/aux/resized.failed

    return 1
}

# Helper to log filesystem resize error and rename pending marker (stage 2)
resize_fs_err()
{
    print_end 1
    logger $opt -p user.err -t "$nm" "$1"
    rm -f /mnt/aux/resized.pending
    echo "1" > /mnt/aux/resized.failed

    return 1
}

# Stage 1: Expand the given partition to fill up the rest of storage
# On success, creates "resized.pending" marker and reboots.
# On failure, creates "resized.failed" marker and returns 1.
resize_partition()
{
    label="$1"

    print_start "Resizing /var partition on sdcard, please wait ..."

    devname=$(find_partition_by_label "$label")
    [ -z "$devname" ] && resize_err "Label \"$label\" not found"

    part="/dev/$devname"
    diskname=$(basename "$(readlink -f "/sys/class/block/$devname/..")")
    disk="/dev/$diskname"
    partnum="${devname##*[^0-9]}"

    logger $opt -p user.notice -t "$nm" "Found partition $part (partition $partnum on $disk)"

    start=$(sgdisk -i "$partnum" "$disk" 2>/dev/null | grep "First sector:" | awk '{print $3}')
    [ -z "$start" ] && resize_err "Could not determine start sector for partition $partnum"

    logger $opt -p user.notice -t "$nm" "Expanding partition $partnum from sector $start to end of disk"

    sgdisk -e "$disk" 2>&1 | logger $opt -p user.notice -t "$nm" || \
	resize_err "Failed expanding GPT on $disk"

    sgdisk -d "$partnum" "$disk" >/dev/null 2>&1 || \
	resize_err "Failed deleting partition $partnum on $disk"

    sgdisk -n "$partnum:$start:0" "$disk" >/dev/null 2>&1 || \
	resize_err "Failed recreating partition $partnum on $disk"

    sgdisk -t "$partnum:8300" "$disk" >/dev/null 2>&1 || \
	resize_err "Failed setting partition type on $disk"

    sgdisk -c "$partnum:$label" "$disk" >/dev/null 2>&1 || \
	resize_err "Failed setting partition label on $disk"

    logger $opt -p user.notice -t "$nm" "Partition table updated on $disk"
    partprobe "$disk" 2>/dev/null

    # Mark stage 1 complete, stage 2 will happen after reboot
    echo "1" > /mnt/aux/resized.pending
    sync

    print_end 0 "Resizing /var partition on sdcard, done. Rebooting ..."
    logger $opt -p user.notice -t "$nm" "Partition expanded, rebooting to resize filesystem"

    reboot -f
}

# Stage 2: Complete filesystem resize after partition table has been expanded
# On success, renames "resized.pending" to "resized".
# On failure, renames "resized.pending" to "resized.failed" and returns 1.
resize_filesystem()
{
    label="$1"

    print_start "Expanding /var filesystem, please wait ..."

    devname=$(find_partition_by_label "$label")
    [ -z "$devname" ] && resize_fs_err "Label \"$label\" not found for filesystem resize"

    part="/dev/$devname"

    logger $opt -p user.notice -t "$nm" "Resizing filesystem on $part"
    resize2fs "$part" 2>&1 | logger $opt -p user.notice -t "$nm" || \
	resize_fs_err "Failed resizing filesystem on $part"

    tune2fs -O resize_inode "$part" 2>/dev/null
    print_end 0 "Expanding /var filesystem, done."
    logger $opt -p user.notice -t "$nm" "Filesystem resize complete"

    # Mark resize complete successfully
    mv /mnt/aux/resized.pending /mnt/aux/resized
    sync

    print_restore
    return 0
}

mount_rw()
{
    # If something is already setup, leave it be.
    mountpoint -q "/$1" && return 0

    # If partition doesn't exist (var is optional), signal caller.
    find_partition_by_label "$1" >/dev/null || return 1

    # Check if /var has been resized to fill the sdcard/eMMC
    if [ "$1" = "var" ] && is_mmc; then
	if [ -f /mnt/aux/resized ] || [ -f /mnt/aux/resized.failed ]; then
	    # 3. Already done
	    :
	elif [ -f /mnt/aux/resized.pending ]; then
	    # 2. After reboot we can resize ext4
	    resize_filesystem "$1"
	else
	    # 1. Start with resizing the var partition to the end of the media
	    resize_partition "$1"
	fi
    fi

    # TODO: Also look for UBI partitions

    # Disable periodic fsck, yet keeping safety checks on ext4
    if grep "LABEL=$label" /etc/fstab |grep -q ext4; then
	tune2fs -c 0 -i 0 LABEL="$1" 2>/dev/null
    fi
    mount LABEL="$1" 2>/dev/null && return 0

    return 1
}

# Bind mount /var or /var/lib to ensure applications like podman can set
# up overlay mounts for containers.  Overlays cannot sit on top of other
# overlays.  We always copy the contents of /var or /var/lib from the
# rootfs to get updates, e.g., new packages added in an upgrade.
mount_bind()
{
    src="$1"
    dst="$2"

    mkdir -p "$src"
    cp -af "$dst/." "$src/" ||:

    mount -o bind "$src" "$dst"
}

# mkdir -p -m 0755 only applies to the deepest directory
# shellcheck disable=SC2174
mount_overlay()
{
    tag="$1"
    src="$2"
    dst="$3"
    u="$src/infix/$tag.u"
    w="$src/infix/$tag.w"

    mkdir -p -m 0755 "$u"
    mkdir -p -m 0755 "$w"

    # Ensure that all users in wheel can create the .reset file
    # on /cfg and upload docker images to /var
    chgrp wheel "$(dirname "$u")"

    mount -t overlay "$tag-overlay" "$dst" \
	  -o lowerdir="$dst",upperdir="$u",workdir="$w"
}

# Fall back to console logging if sysklogd is too old
if ! logger -? |grep -q "Log to kernel"; then
    opt="-c"
fi

# On boards with MMC/SD cards, the controller may probe slowly, in particular
# if we netboot (ram load) the device - wait for it if not already available
if is_mmc && ! ls /dev/mmcblk* >/dev/null 2>&1; then
    wait_mmc
fi

# The aux partition must be mounted before everything else since it's used
# for internal bookkeeping.
if ! mount_rw aux >/dev/null 2>&1; then
    logger $opt -p user.warn -t "$nm" \
	   "No auxiliary partition found, software updates not supported."
fi

varsrc=/mnt/var
if ! mount_rw var >/dev/null 2>&1; then
    logger $opt -p user.warn -t "$nm" \
	   "No persistent storage found for /var, only /var/lib is persisted."
    varsrc=/mnt/tmp/infix/var
    vlibsrc=/mnt/cfg/vlib
fi

cfgsrc=/mnt/cfg
etcsrc=/mnt/tmp
if ! mount_rw cfg >/dev/null 2>&1; then
    err=1

    logger $opt -p user.crit -t "$nm" \
	   "No persistent storage found for /cfg, falling back to tmpfs."

    cfgsrc=/mnt/tmp

    # Even if /mnt/var isn't available, if /mnt/cfg isn't either, then
    # there's no point in overlaying one ramdisk on top of another.
    vlibsrc=
fi

if check_factory; then
    factory_reset
fi

mount_overlay cfg  "$cfgsrc" /cfg
mount_overlay etc  "$etcsrc" /etc
mount_overlay home "$cfgsrc" /home
mount_overlay root "$cfgsrc" /root
mount_bind         "$varsrc" /var

[ -n "$vlibsrc" ] && mount_bind "$vlibsrc" /var/lib

for tag in $(ls /sys/bus/virtio/devices/*/mount_tag 2>/dev/null); do
    if [ "$(cat $tag | tr -d '\0')" = hostfs ]; then
	mount hostfs
	break
    fi
done

exit $err
