#!/bin/bash

set -euo pipefail
shopt -s nullglob

detect_target_device() {
    # Pick the first non-removable, non-virtual block device. Prefer NVMe
    # (the 410b/520 default), fall back to SATA/SCSI on older appliances.
    local usb_wic_installer_disk
    usb_wic_installer_disk="$(detect_usb_wic_installer_disk || true)"

    local d
    for d in /sys/block/nvme*n* /sys/block/sd*; do
        d="${d##*/}"

        # Only install to whole block devices that exist under /dev.
        # Partitions, missing device nodes, and stale sysfs entries are not
        # valid targets for writing a full disk image.
        [ -b "/dev/$d" ] || continue

        # Real hardware disks have a backing device in sysfs. Virtual block
        # devices such as loop, dm, and zram must never be auto-selected for
        # the destructive installer write.
        [ -e "/sys/block/$d/device" ] || continue

        # Removable media is usually the installer USB stick. It is not a
        # reliable target for Edge OS, and overwriting it can kill the running
        # installer before the appliance disk has been installed.
        [ "$(cat "/sys/block/$d/removable" 2>/dev/null)" = "0" ] || continue

        # Some USB media reports itself as non-removable. When booting the
        # installer.wic image, skip the disk that contains its installroot
        # partition even if the kernel exposes it as /dev/sdX.
        [ "$d" != "$usb_wic_installer_disk" ] || continue

        echo "/dev/$d"
        return 0
    done
    return 1
}

detect_usb_wic_installer_disk() {
    # Only the bootable installer.wic image uses root=PARTLABEL=installroot.
    # In-place installer flows run from an initramfs and must still be allowed
    # to overwrite the disk they were staged on.
    local root
    local part
    local dev

    root="$(awk 'BEGIN { RS=" "; FS="=" } $1 == "root" { print $2; exit }' < /proc/cmdline)"
    [ "$root" = "PARTLABEL=installroot" ] || return 1

    part="$(readlink -f /dev/disk/by-partlabel/installroot 2>/dev/null)"

    dev="${part##*/}"
    [ -e "/sys/class/block/$dev/partition" ] || return 1
    basename "$(dirname "$(readlink -f "/sys/class/block/$dev")")"
}

verify_connectivity() {
    local url="$1"

    echo >&2 "Verifying internet connectivity to $url"

    for attempt in {1..100}; do
        if wget --spider -T 2 "$url" 2>/dev/null; then
            echo >&2 "Internet connectivity verified"
            return 0
        fi

        echo >&2 "[$attempt/100] Connectivity check failed"

        sleep 2
    done

    echo >&2 "Error: No internet connectivity"
    echo >&2 "Please check network configuration and internet access"
    return 1
}

PAYLOAD_DIR="opt/edgeos/payload"

# pv redraws its progress bar with carriage returns; through the journal each
# redraw becomes a separate log line, so write to a console device directly.
# Prefer a serial port: IPMI SOL is ttyS1 and /dev/console never points there.
CONSOLE=/dev/stderr
if [ -c /dev/console ] && (: >/dev/console) 2>/dev/null; then
    CONSOLE=/dev/console
fi
for c in $(cat /sys/class/tty/console/active 2>/dev/null) ttyS1 ttyS0; do
    case "$c" in
    ttyS*)
        if [ -c "/dev/$c" ] && (: >"/dev/$c") 2>/dev/null; then
            CONSOLE="/dev/$c"
            break
        fi
        ;;
    esac
done

disk_of_partition() {
    local part base parent
    part="$1"
    base="${part##*/}"
    if [ -e "/sys/class/block/$base/partition" ]; then
        parent="$(basename "$(dirname "$(readlink -f "/sys/class/block/$base")")")"
        echo "/dev/$parent"
    else
        echo "$part"
    fi
}

installer_version() {
    [ -f /etc/edge-installer.conf ] || { echo "unknown"; return; }
    local v
    v=$(awk -F= '/^EDGE_INSTALLER_VERSION=/ { print $2; exit }' /etc/edge-installer.conf)
    echo "${v:-unknown}"
}

print_install_banner() {
    {
        echo
        echo "=================================================="
        echo "  Installing Edge OS"
        echo "    Image:   ${1:-unknown}"
        echo "    Variant: ${2:-unknown}"
        echo "    Version: ${3:-unknown}"
        echo "    Target:  ${4}"
        echo "=================================================="
        echo
    } > "$CONSOLE"
}

write_image() {
    local src="$1" device="$2" m
    # A read-write mount on the target disk flushes ext4 metadata at shutdown,
    # after dd replaced the partition layout, corrupting the new image. The
    # installer-bundle flow runs from this disk and cannot unmount its own
    # root, so drop everything to read-only instead.
    for m in $(awk -v d="^$device" '$1 ~ d && $2 != "/" {print $2}' /proc/mounts | sort -r); do
        umount "$m" 2>/dev/null || mount -o remount,ro "$m" 2>/dev/null || true
    done
    if awk -v d="^$device" '$1 ~ d && $2 == "/" {found=1} END {exit !found}' /proc/mounts; then
        mount -o remount,ro / 2>/dev/null || true
    fi
    sync
    if awk -v d="^$device" '$1 ~ d && $4 !~ /(^|,)ro(,|$)/ {found=1} END {exit !found}' /proc/mounts; then
        echo >&2 "Error: $device still has read-write mounts, refusing to overwrite:"
        awk -v d="^$device" '$1 ~ d {print}' /proc/mounts >&2
        exit 1
    fi
    pv -f -i 1 "$src" 2>"$CONSOLE" | unzstd | dd of="$device" bs=4M
    sync # dd leaves tens of GB in the page cache; shutdown may not drain it
}

install_from_local_payload() {
    local device="$1"
    local payload="" src_part="" f
    local mnt="/run/edge-installer-payload"

    # The USB installer image boots with the payload baked into its rootfs
    # (edge-payload.bbclass). In-place flows (RAUC installer bundle, NimOS
    # upgrader) boot the installer initramfs instead, with the payload staged
    # on a disk partition, so scan and mount if the running rootfs has none.
    for f in /${PAYLOAD_DIR}/*.rootfs.wic.zst; do
        payload="$f"
        break
    done

    if [ -z "$payload" ]; then
        mkdir -p "$mnt"
        local part
        for part in /dev/nvme*n*p* /dev/sd*[0-9] /dev/mmcblk*p*; do
            [ -b "$part" ] || continue
            # noload: a read-only ext4 mount still writes journal recovery
            mount -o ro,noload "$part" "$mnt" 2>/dev/null || mount -o ro "$part" "$mnt" 2>/dev/null || continue
            for f in "$mnt"/${PAYLOAD_DIR}/*.rootfs.wic.zst; do
                payload="$f"
                src_part="$part"
                break
            done
            [ -n "$payload" ] && break
            umount "$mnt" 2>/dev/null || true
        done
    fi

    [ -n "$payload" ] || return 1

    local pbase image variant
    pbase="$(basename "$payload")"
    case "$pbase" in
        edge-core-image*)    image="edge-core" ;;
        edge-connect-image*) image="edge-connect" ;;
        *)                   image="${pbase%%-image*}" ;;
    esac
    case "$pbase" in
        *-image-dev-*) variant="dev" ;;
        *)             variant="prod" ;;
    esac
    print_install_banner "$image" "$variant" "$(installer_version)" "$device"
    echo >&2 "Embedded rootfs payload: $payload"

    # In-place flows stage the payload on the disk we are about to overwrite;
    # copy it to RAM first so dd doesn't destroy it mid-read.
    if [ -n "$src_part" ] && [ "$(disk_of_partition "$src_part")" = "$device" ]; then
        local cache_mnt="/run/edge-installer-cache"
        local cache="$cache_mnt/payload.wic.zst"
        local need_kb=$(( ($(wc -c < "$payload") / 1024) + 262144 ))
        echo >&2 "Payload is on the target disk; caching ${need_kb}KB to RAM"
        mkdir -p "$cache_mnt"
        mount -t tmpfs -o "size=${need_kb}k" tmpfs "$cache_mnt"
        cp "$payload" "$cache"
        umount "$mnt" 2>/dev/null || true
        sync # drain the source mount's writeback before overwriting its disk
        write_image "$cache" "$device"
        umount "$cache_mnt" 2>/dev/null || true
    else
        write_image "$payload" "$device"
        [ -n "$src_part" ] && umount "$mnt" 2>/dev/null || true
    fi

    return 0
}

# Use provided device or auto-detect the appliance's internal disk.
if [ $# -eq 1 ]; then
    device="$1"
    echo >&2 "Using specified installation target: $device"
else
    device=$(detect_target_device || true)
    if [ -z "$device" ]; then
        echo >&2 "Error: No internal disk found for auto-detection"
        echo >&2 "Usage: $0 [device]"
        echo >&2 "  device: target block device (auto-detected if not specified)"
        exit 1
    fi
    echo >&2 "Auto-detected installation target: $device"
fi

# Verify device exists and is a block device
if [ ! -b "$device" ]; then
    echo >&2 "Error: $device is not a valid block device"
    exit 1
fi

if install_from_local_payload "$device"; then
    echo >&2 "Installed Edge OS from embedded payload"
else
    image=""
    version=""
    env=""

    if [ -f /etc/edge-installer.conf ]; then
        # shellcheck disable=SC1091
        . /etc/edge-installer.conf
        image="${EDGE_INSTALLER_IMAGE:-}"
        version="${EDGE_INSTALLER_VERSION:-}"
    fi

    # Fall back to /proc/cmdline (PXE flow). Each lookup only runs if the conf
    # file didn't supply that field, so a partial conf still works.
    [ -n "$version" ] || version=$(awk 'BEGIN { RS=" "; FS="=" } /edgeos.version/ { print $2 }' < /proc/cmdline)
    [ -n "$env" ]     || env=$(awk 'BEGIN { RS=" "; FS="=" } /edgeos.env/ { print $2 }' < /proc/cmdline)
    [ -n "$image" ]   || image=$(awk 'BEGIN { RS=" "; FS="=" } /edgeos.image/ { print $2 }' < /proc/cmdline)

    [ -n "${env:-}" ] || env="${EDGE_INSTALLER_ENV:-}"
    if [ -z "${env:-}" ]; then
        case "$version" in
            ""|R*) env="prod" ;;
            *)     env="dev" ;;
        esac
    fi

    if [ -z "$version" ]; then
        echo "Warning: edgeos.version not found in kernel cmdline, using 'latest'"
        version="latest"
    fi

    if [ -z "$image" ]; then
        echo "Warning: edgeos.image not found in kernel cmdline, defaulting to 'edge-connect'"
        image="edge-connect"
    fi

    print_install_banner "$image" "$env" "$version" "$device"

    if [ "$env" = "dev" ]; then
        artifacts=artifacts.nimbra.dev
        url="https://$artifacts/api/artifacts/edgeos/v1/${image}-image-dev/$version/${image}-image-dev-nimbra.rootfs.wic.zst"
    else
        artifacts=artifacts.nimbra.io
        url="https://$artifacts/api/artifacts/edgeos/v1/${image}-image/$version/${image}-image-nimbra.rootfs.wic.zst"
    fi

    echo >&2 "Will download $url"

    # Verify connectivity before attempting download
    if ! verify_connectivity "https://$artifacts"; then
        exit 1
    fi

    wget "$url" --spider # Check if the file exists - uses HEAD
    wget "$url" -O - | pv -f -i 1 2>"$CONSOLE" | unzstd | dd of="${device}" bs=4M
    sync # flush the page cache before the post-install steps and reboot
fi

echo >&2 "Fixing partition table"
parted --script --fix "${device}"

echo >&2 "Configuring UEFI boot variable"
efibootmgr --delete-bootnum --label "Edge OS" || true
# This should set the boot order but does not seem to work, at least on a 414b
efibootmgr --create --disk "${device}" --part 1 --loader "\EFI\BOOT\BOOTX64.EFI" --label "Edge OS"
# setting bootnext does however
bootnum=$(efibootmgr | awk '/Edge OS/ { print substr($1, 5, 4) }')
efibootmgr --bootnext "$bootnum"
echo >&2 "installation done"
