#!/bin/sh
set -eu

ARTIFACT=
while [ "$#" -gt 0 ]; do
    case "$1" in
        --artifact) ARTIFACT=$2; shift 2 ;;
        *) echo "unknown argument: $1" >&2; exit 2 ;;
    esac
done
if [ "$(id -u)" -ne 0 ] || [ "$(uname -m)" != "aarch64" ]; then
    echo "installer must run as root on the WalnutPi AArch64 host" >&2
    exit 1
fi
if [ -z "$ARTIFACT" ]; then
    echo "usage: install_axp313a_kernel.sh --artifact DIR" >&2
    exit 2
fi
ARTIFACT=$(CDPATH= cd -- "$ARTIFACT" && pwd)
SOURCE_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
RELEASE=6.1.31-matrix-axp313a1
MODULE_TARGET=/lib/modules/$RELEASE
IMAGE_TARGET=/boot/Image-matrix-axp313a1
DTB_TARGET=/boot/sun50i-h616-walnutpi-1b-matrix-axp313a1.dtb
EMMC_DTB_TARGET=/boot/sun50i-h616-walnutpi-1b-emmc-matrix-axp313a1.dtb
BOOT_CMD=/boot/boot.cmd
BOOT_SCR=/boot/boot.scr
ORIGINAL_CMD=/boot/boot.cmd.matrix-original
ORIGINAL_SCR=/boot/boot.scr.matrix-original
COMMITTED=0
BOOT_REPLACED=0
SYSTEM_FILES_INSTALLED=0

cleanup_failed_install() {
    code=$?
    trap - EXIT HUP INT TERM
    if [ "$code" -ne 0 ] && [ "$COMMITTED" -eq 0 ]; then
        if [ "$BOOT_REPLACED" -eq 1 ] && [ -f "$ORIGINAL_CMD" ] && [ -f "$ORIGINAL_SCR" ]; then
            install -m 0644 "$ORIGINAL_CMD" "$BOOT_CMD"
            install -m 0644 "$ORIGINAL_SCR" "$BOOT_SCR"
        fi
        rm -f -- "$IMAGE_TARGET" "$DTB_TARGET" "$EMMC_DTB_TARGET" \
            "/boot/System.map-$RELEASE" "/boot/config-$RELEASE" /boot/matrix-kernel-good
        case "$MODULE_TARGET" in /lib/modules/6.1.31-matrix-axp313a1) rm -rf -- "$MODULE_TARGET" ;; esac
        if [ "$SYSTEM_FILES_INSTALLED" -eq 1 ]; then
            rm -f -- /etc/systemd/system/matrix-axp313a-health.service \
                /etc/systemd/system/matrix-axp313a-rollback.service \
                /opt/matrix-screen-controller-system/axp313a_kernel_health.py
        fi
        systemctl daemon-reload >/dev/null 2>&1 || true
        sync
    fi
    exit "$code"
}
trap cleanup_failed_install EXIT HUP INT TERM

if [ ! -f "$ARTIFACT/SHA256SUMS" ]; then
    echo "candidate artifact manifest is missing" >&2
    exit 1
fi
(cd "$ARTIFACT" && sha256sum -c SHA256SUMS)
for path in \
    "$ARTIFACT/boot/Image-matrix-axp313a1" \
    "$ARTIFACT/boot/sun50i-h616-walnutpi-1b-matrix-axp313a1.dtb" \
    "$ARTIFACT/boot/sun50i-h616-walnutpi-1b-emmc-matrix-axp313a1.dtb" \
    "$ARTIFACT/lib/modules/$RELEASE"
do
    if [ ! -e "$path" ]; then echo "candidate artifact is incomplete: $path" >&2; exit 1; fi
done
if [ -e "$MODULE_TARGET" ] || [ -e "$IMAGE_TARGET" ] || [ -e "$DTB_TARGET" ] || [ -e "$EMMC_DTB_TARGET" ]; then
    echo "candidate kernel paths already exist; refusing to overwrite" >&2
    exit 1
fi
if [ ! -f "$BOOT_CMD" ] || [ ! -f "$BOOT_SCR" ]; then
    echo "original WalnutPi boot files are missing" >&2
    exit 1
fi
command -v mkimage >/dev/null
command -v fdtget >/dev/null
command -v strings >/dev/null

# Refuse before the first candidate boot write unless the complete candidate,
# rollback set, temporary scripts and fixed safety reserve all fit.
PYTHONPATH="$SOURCE_ROOT" python3 "$SOURCE_ROOT/scripts/boot_space.py" \
    --boot-root /boot --artifact-root "$ARTIFACT"

if [ "$(cat "$ARTIFACT/KERNEL_RELEASE" 2>/dev/null || true)" != "$RELEASE" ]; then
    echo "candidate kernel release marker is wrong" >&2
    exit 1
fi
if [ "$(cat "$ARTIFACT/SOURCE_COMMIT" 2>/dev/null || true)" != \
    "30ff3fd5cf45417622b447a12e7a947402ffe34d" ]; then
    echo "candidate vendor source commit marker is wrong" >&2
    exit 1
fi
if ! strings "$ARTIFACT/boot/Image-matrix-axp313a1" | \
    grep -Fq "Linux version $RELEASE"; then
    echo "candidate Image does not contain the expected release" >&2
    exit 1
fi

verify_dtb() {
    dtb=$1
    pmic=/soc/i2c@7081400/pmic@36
    regulators=$pmic/regulators
    if [ "$(fdtget -t s "$dtb" "$pmic" compatible)" != "x-powers,axp313a" ]; then
        echo "candidate DTB does not use the exact AXP313A compatible: $dtb" >&2
        return 1
    fi
    for specification in \
        "aldo1 1800000 1800000 vcc1v8" \
        "dldo1 3300000 3300000 vcc3v3" \
        "dcdc1 810000 1160000 vdd-gpu-sys" \
        "dcdc2 810000 1080000 vdd-cpu" \
        "dcdc3 1100000 1100000 vdd-dram"
    do
        set -- $specification
        node=$regulators/$1
        [ "$(fdtget -t i "$dtb" "$node" regulator-min-microvolt)" = "$2" ]
        [ "$(fdtget -t i "$dtb" "$node" regulator-max-microvolt)" = "$3" ]
        [ "$(fdtget -t s "$dtb" "$node" regulator-name)" = "$4" ]
        fdtget "$dtb" "$node" regulator-always-on >/dev/null
    done
    cpu_supply=$(fdtget -t x "$dtb" /cpus/cpu@0 cpu-supply)
    dcdc2_phandle=$(fdtget -t x "$dtb" "$regulators/dcdc2" phandle)
    if [ "$cpu_supply" != "$dcdc2_phandle" ]; then
        echo "candidate DTB CPU0 supply is not DCDC2: $dtb" >&2
        return 1
    fi
}

verify_dtb "$ARTIFACT/boot/sun50i-h616-walnutpi-1b-matrix-axp313a1.dtb"
verify_dtb "$ARTIFACT/boot/sun50i-h616-walnutpi-1b-emmc-matrix-axp313a1.dtb"

install -m 0644 "$ARTIFACT/boot/Image-matrix-axp313a1" "$IMAGE_TARGET"
install -m 0644 "$ARTIFACT/boot/sun50i-h616-walnutpi-1b-matrix-axp313a1.dtb" "$DTB_TARGET"
install -m 0644 "$ARTIFACT/boot/sun50i-h616-walnutpi-1b-emmc-matrix-axp313a1.dtb" "$EMMC_DTB_TARGET"
install -m 0644 "$ARTIFACT/boot/System.map-$RELEASE" "/boot/System.map-$RELEASE"
install -m 0644 "$ARTIFACT/boot/config-$RELEASE" "/boot/config-$RELEASE"
cp -a -- "$ARTIFACT/lib/modules/$RELEASE" "$MODULE_TARGET"
depmod "$RELEASE"

if [ ! -f "$ORIGINAL_CMD" ]; then install -m 0644 "$BOOT_CMD" "$ORIGINAL_CMD"; fi
if [ ! -f "$ORIGINAL_SCR" ]; then install -m 0644 "$BOOT_SCR" "$ORIGINAL_SCR"; fi
sha256sum "$ORIGINAL_CMD" "$ORIGINAL_SCR" > /boot/matrix-original.SHA256SUMS

install -d -m 0755 /opt/matrix-screen-controller-system
install -m 0755 "$SOURCE_ROOT/scripts/axp313a_kernel_health.py" \
    /opt/matrix-screen-controller-system/axp313a_kernel_health.py
install -m 0644 "$SOURCE_ROOT/systemd/matrix-axp313a-health.service" \
    /etc/systemd/system/matrix-axp313a-health.service
install -m 0644 "$SOURCE_ROOT/systemd/matrix-axp313a-rollback.service" \
    /etc/systemd/system/matrix-axp313a-rollback.service
SYSTEM_FILES_INSTALLED=1

cmd_tmp=$(mktemp /boot/.boot.cmd.matrix.XXXXXX)
scr_tmp=$(mktemp /boot/.boot.scr.matrix.XXXXXX)
python3 "$SOURCE_ROOT/scripts/render_dual_kernel_boot.py" "$ORIGINAL_CMD" "$cmd_tmp"
mkimage -C none -A arm -T script -d "$cmd_tmp" "$scr_tmp" >/dev/null
install -m 0644 "$cmd_tmp" "$BOOT_CMD"
install -m 0644 "$scr_tmp" "$BOOT_SCR"
BOOT_REPLACED=1
rm -f -- "$cmd_tmp" "$scr_tmp"
printf '%s\n' "$RELEASE" > /boot/matrix-kernel-good
sync
systemctl daemon-reload
systemctl enable matrix-axp313a-health.service >/dev/null
COMMITTED=1
trap - EXIT HUP INT TERM
echo "AXP313A candidate installed; reboot is required"