#!/usr/bin/env bash
# Read-only RK3588 media hardware inventory with optional short workloads.

set -uo pipefail
export LC_ALL=C

JSON=0
QUICK=0
STRICT=0
MEDIA_DIR=""

SYS_ROOT=${RK1_SYSFS_ROOT:-/sys}
PROC_ROOT=${RK1_PROCFS_ROOT:-/proc}
DEV_ROOT=${RK1_DEV_ROOT:-/dev}
RKNN_HOME=${RKNN_HOME:-/opt/rknn/current}

declare -a CHECK_GROUP=()
declare -a CHECK_NAME=()
declare -a CHECK_STATUS=()
declare -a CHECK_MESSAGE=()

usage() {
    cat <<'EOF'
Usage: rk1-media-selftest [OPTIONS]

Options:
  --json               Emit one JSON document instead of readable text.
  --quick              Inventory only; skip GPU, codec, and NPU workloads.
  --strict             Return nonzero for warnings as well as failures.
  --media-dir PATH     Exercise RKMPP decoding with named sample files.
                       Recognized names contain h264, hevc, vp9, or av1.
  -h, --help           Show this help.

The default run is non-destructive. It executes only short inference, Vulkan,
and encoder probes when the relevant runtime and devices are present. It never
changes clocks, governors, display modes, storage, or network configuration.
EOF
}

invocation_error() {
    printf 'rk1-media-selftest: %s\n' "$*" >&2
    exit 2
}

while (($#)); do
    case "$1" in
        --json)
            JSON=1
            shift
            ;;
        --quick)
            QUICK=1
            shift
            ;;
        --strict)
            STRICT=1
            shift
            ;;
        --media-dir)
            (($# >= 2)) || invocation_error "--media-dir requires a path"
            MEDIA_DIR=$2
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            invocation_error "unknown option: $1"
            ;;
    esac
done

if [[ -n "$MEDIA_DIR" && ! -d "$MEDIA_DIR" ]]; then
    invocation_error "media directory is not a directory: $MEDIA_DIR"
fi

add_check() {
    CHECK_GROUP+=("$1")
    CHECK_NAME+=("$2")
    CHECK_STATUS+=("$3")
    CHECK_MESSAGE+=("$4")
}

read_text() {
    local path=$1
    if [[ -r "$path" ]]; then
        tr -d '\000' <"$path" 2>/dev/null || true
    fi
}

short_message() {
    local value=$1
    value=${value//$'\r'/}
    value=${value//$'\n'/'; '}
    value=${value//$'\t'/ }
    printf '%s' "${value:0:320}"
}

driver_binding_count() {
    local pattern=$1
    local directory entry base
    local count=0
    shopt -s nullglob
    for directory in "$SYS_ROOT"/bus/platform/drivers/$pattern; do
        [[ -d "$directory" ]] || continue
        for entry in "$directory"/*; do
            base=${entry##*/}
            case "$base" in
                bind|unbind|uevent|module|new_id|remove_id) continue ;;
            esac
            if [[ -L "$entry" || -d "$entry" ]]; then
                ((count++))
            fi
        done
    done
    shopt -u nullglob
    printf '%d' "$count"
}

find_sample() {
    local token=$1
    [[ -n "$MEDIA_DIR" ]] || return 1
    find "$MEDIA_DIR" -maxdepth 1 -type f \
        \( -iname "*${token}*.mkv" -o -iname "*${token}*.mp4" \
           -o -iname "*${token}*.webm" -o -iname "*${token}*.ts" \) \
        -print -quit 2>/dev/null
}

# System identity
machine=$(uname -m 2>/dev/null || printf unknown)
if [[ "$machine" == "aarch64" ]]; then
    add_check system architecture pass "aarch64 userspace"
else
    add_check system architecture fail "expected aarch64, found $machine"
fi

model=$(read_text "$PROC_ROOT/device-tree/model")
if [[ "$model" == *"Turing"*"RK1"* ]]; then
    add_check system board pass "$model"
elif [[ -n "$model" ]]; then
    add_check system board warn "device-tree model is $model"
else
    add_check system board warn "device-tree model is unavailable"
fi

# HDMI / DRM connector state
shopt -s nullglob
hdmi_connectors=("$SYS_ROOT"/class/drm/card*-HDMI-A-*)
shopt -u nullglob
if ((${#hdmi_connectors[@]} == 0)); then
    add_check hdmi connector warn "no DRM HDMI connector is registered"
else
    connected_count=0
    for connector in "${hdmi_connectors[@]}"; do
        connector_name=${connector##*/}
        state=$(read_text "$connector/status")
        case "$state" in
            connected)
                ((connected_count++))
                add_check hdmi "$connector_name" pass "connected"
                if [[ -s "$connector/edid" ]]; then
                    edid_bytes=$(wc -c <"$connector/edid" 2>/dev/null || printf 0)
                    add_check hdmi "$connector_name-edid" pass \
                        "$edid_bytes bytes of EDID data"
                else
                    add_check hdmi "$connector_name-edid" warn \
                        "connected but EDID is empty or unreadable"
                fi
                if [[ -r "$connector/modes" ]] &&
                   grep -qx '3840x2160' "$connector/modes" 2>/dev/null; then
                    add_check hdmi "$connector_name-4k" pass \
                        "EDID advertises a 3840x2160 mode"
                else
                    add_check hdmi "$connector_name-4k" warn \
                        "3840x2160 is not present in the connector mode list"
                fi
                ;;
            disconnected)
                add_check hdmi "$connector_name" warn \
                    "disconnected (expected when no display is attached)"
                ;;
            *)
                add_check hdmi "$connector_name" warn \
                    "connector state is ${state:-unknown}"
                ;;
        esac
    done
fi

hdmi_driver_count=$(driver_binding_count 'dwhdmi*')
hdmi_phy_count=$(driver_binding_count '*hdptx*hdmi*')
if ((hdmi_driver_count > 0 && hdmi_phy_count > 0)); then
    add_check hdmi drivers pass \
        "$hdmi_driver_count controller and $hdmi_phy_count HDMI PHY binding(s)"
else
    add_check hdmi drivers warn \
        "HDMI controller bindings=$hdmi_driver_count, PHY bindings=$hdmi_phy_count"
fi

# GPU / Vulkan
gpu_driver=""
for candidate in panthor panfrost mali; do
    if (($(driver_binding_count "$candidate") > 0)); then
        gpu_driver=$candidate
        break
    fi
done
case "$gpu_driver" in
    panthor)
        add_check gpu kernel-driver pass "Panthor is bound to the GPU"
        ;;
    panfrost)
        add_check gpu kernel-driver warn "Panfrost is bound; the image expects Panthor"
        ;;
    mali)
        add_check gpu kernel-driver warn "proprietary Mali kernel driver is bound"
        ;;
    *)
        add_check gpu kernel-driver warn "no bound Mali GPU driver was found"
        ;;
esac

shopt -s nullglob
gpu_devfreq=("$SYS_ROOT"/class/devfreq/*.gpu)
render_nodes=("$DEV_ROOT"/dri/renderD*)
shopt -u nullglob
if ((${#gpu_devfreq[@]} > 0)); then
    add_check gpu devfreq pass "${#gpu_devfreq[@]} GPU devfreq device(s)"
else
    add_check gpu devfreq warn "GPU devfreq node is missing"
fi
if ((${#render_nodes[@]} > 0)); then
    add_check gpu render-node pass "${#render_nodes[@]} DRM render node(s)"
else
    add_check gpu render-node warn "no accessible DRM render node"
fi

if ((QUICK)); then
    add_check gpu vulkan skip "active Vulkan probe disabled by --quick"
elif command -v vulkaninfo >/dev/null 2>&1; then
    vulkan_output=$(timeout 20s vulkaninfo --summary 2>&1)
    vulkan_rc=$?
    if ((vulkan_rc == 0)); then
        gpu_name=$(printf '%s\n' "$vulkan_output" |
            sed -n 's/^[[:space:]]*deviceName[[:space:]]*=[[:space:]]*//p' |
            head -1)
        add_check gpu vulkan pass "${gpu_name:-vulkaninfo completed}"
    else
        add_check gpu vulkan fail \
            "vulkaninfo failed: $(short_message "$vulkan_output")"
    fi
else
    add_check gpu vulkan skip "vulkaninfo is not installed"
fi

# VPU bindings and FFmpeg integration
decoder_bindings=$(driver_binding_count 'mpp_rkvdec*')
encoder_bindings=$(driver_binding_count 'mpp_rkvenc*')
av1_bindings=$(driver_binding_count '*av1*')
if ((decoder_bindings > 0)); then
    add_check vpu decoder-driver pass "$decoder_bindings RKMPP decoder binding(s)"
else
    add_check vpu decoder-driver warn "RKMPP decoder bindings are missing"
fi
if ((encoder_bindings > 0)); then
    add_check vpu encoder-driver pass "$encoder_bindings RKMPP encoder binding(s)"
else
    add_check vpu encoder-driver warn "RKMPP encoder bindings are missing"
fi
if ((av1_bindings > 0)); then
    add_check vpu av1-driver pass "$av1_bindings AV1 decoder binding(s)"
else
    add_check vpu av1-driver warn "a separate AV1 driver binding was not found"
fi

FFMPEG=""
for candidate in /opt/rkmedia/bin/ffmpeg-rk /usr/local/bin/ffmpeg-rk; do
    if [[ -x "$candidate" ]]; then
        FFMPEG=$candidate
        break
    fi
done
if [[ -z "$FFMPEG" ]] && command -v ffmpeg-rk >/dev/null 2>&1; then
    FFMPEG=$(command -v ffmpeg-rk)
fi

if [[ -z "$FFMPEG" ]]; then
    add_check vpu ffmpeg-rkmpp skip "ffmpeg-rk is not installed"
    add_check rga ffmpeg-filter skip "ffmpeg-rk is not installed"
else
    hwaccels=$($FFMPEG -hide_banner -hwaccels 2>&1)
    if grep -qE '(^|[[:space:]])rkmpp($|[[:space:]])' <<<"$hwaccels"; then
        add_check vpu ffmpeg-rkmpp pass "$FFMPEG advertises the rkmpp hwaccel"
    else
        add_check vpu ffmpeg-rkmpp fail "$FFMPEG does not advertise rkmpp"
    fi

    encoders=$($FFMPEG -hide_banner -encoders 2>&1)
    for codec in h264 hevc mjpeg; do
        if grep -q "${codec}_rkmpp" <<<"$encoders"; then
            add_check vpu "encode-$codec-capability" pass \
                "${codec}_rkmpp is registered"
        else
            add_check vpu "encode-$codec-capability" fail \
                "${codec}_rkmpp is not registered"
        fi
    done

    filters=$($FFMPEG -hide_banner -filters 2>&1)
    if grep -q 'scale_rkrga' <<<"$filters"; then
        add_check rga ffmpeg-filter pass "scale_rkrga is registered"
    else
        add_check rga ffmpeg-filter fail "scale_rkrga is not registered"
    fi

    if ((QUICK)); then
        add_check vpu encode-workload skip "active encoder probe disabled by --quick"
    else
        for codec in h264 hevc mjpeg; do
            encode_output=$(timeout 30s "$FFMPEG" -nostdin -v error \
                -f lavfi -i 'color=size=320x240:rate=30:duration=0.2' \
                -frames:v 3 -c:v "${codec}_rkmpp" -f null - 2>&1)
            encode_rc=$?
            if ((encode_rc == 0)); then
                add_check vpu "encode-$codec-workload" pass \
                    "three frames submitted successfully"
            else
                add_check vpu "encode-$codec-workload" fail \
                    "hardware encode failed: $(short_message "$encode_output")"
            fi
        done
    fi

    if [[ -n "$MEDIA_DIR" && $QUICK -eq 0 ]]; then
        first_decode_sample=""
        for codec in h264 hevc vp9 av1; do
            sample=$(find_sample "$codec" || true)
            if [[ -z "$sample" ]]; then
                add_check vpu "decode-$codec-workload" skip \
                    "no sample containing '$codec' in $MEDIA_DIR"
                continue
            fi
            [[ -n "$first_decode_sample" ]] || first_decode_sample=$sample
            decode_output=$(timeout 180s "$FFMPEG" -nostdin -v error \
                -hwaccel rkmpp -i "$sample" -map 0:v:0 -f null - 2>&1)
            decode_rc=$?
            if ((decode_rc == 0)); then
                add_check vpu "decode-$codec-workload" pass \
                    "decoded ${sample##*/} through RKMPP"
            else
                add_check vpu "decode-$codec-workload" fail \
                    "decode failed: $(short_message "$decode_output")"
            fi
        done

        if [[ -n "$first_decode_sample" ]] && grep -q 'scale_rkrga' <<<"$filters"; then
            rga_output=$(timeout 60s "$FFMPEG" -nostdin -v error \
                -hwaccel rkmpp -hwaccel_output_format drm_prime \
                -i "$first_decode_sample" -map 0:v:0 -frames:v 30 \
                -vf 'scale_rkrga=w=1280:h=720:format=nv12' -f null - 2>&1)
            rga_rc=$?
            if ((rga_rc == 0)); then
                add_check rga workload pass "RKMPP-to-RGA scaling completed"
            else
                add_check rga workload fail \
                    "RGA workload failed: $(short_message "$rga_output")"
            fi
        else
            add_check rga workload skip "no decode sample is available for RGA"
        fi
    elif [[ -z "$MEDIA_DIR" ]]; then
        add_check vpu decode-workload skip "use --media-dir to test hardware decoding"
        add_check rga workload skip "use --media-dir to test zero-copy RGA scaling"
    fi
fi

rga2_bindings=$(driver_binding_count 'rga2')
rga3_bindings=$(driver_binding_count 'rga3')
if ((rga2_bindings + rga3_bindings > 0)); then
    add_check rga kernel-driver pass \
        "RGA2 bindings=$rga2_bindings, RGA3 bindings=$rga3_bindings"
else
    add_check rga kernel-driver warn "no RGA2/RGA3 binding was found"
fi

# NPU runtime and per-core inference
npu_bindings=$(driver_binding_count 'rknpu')
((npu_bindings += $(driver_binding_count 'RKNPU')))
shopt -s nullglob
npu_devfreq=("$SYS_ROOT"/class/devfreq/*.npu)
shopt -u nullglob
if ((npu_bindings > 0 || ${#npu_devfreq[@]} > 0)); then
    add_check npu kernel-driver pass \
        "RKNPU bindings=$npu_bindings, devfreq nodes=${#npu_devfreq[@]}"
    npu_present=1
else
    add_check npu kernel-driver warn "RKNPU driver/device was not found"
    npu_present=0
fi

runtime_so="$RKNN_HOME/lib/librknnrt.so"
model_path="$RKNN_HOME/share/models/rk3588/mobilenet_v1.rknn"
npu_test="$RKNN_HOME/bin/rknn-inference-test"
if [[ -r "$runtime_so" ]]; then
    runtime_hash=$(sha256sum "$runtime_so" 2>/dev/null)
    runtime_hash=${runtime_hash%% *}
    if [[ "$runtime_hash" == \
          "d31fc19c85b85f6091b2bd0f6af9d962d5264a4e410bfb536402ec92bac738e8" ]]; then
        add_check npu runtime pass "RKNN Runtime 2.3.2 checksum matches"
    else
        add_check npu runtime fail "unexpected librknnrt.so checksum $runtime_hash"
    fi
else
    add_check npu runtime warn "RKNN runtime is not installed at $runtime_so"
fi

if ((QUICK)); then
    add_check npu inference skip "active inference disabled by --quick"
elif ((npu_present == 0)); then
    add_check npu inference skip "NPU driver is absent; inference was not attempted"
elif [[ ! -x "$npu_test" || ! -r "$model_path" ]]; then
    add_check npu inference skip "test binary or pinned RK3588 model is unavailable"
else
    for core in 0 1 2 all; do
        npu_output=$(timeout 45s "$npu_test" --model "$model_path" \
            --core "$core" --iterations 1 2>&1)
        npu_rc=$?
        if ((npu_rc == 0)) && grep -q 'RKNN_RESULT status=pass' <<<"$npu_output"; then
            add_check npu "core-$core" pass "$(short_message "$npu_output")"
        else
            add_check npu "core-$core" fail \
                "inference failed: $(short_message "$npu_output")"
        fi
    done
fi

# HDMI ALSA devices
asound_cards=$(read_text "$PROC_ROOT/asound/cards")
if [[ -n "$asound_cards" ]] && grep -qiE 'hdmi|rockchiphdmi' <<<"$asound_cards"; then
    add_check audio hdmi-card pass "an HDMI ALSA card is registered"
elif [[ -n "$asound_cards" ]]; then
    add_check audio hdmi-card warn "ALSA cards exist, but none is labeled HDMI"
else
    add_check audio hdmi-card warn "ALSA card inventory is unavailable"
fi
if command -v aplay >/dev/null 2>&1; then
    aplay_output=$(aplay -l 2>&1)
    if grep -qi 'hdmi' <<<"$aplay_output"; then
        add_check audio pcm-device pass "aplay lists an HDMI PCM device"
    else
        add_check audio pcm-device warn \
            "aplay did not list an HDMI PCM device: $(short_message "$aplay_output")"
    fi
else
    add_check audio pcm-device skip "aplay is not installed"
fi

# Link state only: no pings, DNS requests, or network mutations.
shopt -s nullglob
interfaces=("$SYS_ROOT"/class/net/*)
shopt -u nullglob
interface_names=()
up_names=()
for interface in "${interfaces[@]}"; do
    name=${interface##*/}
    [[ "$name" == lo ]] && continue
    interface_names+=("$name")
    state=$(read_text "$interface/operstate")
    carrier=$(read_text "$interface/carrier")
    if [[ "$state" == up || "$carrier" == 1 ]]; then
        up_names+=("$name")
    fi
done
if ((${#interface_names[@]} == 0)); then
    add_check network interfaces warn "no non-loopback interface was found"
elif ((${#up_names[@]} > 0)); then
    add_check network link pass "up: ${up_names[*]}"
else
    add_check network link warn \
        "interfaces present but down: ${interface_names[*]}"
fi

# Root filesystem, eMMC, and NVMe presence. No serial numbers are read.
if command -v findmnt >/dev/null 2>&1; then
    root_source=$(findmnt -n -o SOURCE / 2>/dev/null || true)
    root_fstype=$(findmnt -n -o FSTYPE / 2>/dev/null || true)
    if [[ -n "$root_source" ]]; then
        add_check storage rootfs pass "$root_source ($root_fstype)"
    else
        add_check storage rootfs warn "root filesystem source is unavailable"
    fi
else
    add_check storage rootfs skip "findmnt is not installed"
fi

shopt -s nullglob
emmc_devices=("$SYS_ROOT"/block/mmcblk*)
nvme_devices=("$SYS_ROOT"/block/nvme*n1)
shopt -u nullglob
if ((${#emmc_devices[@]} > 0)); then
    add_check storage emmc pass "${#emmc_devices[@]} MMC block device(s)"
else
    add_check storage emmc warn "no MMC/eMMC block device was found"
fi
if ((${#nvme_devices[@]} > 0)); then
    add_check storage nvme pass "${#nvme_devices[@]} NVMe namespace(s)"
else
    add_check storage nvme warn "no NVMe namespace was found"
fi

pass_count=0
warn_count=0
fail_count=0
skip_count=0
for status in "${CHECK_STATUS[@]}"; do
    case "$status" in
        pass) ((pass_count++)) ;;
        warn) ((warn_count++)) ;;
        fail) ((fail_count++)) ;;
        skip) ((skip_count++)) ;;
    esac
done

if ((fail_count > 0)); then
    overall=fail
elif ((warn_count > 0)); then
    overall=warn
else
    overall=pass
fi

json_escape() {
    local value=$1
    value=${value//\\/\\\\}
    value=${value//\"/\\\"}
    value=${value//$'\n'/\\n}
    value=${value//$'\r'/\\r}
    value=${value//$'\t'/\\t}
    printf '%s' "$value"
}

if ((JSON)); then
    printf '{"schema_version":1,"overall":"%s","quick":%s,' \
        "$overall" "$([[ $QUICK -eq 1 ]] && printf true || printf false)"
    printf '"summary":{"pass":%d,"warn":%d,"fail":%d,"skip":%d},' \
        "$pass_count" "$warn_count" "$fail_count" "$skip_count"
    printf '"checks":['
    for ((i = 0; i < ${#CHECK_NAME[@]}; ++i)); do
        ((i == 0)) || printf ','
        printf '{"group":"%s","name":"%s","status":"%s","message":"%s"}' \
            "$(json_escape "${CHECK_GROUP[i]}")" \
            "$(json_escape "${CHECK_NAME[i]}")" \
            "${CHECK_STATUS[i]}" \
            "$(json_escape "${CHECK_MESSAGE[i]}")"
    done
    printf ']}\n'
else
    printf '%-9s %-24s %-6s %s\n' GROUP CHECK STATUS DETAIL
    printf '%-9s %-24s %-6s %s\n' '---------' '------------------------' \
        '------' '------'
    for ((i = 0; i < ${#CHECK_NAME[@]}; ++i)); do
        printf '%-9s %-24s %-6s %s\n' \
            "${CHECK_GROUP[i]}" "${CHECK_NAME[i]}" \
            "${CHECK_STATUS[i]^^}" "${CHECK_MESSAGE[i]}"
    done
    printf '\nOverall: %s (pass=%d warn=%d fail=%d skip=%d)\n' \
        "${overall^^}" "$pass_count" "$warn_count" "$fail_count" "$skip_count"
fi

if ((fail_count > 0 || (STRICT && warn_count > 0))); then
    exit 1
fi
exit 0