#!/bin/sh

#set -x

# 1. Bash scripts are not supported in busybox
# 2. There is no lsblk command in the busybox environment, only blkid command
# 3. In the busybox environment, blkid does not support additional parameters, only blkid or blkid [device]
# 4. In POSIX sh, 'declare' and 'local' is undefined
# 5. grep supports -E parameters in busybox environment, not -P parameters

VERSION="1.0.0"

show_version() {
    echo "$VERSION"
}

get_fstype() {
    local DEVICE
    DEVICE=$1
    local type_val=$(blkid "$DEVICE" | grep -oE 'TYPE="[^"]+"')
    type_val=${type_val#TYPE=\"}
    type_val=${type_val%\"}
    echo $type_val
}

resolve_device() {
    DEV="$1"

	case "$DEV" in
	LABEL=* | UUID=* | PARTLABEL=* | PARTUUID=*)
		# For disk encryption, take the last result from blkid output for the same disk
		DEV=$(blkid -t "${DEV}" -o device | tail -n 1) || return 1
		;;
	esac
	[ -e "$DEV" ] && echo "$DEV"
}

# Add a function to wait for device using udevadm
wait_for_device() {
    local device=$1
    local timeout=10
    local dev_path=""
    
    case "$device" in
    LABEL=*)
        dev_path="/dev/disk/by-label/${device#LABEL=}"
        ;;
    UUID=*)
        dev_path="/dev/disk/by-uuid/${device#UUID=}"
        ;;
    PARTLABEL=*)
        dev_path="/dev/disk/by-partlabel/${device#PARTLABEL=}"
        ;;
    PARTUUID=*)
        dev_path="/dev/disk/by-partuuid/${device#PARTUUID=}"
        ;;
    *)
        # For direct device paths
        dev_path="$device"
        ;;
    esac
    
    if [ -n "$dev_path" ] && [ ! -e "$dev_path" ]; then
        echo "Waiting for device $device to appear..." >&2
        if ! udevadm wait --timeout=$timeout "$dev_path" >/dev/null 2>&1; then
            echo "Timeout waiting for device $device ($dev_path)" >&2
            return 1
        fi
    fi
}

try_get_ostree() {
    local rootmnt
    rootmnt=$1
    # get ostree from cmdline
    local OSTREE BOOT_IMAGE
    for param in $(cat /proc/cmdline)
    do
        case "${param}" in
        ostree=*)
            OSTREE="${param#ostree=}"
            ;;
        BOOT_IMAGE=*)
            BOOT_IMAGE="${param#BOOT_IMAGE=}"
            ;;
        esac
    done

    # If OSTREE is set and is not "auto", use it directly
    if [ -n "$OSTREE" ] && [ "$OSTREE" != "auto" ] ;then
        echo $OSTREE
        return
    fi

    # If BOOT_IMAGE exists and matches the required format, extract the hash value
    if [ -n "$BOOT_IMAGE" ]; then
        # Check if BOOT_IMAGE matches the format /boot/immutable/vmlinuz-*
        if echo "$BOOT_IMAGE" | grep -q "^/boot/immutable/.*/vmlinuz-"; then
            local deploy_id
            deploy_id=$(echo "$BOOT_IMAGE" | sed -E 's|^/boot/immutable/([^/]+)(\.[0-9]+)/vmlinuz-.*|\1\2|')
            if [ -n "$deploy_id" ]; then
                OSTREE="/ostree/data/$deploy_id/checkout"
                echo $OSTREE
                return
            fi
        fi
    fi

    # get ostree from status.list
    if [ -f $rootmnt/persistent/ostree/data/status.list ];then
        local deploy_id=$(cat $rootmnt/persistent/ostree/data/status.list| head -n 1)
        echo "/ostree/data/$deploy_id/checkout"
    else
        echo ""
    fi
}

is_mountpoint() {
    path=$1
    mount_infos=$(grep -E "[0-9]{1,2}:[0-9]{1,2}[[:space:]]/[^[:space:]]*[[:space:]]$path[[:space:]]" /proc/self/mountinfo || echo "")
    if [ -z "$mount_infos" ];then
        return 1
    fi
    return 0
}

recursive_umount() {
    local rootmnt=$1

    # Get all mount point information and filter out those under /rootmnt
    mount_points=$(grep -E "[0-9]{1,2}:[0-9]{1,2}[[:space:]]/[^[:space:]]*[[:space:]]$rootmnt(/[^[:space:]]*)?[[:space:]]" /proc/self/mountinfo | awk '{print $5}')

    # Remove duplicate mount points
    mount_points=$(echo "$mount_points" | sort -u)

    # Sort mount points by depth to ensure child mounts are unmounted first
    sorted_mount_points=$(echo "$mount_points" | awk -F'/' '{print NF, $0}' | sort -rn | awk '{print $2}')
    echo "mount points: $sorted_mount_points"

    # Unmount each mount point
    for mount_point in $sorted_mount_points; do
        while is_mountpoint "$mount_point"; do
            echo "Unmounting $mount_point"
            umount "$mount_point"
        done
    done
}

ext_mount() {
    local rootmnt=$1
    local OSTREE=$2
    local lowerdir=$3
    local datadir=$rootmnt/persistent
    local extDeploy=$(realpath $datadir/$OSTREE)
    [ ! -d $datadir/overlay/work ] && mkdir -p $datadir/overlay/work
    [ ! -d $datadir/overlay/merged ] && mkdir -p $datadir/overlay/merged
    local otDeploy=$datadir/overlay/merged
    # remove opt symlink if exists
    if [ -L "$lowerdir/opt" ]; then
        rm "$lowerdir/opt"
    fi
    mkdir -p "$lowerdir/opt"
    modprobe overlay
    mount -t overlay -o lowerdir="$lowerdir",upperdir="$extDeploy",workdir="$datadir/overlay/work" overlay $otDeploy
    mount --make-private $otDeploy
    mount -o bind,rw $otDeploy/etc $otDeploy/etc
    mount --bind $datadir/ostree/deploy/deepin/var $otDeploy/var
    mkdir -p "$otDeploy/var/mnt" || true
    mkdir -p $datadir/overlay/usr-work
    mkdir -p $datadir/overlay/usr-upper
    # Handle /opt
    mkdir -p $datadir/overlay/opt-work
    mkdir -p $datadir/overlay/opt-upper
  
    mount -t overlay -o rw,relatime,lowerdir=$extDeploy/usr/opt:$lowerdir/usr/opt,upperdir=$datadir/overlay/opt-upper,workdir=$datadir/overlay/opt-work opt-overlay $otDeploy/opt
    mount -t overlay -o ro,relatime,lowerdir=$otDeploy/usr,upperdir=$datadir/overlay/usr-upper,workdir=$datadir/overlay/usr-work usr-overlay $otDeploy/usr
    mount -o bind $datadir $otDeploy/persistent
    mount -o bind $rootmnt $otDeploy/sysroot
    mount --bind -o remount,ro $otDeploy/sysroot $otDeploy/sysroot
    # chroot to roomnt and start init
    mount --make-private $rootmnt
    mount --move $otDeploy $rootmnt
}

run_hooks() {
    if [ "$MOUNT_ROOT_HOOK_ENABLED" != 1 ]; then
        return
    fi

    local hooks_dir="/usr/share/deepin-immutable-mount-root-hooks/$1"
    shift
    if [ ! -d "$hooks_dir" ]; then
        echo "Hooks directory $hooks_dir does not exist" >&2
        return
    fi
    # iterate hooks_dir
    for hook_file in "$hooks_dir"/*; do
        if [ -x "$hook_file" ]; then
            # check hook file is owned by root
            local file_stat
            file_stat=$(stat -c '%u:%g' "$hook_file" 2>/dev/null) || {
                echo "Unable to get file status for $hook_file" >&2
                continue
            }
            if [ "$file_stat" != "0:0" ]; then
                echo "Skipping hook $hook_file: not owned by root (uid:gid=$file_stat)" >&2
                continue
            fi
            echo "Executing hook: $hook_file"
            "$hook_file" "$@" || echo "Failed to run hook $hook_file" >&2
        fi
    done
}

data_layer_mount() {
    local rootmnt=$1
    local OSTREE=$2
    local baseDeploy=$3

    local datadir=$rootmnt/persistent
    local extDeploy=$(realpath $datadir/$OSTREE)

    # remove opt symlink if exists
    if [ -L "$rootmnt/opt" ]; then
        rm "$rootmnt/opt"
    fi

    if [ -L "$rootmnt/media" ]; then
        rm "$rootmnt/media"
    fi

    mkdir -p $rootmnt/boot \
            $rootmnt/dev \
            $rootmnt/etc \
            $rootmnt/home \
            $rootmnt/mnt \
            $rootmnt/media \
            $rootmnt/opt \
            $rootmnt/proc \
            $rootmnt/root \
            $rootmnt/run \
            $rootmnt/srv \
            $rootmnt/sys \
            $rootmnt/sysroot \
            $rootmnt/tmp \
            $rootmnt/usr \
            $rootmnt/var

    chmod 1777 $rootmnt/tmp

    # Create symlinks for bin, lib, lib64, and sbin if they don't exist
    [ ! -e "$rootmnt/bin" ] && [ ! -L "$rootmnt/bin" ] && ln -s usr/bin $rootmnt/bin
    [ ! -e "$rootmnt/lib" ] && [ ! -L "$rootmnt/lib" ] && ln -s usr/lib $rootmnt/lib
    [ ! -e "$rootmnt/lib64" ] && [ ! -L "$rootmnt/lib64" ] && ln -s usr/lib64 $rootmnt/lib64
    [ ! -e "$rootmnt/sbin" ] && [ ! -L "$rootmnt/sbin" ] && ln -s usr/sbin $rootmnt/sbin

    # When the sysroot partition is the actual root or just a regular directory, 
    # binding /sysroot/boot to /boot might change the actual /boot directory, causing issues. 
    # Use a symlink to keep the data in sync.
    if grep -q "/sysroot/boot /boot none defaults,bind,rw 0 0" /etc/fstab; then
        ln -s /boot $rootmnt/sysroot/boot
    fi

    # Remove the last slash and everything after it
    data_commitid_idx="${extDeploy%/*}"
    # Get the content after the last slash
    data_commitid_idx="${data_commitid_idx##*/}"
    data_overlay_root=$rootmnt/persistent/overlay/data/$data_commitid_idx
    mkdir -p $data_overlay_root

    cfgFile=$(dirname $extDeploy)/config
    if [ -f "$cfgFile" ]; then
        while IFS='=' read -r key value; do
            # Ignore empty lines and comment lines
            if [ -z "$key" ] || [ "${key#\#}" != "$key" ]; then
                continue
            fi

            # Remove leading and trailing spaces
            key=$(echo "$key" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
            value=$(echo "$value" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
            # Only process upper and lower configuration items
            if [ "$key" = "upper" ] || [ "$key" = "lower" ]; then
                eval "cfg_$key=$value"
            fi
        done < "$cfgFile"
    fi
    if [ -z "$cfg_upper" ]; then
        cfg_upper=default
    fi
    if [ -z "$cfg_lower" ]; then
        cfg_lower=none
    fi

    modprobe overlay
    for mountParam in usr:ro opt:rw etc:rw;do
        # ash in busybox does not support <<< syntax.
        # `echo "$mountParam" | IFS=':' read -r dir perm` will also have problems, 
        # because when read runs in a pipeline, it will run a shell process separately, resulting in the external failure to use dir and perm
        IFS=':' read -r dir perm << EOF
$mountParam
EOF
        echo "do data layer mount $dir"
        mkdir -p "$baseDeploy/$dir"

        if [ "$cfg_upper" = "default" ]; then
            upper_dir=$data_overlay_root/$dir-upper
        else
            upper_dir=$rootmnt/persistent/overlay/data/layer-$cfg_upper/$dir
        fi
        if [ "$cfg_lower" = "none" ]; then
            lower_dir=$extDeploy/$dir:$baseDeploy/$dir
        else
            if [ "$cfg_lower" = "default" ]; then
                first_lower_dir=$data_overlay_root/$dir-upper
                mkdir -p "$first_lower_dir"
                lower_dir=$first_lower_dir:$extDeploy/$dir:$baseDeploy/$dir
            else
                first_lower_dir=$rootmnt/persistent/overlay/data/layer-$cfg_lower/$dir
                mkdir -p "$first_lower_dir"
                lower_dir=$first_lower_dir:$extDeploy/$dir:$baseDeploy/$dir
            fi
            if [ "$dir" = usr ]; then
                # CRITICAL WARNING: DO NOT CHANGE THE ORDER OF DIRECTORIES IN lower_dir!
                # The overlay merge logic in overlay merger depends on this specific order.
                # Changing this order will break the overlay merge functionality!
                run_hooks before_mount_usr_overlay "$(dirname $extDeploy)" "$lower_dir" "$upper_dir"
            fi
        fi
        workDir=$data_overlay_root/$dir-work
        mkdir -p "$upper_dir" "$workDir" $rootmnt/$dir
        LIBMOUNT_FORCE_MOUNT2=always mount -t overlay -o "$perm,relatime,lowerdir=$lower_dir,upperdir=$upper_dir,workdir=$workDir" $dir-overlay $rootmnt/$dir
    done

    # make etc without hardlink, some apps not work well with hardlink
    hardlink_count=$(stat -c "%h" $rootmnt/etc/deepin-immutable-ctl/deepin-immutable-ctl.conf)
    if [ "$hardlink_count" -gt 1 ] && [ -f /usr/bin/cp.fix ]; then
        tmp_copyed=$rootmnt/persistent/ostree/data/tmp
        mkdir -p $tmp_copyed
        /usr/bin/cp.fix --preserve=all -r $rootmnt/etc $tmp_copyed || true
        rm -rf $rootmnt/etc
        /usr/bin/cp.fix --preserve=all -r $tmp_copyed/etc $rootmnt || true
        rm -rf $tmp_copyed
    fi

    mount -o bind $rootmnt/sysroot $rootmnt/sysroot
    mkdir -p $rootmnt/sysroot/ostree
    mount -o bind $rootmnt/ostree $rootmnt/sysroot/ostree
    mount -o remount,bind,ro $rootmnt/sysroot/ostree $rootmnt/sysroot/ostree

    mount -o bind $rootmnt/ostree $rootmnt/ostree
    mount -o remount,bind,ro $rootmnt/ostree $rootmnt/ostree

    mount -o bind $rootmnt/persistent/ostree $rootmnt/persistent/ostree
    mount -o remount,bind,ro $rootmnt/persistent/ostree $rootmnt/persistent/ostree

    mount -o bind $datadir/ostree/deploy/deepin/var $rootmnt/var
    mount -o remount,bind,rw $rootmnt/var $rootmnt/var

    run_hooks before_mount_var_overlay "$data_overlay_root"

    if [ -d "$data_overlay_root/var-upper" ]; then
        mkdir -p "$data_overlay_root/var-work"
        mount -t overlay -o "rw,relatime,lowerdir=$rootmnt/var,upperdir=$data_overlay_root/var-upper,workdir=$data_overlay_root/var-work" var-overlay "$rootmnt/var"
    fi
    # chroot to roomnt and start init
}

ostree_mount_root() {
    local rootmnt=$1
    local OSTREE=$2

    local datadir=$rootmnt/persistent

    local extDeploy=$(realpath $datadir/$OSTREE)
    local sys_commit_idx
    local ostree_parent_path="$extDeploy/usr/share/deepin-immutable-ctl/state/ostree-parent"
    sys_commit_idx=$(cat "$ostree_parent_path")

    sysroot="$rootmnt"
    ostree_path=/ostree/deploy/deepin/deploy/$sys_commit_idx

    case $extDeploy in
        */ostree/data/*)
            data_layer_mount $rootmnt "$OSTREE" "$sysroot$ostree_path"
            ;;
        *)
            ext_mount $rootmnt "$OSTREE" "$sysroot$ostree_path"
            ;;
    esac
}

# Parse mountroot --root=uuid=xxx --persistent=uuid=xxx --rootflags=xxxx --persistentflags=xxxx --ostree=xxxx /root using getopt.
mountroot() {
    local ROOT ROOTFLAGS PERSISTENT PERSISTENTFLAGS OSTREE rootmnt
    local OPTIND
    local OPTS
    local UMOUNT=false
    local HAS_PERSISTENT=false

    OPTS=$(getopt -o r:p:f:F:o:uhv --long root:,persistent:,rootflags:,persistentflags:,ostree:,umount,help,version -- "$@")
    if [ $? != 0 ]; then
        echo "Failed parsing options." >&2
        return 1
    fi

    eval set -- "$OPTS"

    while true; do
        case "$1" in
            -r | --root)
                ROOT=$2
                shift 2
                ;;
            -p | --persistent)
                HAS_PERSISTENT=true
                PERSISTENT=$2
                shift 2
                ;;
            -f | --rootflags)
                ROOTFLAGS="-o $2"
                shift 2
                ;;
            -F | --persistentflags)
                PERSISTENTFLAGS="-o $2"
                shift 2
                ;;
            -o | --ostree)
                OSTREE=$2
                shift 2
                ;;
            -u | --umount)
                UMOUNT=true
                shift 1
                ;;
            -v | --version)
                show_version
                return 0
                ;;
            -h | --help)
                echo "Usage: $0 [OPTIONS] rootmnt"
                echo
                echo "Options:"
                echo "  -r, --root=DEVICE          Root device (e.g., UUID=xxx)"
                echo "  -p, --persistent=DEVICE    Persistent device (e.g., UUID=xxx)"
                echo "  -f, --rootflags=FLAGS      Flags for root device (e.g., rw)"
                echo "  -F, --persistentflags=FLAGS Flags for persistent device (e.g., rw)"
                echo "  -o, --ostree=PATH          OSTree path"
                echo "  -u, --umount               Unmount all mount points under rootmnt"
                echo "  -v, --version              Show version"
                echo "  -h, --help                 Show this help message"
                return 0
                ;;
            --)
                shift
                break
                ;;
            *)
                echo "Invalid option: $1" >&2
                return 1
                ;;
        esac
    done

    # get rootmnt param
    if [ $# -eq 1 ]; then
        rootmnt=$1
    elif [ $# -eq 2 ];then
        rootmnt=$1
        OSTREE=$2
    else
        echo "Missing rootmnt argument." >&2
        return 1
    fi

    if $UMOUNT; then
        recursive_umount "$rootmnt"
        return 0
    fi

    mkdir -p $rootmnt || return 1

    if [ -n "$ROOT" ];then
        # Wait for device to be available
        wait_for_device "$ROOT"
        ROOT=$(resolve_device "$ROOT")
        if [ -z "$ROOT" ]; then
            echo "Failed to resolve root device: $ROOT" >&2
            return 1
        fi
        ROOTFSTYPE=$(get_fstype "$ROOT")
        if ! is_mountpoint "$rootmnt";then
            echo "mount -t $ROOTFSTYPE $ROOTFLAGS $ROOT $rootmnt"
            mount -t "$ROOTFSTYPE" $ROOTFLAGS "$ROOT" "$rootmnt"
        fi
    fi

    mount -o remount,rw $rootmnt

    mkdir -p "$rootmnt/persistent"
    if [ -n "$PERSISTENT" ]; then
        # Wait for device to be available
        wait_for_device "$PERSISTENT"
        PERSISTENT=$(resolve_device "$PERSISTENT")
        if [ -z "$PERSISTENT" ]; then
            echo "Failed to resolve persistent device: $PERSISTENT" >&2
            return 1
        fi
        if [ -n "$PERSISTENT" ];then
            PERSISTENTFSTYPE=$(get_fstype "$PERSISTENT")
            if ! is_mountpoint "$rootmnt/persistent";then
                echo "mount -t $PERSISTENTFSTYPE $PERSISTENTFLAGS $PERSISTENT $rootmnt/persistent"
                mount -t "$PERSISTENTFSTYPE" $PERSISTENTFLAGS "$PERSISTENT" "$rootmnt/persistent"
            fi
        fi
    fi

    if [ -z "$OSTREE" ]; then
        OSTREE=$(try_get_ostree "$rootmnt")
    fi
    echo "OSTREE=$OSTREE"
    ostree_mount_root "$rootmnt" "$OSTREE"
}

# 1. mount root device to $mnt
# 2. mount data device to $mnt/persistent

#mount /dev/vda4 $mnt
#mount /dev/vda5 $mnt/persistent

#ostree_mount_root $ostree_arg
mountroot $@

# you can do chroot here