#!/usr/bin/env bash
# Install the pinned RKNN C runtime, test model, and RK1 diagnostic tools.

set -Eeuo pipefail
umask 022

SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)
# shellcheck source=rknn-version.env
source "$SCRIPT_DIR/rknn-version.env"

ROOTFS="/"
SOURCE_DIR=""
CC_BIN=${CC:-cc}
KEEP_WORK=0
WORK_DIR=""

usage() {
    cat <<'EOF'
Usage: install-rknn-runtime.sh [OPTIONS]

Options:
  --rootfs PATH       Install below PATH (default: /).
  --source-dir PATH   Use an existing rknn-toolkit2 v2.3.2 checkout.
  --cc COMMAND        C compiler for the target (default: $CC or cc).
  --keep-work         Preserve a downloaded checkout for inspection.
  -h, --help          Show this help.

Without --source-dir, the script fetches only the required paths from the
immutable upstream commit. Network access and git are then required.
EOF
}

die() {
    printf 'install-rknn-runtime: error: %s\n' "$*" >&2
    exit 1
}

note() {
    printf 'install-rknn-runtime: %s\n' "$*" >&2
}

while (($#)); do
    case "$1" in
        --rootfs)
            (($# >= 2)) || die "--rootfs requires a path"
            ROOTFS=$2
            shift 2
            ;;
        --source-dir)
            (($# >= 2)) || die "--source-dir requires a path"
            SOURCE_DIR=$2
            shift 2
            ;;
        --cc)
            (($# >= 2)) || die "--cc requires a compiler command"
            CC_BIN=$2
            shift 2
            ;;
        --keep-work)
            KEEP_WORK=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            die "unknown option: $1"
            ;;
    esac
done

[[ -d "$ROOTFS" ]] || die "rootfs is not a directory: $ROOTFS"
ROOTFS=$(readlink -f -- "$ROOTFS")
[[ -n "$ROOTFS" ]] || die "could not resolve rootfs"

if [[ "$ROOTFS" == "/" && ${EUID:-$(id -u)} -ne 0 ]]; then
    die "installing into / requires root"
fi

for tool in install ln readlink sha256sum "$CC_BIN"; do
    command -v "$tool" >/dev/null 2>&1 || die "required command not found: $tool"
done

cleanup() {
    if [[ -n "$WORK_DIR" && -d "$WORK_DIR" && $KEEP_WORK -eq 0 ]]; then
        rm -rf -- "$WORK_DIR"
    elif [[ -n "$WORK_DIR" && -d "$WORK_DIR" ]]; then
        note "preserved work directory: $WORK_DIR"
    fi
}
trap cleanup EXIT

verify_sha256() {
    local expected=$1
    local path=$2
    local actual
    [[ -f "$path" ]] || die "required upstream asset is missing: $path"
    actual=$(sha256sum -- "$path")
    actual=${actual%% *}
    [[ "$actual" == "$expected" ]] ||
        die "checksum mismatch for $path (expected $expected, got $actual)"
}

if [[ -z "$SOURCE_DIR" ]]; then
    command -v git >/dev/null 2>&1 || die "git is required without --source-dir"
    WORK_DIR=$(mktemp -d "${TMPDIR:-/tmp}/rknn-v232.XXXXXXXX")
    SOURCE_DIR="$WORK_DIR/source"
    note "fetching RKNN Toolkit2 commit $RKNN_COMMIT"
    git init -q "$SOURCE_DIR"
    git -C "$SOURCE_DIR" remote add origin "$RKNN_REPOSITORY"
    git -C "$SOURCE_DIR" config remote.origin.promisor true
    git -C "$SOURCE_DIR" config remote.origin.partialclonefilter blob:none
    git -C "$SOURCE_DIR" sparse-checkout init --cone
    git -C "$SOURCE_DIR" sparse-checkout set \
        rknpu2/runtime/Linux/librknn_api \
        rknpu2/examples/rknn_api_demo/model
    git -C "$SOURCE_DIR" fetch -q --depth 1 --filter=blob:none origin \
        "refs/tags/$RKNN_TAG"
    git -C "$SOURCE_DIR" checkout -q --detach FETCH_HEAD
fi

SOURCE_DIR=$(readlink -f -- "$SOURCE_DIR")
[[ -d "$SOURCE_DIR" ]] || die "source directory is not a directory"
if [[ -d "$SOURCE_DIR/.git" ]]; then
    source_commit=$(git -C "$SOURCE_DIR" rev-parse HEAD)
    [[ "$source_commit" == "$RKNN_COMMIT" ]] ||
        die "source checkout is $source_commit, expected $RKNN_COMMIT"
fi

HEADER="$SOURCE_DIR/rknpu2/runtime/Linux/librknn_api/include/rknn_api.h"
RUNTIME_SO="$SOURCE_DIR/rknpu2/runtime/Linux/librknn_api/aarch64/librknnrt.so"
MODEL="$SOURCE_DIR/rknpu2/examples/rknn_api_demo/model/RK3588/mobilenet_v1.rknn"
DEMO_IMAGE="$SOURCE_DIR/rknpu2/examples/rknn_api_demo/model/dog_224x224.jpg"
LICENSE_FILE="$SOURCE_DIR/LICENSE"

verify_sha256 "$RKNN_HEADER_SHA256" "$HEADER"
verify_sha256 "$RKNN_RUNTIME_AARCH64_SHA256" "$RUNTIME_SO"
verify_sha256 "$RKNN_MOBILENET_RK3588_SHA256" "$MODEL"
verify_sha256 "$RKNN_DEMO_IMAGE_SHA256" "$DEMO_IMAGE"
verify_sha256 "$RKNN_LICENSE_SHA256" "$LICENSE_FILE"

if [[ -z "$WORK_DIR" ]]; then
    WORK_DIR=$(mktemp -d "${TMPDIR:-/tmp}/rknn-v232-build.XXXXXXXX")
fi
TEST_BINARY="$WORK_DIR/rknn-inference-test"

note "building the ARM64 RKNN smoke test with $CC_BIN"
"$CC_BIN" \
    -std=c11 -D_POSIX_C_SOURCE=200809L -O2 -Wall -Wextra -Wpedantic \
    -fstack-protector-strong -D_FORTIFY_SOURCE=2 \
    -I"$(dirname -- "$HEADER")" \
    "$SCRIPT_DIR/rknn-inference-test.c" \
    -L"$(dirname -- "$RUNTIME_SO")" -Wl,--as-needed \
    -Wl,-z,relro,-z,now -Wl,-rpath,/opt/rknn/current/lib \
    -lrknnrt -lm -o "$TEST_BINARY"

VERSION_ROOT="$ROOTFS/opt/rknn/$RKNN_VERSION"
install -d -m 0755 \
    "$VERSION_ROOT/bin" \
    "$VERSION_ROOT/include" \
    "$VERSION_ROOT/lib" \
    "$VERSION_ROOT/share/models/rk3588" \
    "$VERSION_ROOT/share/demo" \
    "$VERSION_ROOT/share/licenses" \
    "$ROOTFS/etc/ld.so.conf.d" \
    "$ROOTFS/usr/local/bin" \
    "$ROOTFS/usr/local/sbin"

install -m 0755 "$TEST_BINARY" "$VERSION_ROOT/bin/rknn-inference-test"
install -m 0644 "$HEADER" "$VERSION_ROOT/include/rknn_api.h"
install -m 0644 "$RUNTIME_SO" "$VERSION_ROOT/lib/librknnrt.so"
ln -sfn librknnrt.so "$VERSION_ROOT/lib/librknn_api.so"
install -m 0644 "$MODEL" \
    "$VERSION_ROOT/share/models/rk3588/mobilenet_v1.rknn"
install -m 0644 "$DEMO_IMAGE" "$VERSION_ROOT/share/demo/dog_224x224.jpg"
install -m 0644 "$LICENSE_FILE" \
    "$VERSION_ROOT/share/licenses/rknn-toolkit2-LICENSE"

cat >"$VERSION_ROOT/manifest.env" <<EOF
RKNN_VERSION=$RKNN_VERSION
RKNN_TAG=$RKNN_TAG
RKNN_COMMIT=$RKNN_COMMIT
RKNN_REPOSITORY=$RKNN_REPOSITORY
RKNN_RUNTIME_AARCH64_SHA256=$RKNN_RUNTIME_AARCH64_SHA256
RKNN_MOBILENET_RK3588_SHA256=$RKNN_MOBILENET_RK3588_SHA256
RKNN_DEMO_IMAGE_SHA256=$RKNN_DEMO_IMAGE_SHA256
EOF
chmod 0644 "$VERSION_ROOT/manifest.env"

ln -sfn "$RKNN_VERSION" "$ROOTFS/opt/rknn/current"
ln -sfn /opt/rknn/current/bin/rknn-inference-test \
    "$ROOTFS/usr/local/bin/rknn-inference-test"
printf '%s\n' '/opt/rknn/current/lib' >"$ROOTFS/etc/ld.so.conf.d/rknn.conf"
chmod 0644 "$ROOTFS/etc/ld.so.conf.d/rknn.conf"

install -m 0755 "$SCRIPT_DIR/rk1-media-selftest" \
    "$ROOTFS/usr/local/bin/rk1-media-selftest"
install -m 0755 "$SCRIPT_DIR/rk1-media-diagnostics" \
    "$ROOTFS/usr/local/sbin/rk1-media-diagnostics"

if command -v ldconfig >/dev/null 2>&1 &&
   [[ -e "$ROOTFS/etc/ld.so.conf" && -d "$ROOTFS/lib" ]]; then
    if ! ldconfig -r "$ROOTFS"; then
        note "warning: ldconfig failed; the smoke test still has an embedded rpath"
    fi
fi

note "installed RKNN Runtime $RKNN_VERSION below $ROOTFS/opt/rknn/current"