#!/usr/bin/env bash
set -Eeuo pipefail

QNN_BASE_DIR=${QNN_BASE_DIR:-/opt/qualcomm/qairt}
QNN_STATE_DIR=${QNN_STATE_DIR:-/var/lib/dragon-q8b-qnn}
QNN_PACKAGED_CONF=${QNN_PACKAGED_CONF:-/usr/lib/dragon-q8b/qnn.conf}
QNN_CONFIG_FILE=${QNN_CONFIG_FILE:-/etc/dragon-q8b/qnn.conf}
QNN_CURRENT_LINK="$QNN_BASE_DIR/current"
qnn_work_dir=
qnn_install_tmp=

# Fallbacks used when running from git without the packaged conf. RPM builds
# replace these from config/dragon-q8b.env.
QAIRT_VERSION=2.49.40.260810
QAIRT_DOWNLOAD_URL=https://softwarecenter.qualcomm.com/api/download/software/sdks/Qualcomm_AI_Runtime_Community/All/2.49.40.260810/v2.49.40.260810.zip
QAIRT_ARCHIVE_SHA256=8b1e00d03806b026c41dbc51aec40d6b6c684fff34806409dc2283d00926629a
QAIRT_LICENSE_SHA256=ec1dccfdcba5c6e64126e84199b8362bf4999107bfa567ebe831dbb4c461692b
QAIRT_TARGET=aarch64-oe-linux-gcc11.2
QAIRT_DSP_ARCH=68

load_qnn_conf() {
    local file=$1
    [[ -r "$file" ]] || return 0
    # shellcheck disable=SC1090
    source "$file"
}

load_qnn_conf "$QNN_PACKAGED_CONF"
load_qnn_conf "$QNN_CONFIG_FILE"

usage() {
    cat <<'EOF'
Usage: dragon-q8b-qnn COMMAND [OPTIONS]

Install and validate the official Qualcomm AI Runtime (QAIRT/QNN) Community
Edition for the Radxa Dragon Q8B. The SDK is downloaded directly from
Qualcomm and is not distributed by this RPM.

Commands:
  status                    Show the configured and installed QAIRT versions
  install --accept-license  Download, verify, and install the pinned QAIRT SDK
  verify                    Run file, FastRPC, and QNN platform validation
  license                   Show the license acceptance and source information

Install options:
  --archive FILE            Use an already-downloaded official QAIRT archive
  --accept-license          Accept the Qualcomm AI Stack license bundled in it
  -h, --help                Show this help text

After the first --accept-license, later RPM pin updates install silently when
the bundled LICENSE.pdf hash is unchanged.
EOF
}

die() { echo "dragon-q8b-qnn: $*" >&2; exit 1; }

cleanup() {
    case "$qnn_work_dir" in
        /var/tmp/dragon-q8b-qairt.*) rm -rf -- "$qnn_work_dir" ;;
    esac
    case "$qnn_install_tmp" in
        "$QNN_BASE_DIR"/.install.*) rm -rf -- "$qnn_install_tmp" ;;
    esac
}
trap cleanup EXIT

release_complete() {
    local root=$1
    [[ -x "$root/bin/$QAIRT_TARGET/qnn-platform-validator" ]] &&
        [[ -r "$root/lib/$QAIRT_TARGET/libQnnHtp.so" ]] &&
        [[ -r "$root/lib/$QAIRT_TARGET/libQnnHtpV${QAIRT_DSP_ARCH}Stub.so" ]] &&
        [[ -r "$root/lib/hexagon-v${QAIRT_DSP_ARCH}/unsigned/libQnnHtpV${QAIRT_DSP_ARCH}Skel.so" ]] &&
        [[ -r "$root/LICENSE.pdf" ]]
}

runtime_complete() {
    release_complete "$QNN_CURRENT_LINK"
}

installed_version() {
    if [[ -r "$QNN_STATE_DIR/installed-version" ]]; then
        cat "$QNN_STATE_DIR/installed-version"
    fi
}

accepted_license_hash() {
    if [[ -r "$QNN_STATE_DIR/accepted-license-sha256" ]]; then
        cat "$QNN_STATE_DIR/accepted-license-sha256"
    fi
}

show_license() {
    cat <<EOF
Qualcomm AI Runtime Community Edition $QAIRT_VERSION is governed by the
Qualcomm AI Stack license included as LICENSE.pdf in the official archive.

The license permits development and application-incorporated object-code
distribution, but prohibits standalone redistribution of the SDK. This RPM
therefore installs only an open-source integration helper. The SDK is fetched
directly from Qualcomm when you explicitly pass --accept-license.

Product page:
  https://www.qualcomm.com/developer/software/qualcomm-ai-engine-direct-sdk
Catalog:
  https://softwarecenter.qualcomm.com/catalog/item/Qualcomm_AI_Runtime_Community
Official archive:
  $QAIRT_DOWNLOAD_URL
EOF
}

show_status() {
    printf 'Configured QAIRT version : %s\n' "$QAIRT_VERSION"
    printf 'Target ABI               : %s\n' "$QAIRT_TARGET"
    printf 'SC8280XP HTP architecture: v%s\n' "$QAIRT_DSP_ARCH"
    if [[ -L "$QNN_CURRENT_LINK" ]] && runtime_complete; then
        printf 'Installed                : yes\n'
        printf 'Active SDK               : %s\n' "$(readlink -f "$QNN_CURRENT_LINK")"
        printf 'Installed version        : %s\n' "$(installed_version)"
    else
        printf 'Installed                : no\n'
        printf 'Install command           : sudo dragon-q8b-qnn install --accept-license\n'
    fi
}

check_board() {
    [[ "$(uname -m)" == aarch64 ]] || die "QAIRT target runtime requires aarch64"
    if [[ "${QNN_SKIP_BOARD_CHECK:-0}" != 1 && -r /proc/device-tree/compatible ]]; then
        tr '\0' '\n' < /proc/device-tree/compatible | grep -Fxq 'radxa,dragon-q8b' ||
            die "this installer supports only the Radxa Dragon Q8B"
    fi
}

archive_license_hash() {
    unzip -p "$1" "qairt/$QAIRT_VERSION/LICENSE.pdf" | sha256sum | awk '{print $1}'
}

verify_archive() {
    local archive=$1 actual license_hash
    actual=$(sha256sum "$archive" | awk '{print $1}')
    [[ "$actual" == "$QAIRT_ARCHIVE_SHA256" ]] ||
        die "archive checksum mismatch: $actual != $QAIRT_ARCHIVE_SHA256"
    unzip -tq "$archive" >/dev/null || die "QAIRT archive is corrupt"
    license_hash=$(archive_license_hash "$archive")
    [[ "$license_hash" == "$QAIRT_LICENSE_SHA256" ]] ||
        die "QAIRT license checksum mismatch"
}

write_ld_so_conf() {
    local dest=/etc/ld.so.conf.d/dragon-q8b-qnn.conf
    install -d -m0755 "$(dirname "$dest")"
    printf '/opt/qualcomm/qairt/current/lib/%s\n' "$QAIRT_TARGET" > "$dest"
}

activate_release() {
    local version=$1 license_hash=$2 link_tmp="$QNN_BASE_DIR/.current.$$"
    ln -sfn "releases/$version" "$link_tmp"
    mv -Tf "$link_tmp" "$QNN_CURRENT_LINK"
    install -d -m0755 "$QNN_STATE_DIR"
    printf '%s\n' "$version" > "$QNN_STATE_DIR/installed-version"
    printf '%s\n' "$license_hash" > "$QNN_STATE_DIR/accepted-license-sha256"
    write_ld_so_conf
    command -v ldconfig >/dev/null 2>&1 && ldconfig
}

install_runtime() {
    local archive=$1 accept_license=$2 source_dir target_dir license_hash stored
    stored=$(accepted_license_hash)
    if [[ "$accept_license" != 1 && -z "$stored" ]]; then
        show_license >&2
        die "pass --accept-license to acknowledge the Qualcomm AI Stack license"
    fi
    [[ $EUID -eq 0 ]] || die "installation requires root privileges"
    check_board

    target_dir="$QNN_BASE_DIR/releases/$QAIRT_VERSION"
    if [[ -d "$target_dir" ]] && release_complete "$target_dir" && \
            [[ "$(installed_version)" == "$QAIRT_VERSION" ]]; then
        activate_release "$QAIRT_VERSION" "${stored:-$QAIRT_LICENSE_SHA256}"
        echo "QAIRT $QAIRT_VERSION is already installed and verified"
        return
    fi

    qnn_work_dir=$(mktemp -d /var/tmp/dragon-q8b-qairt.XXXXXX)
    if [[ -z "$archive" ]]; then
        archive="$qnn_work_dir/v${QAIRT_VERSION}.zip"
        echo "Downloading QAIRT $QAIRT_VERSION from Qualcomm Software Center"
        curl --fail --location --retry 4 --retry-all-errors \
            --proto '=https' --tlsv1.2 "$QAIRT_DOWNLOAD_URL" --output "$archive"
    else
        [[ -r "$archive" ]] || die "archive is not readable: $archive"
        archive=$(cd "$(dirname "$archive")" && pwd)/$(basename "$archive")
    fi

    echo "Verifying official QAIRT archive and bundled license"
    verify_archive "$archive"
    license_hash=$(archive_license_hash "$archive")
    if [[ "$accept_license" != 1 && "$stored" != "$license_hash" ]]; then
        show_license >&2
        die "Qualcomm AI Stack license text changed; pass --accept-license again"
    fi

    install -d -m0755 "$QNN_BASE_DIR/releases"
    if [[ -d "$target_dir" ]]; then
        release_complete "$target_dir" || die "existing QAIRT installation is incomplete: $target_dir"
        activate_release "$QAIRT_VERSION" "$license_hash"
        echo "QAIRT $QAIRT_VERSION is already installed and verified"
        return
    fi

    qnn_install_tmp=$(mktemp -d "$QNN_BASE_DIR/.install.XXXXXX")
    unzip -q "$archive" "qairt/$QAIRT_VERSION/*" -d "$qnn_install_tmp"
    source_dir="$qnn_install_tmp/qairt/$QAIRT_VERSION"
    for required in \
        "bin/$QAIRT_TARGET/qnn-platform-validator" \
        "lib/$QAIRT_TARGET/libQnnHtp.so" \
        "lib/$QAIRT_TARGET/libQnnHtpV${QAIRT_DSP_ARCH}Stub.so" \
        "lib/hexagon-v${QAIRT_DSP_ARCH}/unsigned/libQnnHtpV${QAIRT_DSP_ARCH}Skel.so" \
        LICENSE.pdf; do
        [[ -r "$source_dir/$required" ]] || die "official archive lacks $required"
    done
    mv "$source_dir" "$target_dir"
    rmdir "$qnn_install_tmp/qairt" "$qnn_install_tmp"
    qnn_install_tmp=
    activate_release "$QAIRT_VERSION" "$license_hash"
    echo "Installed QAIRT $QAIRT_VERSION from Qualcomm Software Center"
    echo "Run 'dragon-q8b-qnn verify' on the Dragon Q8B to validate NPU access"
}

verify_runtime() {
    check_board
    runtime_complete || die "QAIRT is not installed or is incomplete"
    compgen -G '/dev/fastrpc-cdsp*' >/dev/null || die "FastRPC CDSP device nodes are missing"
    ldconfig -p 2>/dev/null | grep -F libcdsprpc.so >/dev/null || \
        die "libcdsprpc.so is not registered"

    export QNN_SDK_ROOT="$QNN_CURRENT_LINK"
    export QAIRT_SDK_ROOT="$QNN_CURRENT_LINK"
    export LD_LIBRARY_PATH="$QNN_CURRENT_LINK/lib/$QAIRT_TARGET${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
    export ADSP_LIBRARY_PATH="$QNN_CURRENT_LINK/lib/hexagon-v${QAIRT_DSP_ARCH}/unsigned;/usr/share/qcom/sc8280xp/radxa/dragon-q8b/dsp/cdsp;/usr/lib/rfsa/adsp;/dsp"
    "$QNN_CURRENT_LINK/bin/$QAIRT_TARGET/qnn-platform-validator" \
        --backend dsp --coreVersion --testBackend
}

command=${1:-status}
[[ $# -eq 0 ]] || shift
archive=
accept_license=0
while [[ $# -gt 0 ]]; do
    case "$1" in
        --archive) archive=${2:?missing archive path}; shift 2 ;;
        --accept-license) accept_license=1; shift ;;
        -h|--help) usage; exit 0 ;;
        *) die "unknown option: $1" ;;
    esac
done

case "$command" in
    status) show_status ;;
    install) install_runtime "$archive" "$accept_license" ;;
    verify) verify_runtime ;;
    license) show_license ;;
    -h|--help) usage ;;
    *) usage >&2; die "unknown command: $command" ;;
esac
