#!/usr/bin/env bash

set -Eeuo pipefail

export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

LOG_FILE="/tmp/rutoll-auto-restore.log"
exec > >(tee -a "$LOG_FILE") 2>&1

cmdline_value() {
    local key="$1"
    local arg

    for arg in $(cat /proc/cmdline); do
        case "$arg" in
            "${key}="*)
                printf '%s\n' "${arg#*=}"
                return 0
                ;;
        esac
    done

    return 1
}

is_yes() {
    case "${1,,}" in
        1|yes|true|on)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
}

IMAGE="$(cmdline_value rutoll_image || true)"
IMAGE="${IMAGE:-}"

MIN_GB="$(cmdline_value rutoll_min_gb || true)"
MIN_GB="${MIN_GB:-50}"

TARGET_OVERRIDE="$(cmdline_value rutoll_target || true)"
TARGET_OVERRIDE="${TARGET_OVERRIDE:-auto}"

CONFIRM="$(cmdline_value rutoll_confirm || true)"
CONFIRM="${CONFIRM:-yes}"

ALLOW_USB="$(cmdline_value rutoll_allow_usb || true)"
ALLOW_USB="${ALLOW_USB:-no}"

ALLOW_REMOVABLE="$(cmdline_value rutoll_allow_removable || true)"
ALLOW_REMOVABLE="${ALLOW_REMOVABLE:-no}"

DRY_RUN="$(cmdline_value rutoll_dry_run || true)"
DRY_RUN="${DRY_RUN:-no}"

if [[ -z "$IMAGE" ]]; then
    echo "ERROR: rutoll_image is not set in kernel cmdline."
    echo "Example: rutoll_image=LATM-v1.13"
    exit 10
fi

if ! [[ "$MIN_GB" =~ ^[0-9]+$ ]]; then
    echo "ERROR: rutoll_min_gb must be numeric."
    echo "Current value: $MIN_GB"
    exit 11
fi

MIN_BYTES="$((MIN_GB * 1000 * 1000 * 1000))"

echo
echo "============================================================"
echo " RUTOLL Clonezilla automatic restore"
echo "============================================================"
echo
echo "Image:              $IMAGE"
echo "Minimum disk size:  ${MIN_GB} GB"
echo "Target override:    $TARGET_OVERRIDE"
echo "Allow USB target:   $ALLOW_USB"
echo "Allow removable:    $ALLOW_REMOVABLE"
echo "Confirmation:       $CONFIRM"
echo "Dry run:            $DRY_RUN"
echo

echo "Current block devices:"
lsblk -o NAME,SIZE,TYPE,FSTYPE,MODEL,SERIAL,TRAN,RM,RO,MOUNTPOINTS || true
echo

if ! mountpoint -q /home/partimag; then
    echo "ERROR: /home/partimag is not mounted."
    echo "Check ocs_prerun2 NFS mount command."
    exit 20
fi

if [[ ! -d "/home/partimag/${IMAGE}" ]]; then
    echo "ERROR: Clonezilla image directory not found:"
    echo "/home/partimag/${IMAGE}"
    echo
    echo "Available images:"
    find /home/partimag \
        -mindepth 1 \
        -maxdepth 1 \
        -type d \
        -printf '  %f\n' \
        | sort || true
    exit 21
fi

declare -a DISKS=()

while IFS= read -r line; do
    unset NAME TYPE RM RO TRAN SIZE

    eval "$line"

    [[ "${TYPE:-}" == "disk" ]] || continue
    [[ "${RO:-}" == "0" ]] || continue

    case "${NAME:-}" in
        loop*|ram*|zram*|sr*|fd*)
            continue
            ;;
    esac

    if ! is_yes "$ALLOW_REMOVABLE" && [[ "${RM:-0}" == "1" ]]; then
        continue
    fi

    if ! is_yes "$ALLOW_USB" && [[ "${TRAN:-}" == "usb" ]]; then
        continue
    fi

    if [[ "${SIZE:-0}" -lt "$MIN_BYTES" ]]; then
        continue
    fi

    DISKS+=("/dev/${NAME}")

done < <(
    lsblk -P -b -dn -o NAME,TYPE,RM,RO,TRAN,SIZE
)

if [[ "$TARGET_OVERRIDE" != "auto" ]]; then
    TARGET="/dev/${TARGET_OVERRIDE#/dev/}"

    if [[ ! -b "$TARGET" ]]; then
        echo "ERROR: Requested target does not exist: $TARGET"
        exit 30
    fi

    if [[ "$(lsblk -dn -o TYPE "$TARGET")" != "disk" ]]; then
        echo "ERROR: Requested target is not a whole disk: $TARGET"
        exit 31
    fi

    TARGET_BYTES="$(blockdev --getsize64 "$TARGET")"

    if [[ "$TARGET_BYTES" -lt "$MIN_BYTES" ]]; then
        echo "ERROR: Requested target is smaller than ${MIN_GB} GB."
        echo "Target: $TARGET"
        exit 32
    fi
else
    case "${#DISKS[@]}" in
        0)
            echo "ERROR: No suitable internal disk found."
            echo
            echo "Rules:"
            echo "  - type must be disk"
            echo "  - read-only must be 0"
            echo "  - removable is excluded by default"
            echo "  - USB is excluded by default"
            echo "  - size must be >= ${MIN_GB} GB"
            echo
            lsblk -o NAME,SIZE,TYPE,FSTYPE,MODEL,SERIAL,TRAN,RM,RO,MOUNTPOINTS || true
            exit 40
            ;;

        1)
            TARGET="${DISKS[0]}"
            ;;

        *)
            echo "More than one suitable disk found:"
            echo

            index=1
            for disk in "${DISKS[@]}"; do
                echo "[$index] $disk"
                lsblk -dn -o NAME,SIZE,MODEL,SERIAL,TRAN "$disk" || true
                echo
                index="$((index + 1))"
            done

            echo "[0] ABORT"
            echo

            while true; do
                read -r -p "Select target disk number: " choice

                if [[ "$choice" == "0" ]]; then
                    echo "Restore cancelled."
                    exit 41
                fi

                if [[ "$choice" =~ ^[0-9]+$ ]] \
                    && (( choice >= 1 )) \
                    && (( choice <= ${#DISKS[@]} )); then
                    TARGET="${DISKS[$((choice - 1))]}"
                    break
                fi

                echo "Invalid selection."
            done
            ;;
    esac
fi

TARGET_NAME="${TARGET#/dev/}"
TARGET_BYTES="$(blockdev --getsize64 "$TARGET")"
TARGET_GB="$((TARGET_BYTES / 1000 / 1000 / 1000))"

echo
echo "Selected target disk:"
echo

lsblk -dn -o NAME,SIZE,MODEL,SERIAL,TRAN "$TARGET" || true

echo
echo "Target device: $TARGET"
echo "Target size:   ${TARGET_GB} GB"
echo

if [[ "$TARGET_BYTES" -lt "$MIN_BYTES" ]]; then
    echo "ERROR: Target disk is too small."
    echo "Target:   ${TARGET_GB} GB"
    echo "Required: ${MIN_GB} GB"
    exit 50
fi

echo "WARNING:"
echo "ALL DATA ON $TARGET WILL BE DESTROYED."
echo

if is_yes "$CONFIRM"; then
    EXPECTED="ERASE-${TARGET_NAME}"

    echo "To continue, type exactly:"
    echo
    echo "$EXPECTED"
    echo

    read -r -p "> " ANSWER

    if [[ "$ANSWER" != "$EXPECTED" ]]; then
        echo "Confirmation failed. Restore cancelled."
        exit 51
    fi
fi

OCS_CMD=(
    ocs-sr
    -batch
    -g auto
    -e1 auto
    -e2
    -r
    -j2
    -k1
    -icds
    -scr
    -p reboot
    restoredisk
    "$IMAGE"
    "$TARGET_NAME"
)

echo
echo "Clonezilla command:"
printf ' %q' "${OCS_CMD[@]}"
echo
echo

if is_yes "$DRY_RUN"; then
    echo "DRY RUN enabled. Nothing will be written."
    exit 0
fi

exec "${OCS_CMD[@]}"