#!/usr/bin/bash
# waydroid-nvidia-setup — provision an initialized waydroid install for the
# NVIDIA/Venus stack. Run as root AFTER `waydroid init`. Safe to re-run.
#
#   waydroid-nvidia-setup [--refresh <hz>]
#
# What it does:
#   1. copies the guest driver stack from /usr/lib/waydroid-nvidia/guest into
#      /var/lib/waydroid/nv/guest while preserving Android's lib/lib64 paths
#      (the patched waydroid config generator bind-mounts them into the
#      container),
#   2. writes the stack's properties into waydroid.cfg [properties], sets
#      suspend_action=none (values merge into base.prop on upgrade), and points
#      drm_device at the NVIDIA render node (upstream gpu.py blacklists nvidia
#      during auto-detection, so NVIDIA-only hosts would otherwise get no
#      /dev/dri node in the container and SurfaceFlinger crash-loops),
#   3. regenerates the container config via `waydroid upgrade -o`.
#
# --refresh <hz>: match your monitor's refresh rate (e.g. 144/240/500). Above
# 240 Hz this also enables the vsync-snap window that needs the patched
# surfaceflinger shipped with this package.
set -euo pipefail

die() { echo "waydroid-nvidia-setup: $*" >&2; exit 1; }
[ "$(id -u)" = 0 ] || die "must run as root"

REFRESH=""
while [ $# -gt 0 ]; do
    case "$1" in
        --refresh) REFRESH="${2:?--refresh needs a value}"; shift 2 ;;
        *) die "unknown argument: $1" ;;
    esac
done

CFG=/var/lib/waydroid/waydroid.cfg
GUEST=/usr/lib/waydroid-nvidia/guest
NV=/var/lib/waydroid/nv
NV_GUEST=$NV/guest
[ -f "$CFG" ] || die "waydroid is not initialized — run 'waydroid init' first"
[ -d "$GUEST" ] || die "$GUEST missing (broken package?)"

GUEST_LIBS=(
    vendor/lib/hw/vulkan.virtio.so
    vendor/lib/egl/libEGL_angle.so
    vendor/lib/egl/libGLESv1_CM_angle.so
    vendor/lib/egl/libGLESv2_angle.so
    vendor/lib64/hw/vulkan.virtio.so
    vendor/lib64/egl/libEGL_angle.so
    vendor/lib64/egl/libGLESv1_CM_angle.so
    vendor/lib64/egl/libGLESv2_angle.so
    vendor/lib64/libgbm_mesa_wrapper.so
    vendor/lib64/hw/hwcomposer.waydroid.so
)
GUEST_BINS=(system/bin/surfaceflinger)

verify_android_elf() {
    local file="$1" abi="$2" expected_soname="${3:-}"
    local expected_class expected_machine header dynamic actual_soname
    case "$abi" in
        x86)
            expected_class=ELF32
            expected_machine="Intel 80386"
            ;;
        x86_64)
            expected_class=ELF64
            expected_machine="Advanced Micro Devices X86-64"
            ;;
        *) die "internal error: unsupported Android ABI '$abi'" ;;
    esac

    [ -f "$file" ] || die "required guest payload missing: $file"
    header=$(LC_ALL=C readelf -hW -- "$file" 2>/dev/null) || \
        die "$file is not a complete ELF file"
    if ! grep -Eq "^[[:space:]]*Class:[[:space:]]+${expected_class}[[:space:]]*$" <<<"$header" ||
       ! grep -Eq "^[[:space:]]*Machine:[[:space:]]+${expected_machine}[[:space:]]*$" <<<"$header" ||
       ! grep -Eq '^[[:space:]]*Type:[[:space:]]+DYN([[:space:]]|$)' <<<"$header"; then
        die "$file is not an $expected_class/$expected_machine shared object"
    fi

    if [ -n "$expected_soname" ]; then
        dynamic=$(LC_ALL=C readelf -dW -- "$file" 2>/dev/null) ||
            die "cannot read the dynamic section from $file"
        actual_soname=$(sed -n 's/.*(SONAME).*Library soname: \[\([^]]*\)\].*/\1/p' <<<"$dynamic")
        [ "$actual_soname" = "$expected_soname" ] ||
            die "$file has SONAME '${actual_soname:-<missing>}' (expected '$expected_soname')"
    fi
}

echo "== validating dual-ABI guest payload"
command -v readelf >/dev/null 2>&1 || die "readelf is required (install binutils)"
for rel in "${GUEST_LIBS[@]}"; do
    case "$rel" in
        vendor/lib/*) abi=x86 ;;
        vendor/lib64/*) abi=x86_64 ;;
        *) die "internal error: guest library has no ABI path: $rel" ;;
    esac
    case "$rel" in
        */hw/vulkan.virtio.so) soname=libvulkan_virtio.so ;;
        *) soname=$(basename -- "$rel") ;;
    esac
    verify_android_elf "$GUEST/$rel" "$abi" "$soname"
done
for rel in "${GUEST_BINS[@]}"; do
    verify_android_elf "$GUEST/$rel" x86_64
done
echo "   complete ELF headers, ABI types and library SONAMEs verified"

NVNODE=""
for uevent in /sys/class/drm/renderD*/device/uevent; do
    [ -e "$uevent" ] || continue
    if grep -qx "DRIVER=nvidia" "$uevent"; then
        NVNODE="/dev/dri/$(basename "${uevent%/device/uevent}")"
        break
    fi
done
[ -n "$NVNODE" ] || die "no NVIDIA render node found under /dev/dri (is the NVIDIA driver loaded?)"
echo "== NVIDIA render node: $NVNODE"

# nvidia-drm.modeset=1 is REQUIRED: without it the driver exposes no dma_buf
# support at all (no VK_EXT_external_memory_dma_buf, zero exportable
# modifiers — probe-verified on live hosts 2026-07-22) and the guest can
# only crash-loop wrapping its first buffer.
MODESET=$(cat /sys/module/nvidia_drm/parameters/modeset 2>/dev/null || echo missing)
[ "$MODESET" = "Y" ] || die "nvidia-drm.modeset is not enabled (got: $MODESET).
Add nvidia_drm.modeset=1 to your kernel parameters or modprobe.d and reboot."
echo "== nvidia-drm.modeset: Y"

# The stack hooks the AIDL minigbm gralloc; pre-minigbm vendor images give
# SurfaceFlinger buffers our driver never sees -> crash loop (issue #1:
# system and vendor images version independently, and 'upgrade -o' never
# refreshes them — migrated installs can carry an ancient vendor).
VIMG=/var/lib/waydroid/images/vendor.img
[ -f "$VIMG" ] || die "$VIMG missing — run 'waydroid init' first"
VMNT=$(mktemp -d)
mount -o loop,ro "$VIMG" "$VMNT" || { rmdir "$VMNT"; die "cannot inspect $VIMG (mount failed)"; }
HAVE_MINIGBM=0
[ -f "$VMNT/lib64/hw/gralloc.minigbm_gbm_mesa.so" ] && HAVE_MINIGBM=1
umount "$VMNT"; rmdir "$VMNT"
[ "$HAVE_MINIGBM" = 1 ] || die "vendor image is too old (no minigbm_gbm_mesa gralloc).
Run 'waydroid init -f' to fetch current images (apps/data are kept), then re-run this setup."
echo "== vendor image: minigbm AIDL gralloc present"

echo "== installing guest driver stack into $NV_GUEST"
install -d -m 0755 "$NV_GUEST"
for rel in "${GUEST_LIBS[@]}"; do
    install -Dm 0644 "$GUEST/$rel" "$NV_GUEST/$rel"
done
# Patched SurfaceFlinger is a 64-bit system executable, not an app-ABI library.
for rel in "${GUEST_BINS[@]}"; do
    install -Dm 0755 "$GUEST/$rel" "$NV_GUEST/$rel"
done

echo "== writing waydroid.cfg properties"
REFRESH="$REFRESH" NVNODE="$NVNODE" python3 - "$CFG" <<'EOF'
import configparser, os, sys
cfg_path = sys.argv[1]
cp = configparser.ConfigParser()
cp.optionxform = str  # waydroid props are case-sensitive
cp.read(cfg_path)
if not cp.has_section("properties"):
    cp.add_section("properties")

# Installing the NVIDIA stack obsoletes any old gralloc overrides — the
# classic software-rendering workaround (ro.hardware.gralloc=default) or
# gralloc.gbm.legacy from a migrated install blocks the guest's minigbm
# selection and crash-loops SurfaceFlinger (issue #1). Remove them, loudly.
for k in list(cp["properties"].keys()):
    if k == "ro.hardware.gralloc" or k.startswith("gralloc."):
        print("   removing stale gralloc override: {} = {}".format(
            k, cp["properties"][k]))
        cp.remove_option("properties", k)
props = {
    "ro.hardware.egl": "angle",
    "ro.hardware.vulkan": "virtio",
    "mesa.vn.debug": "vtest",
    "mesa.vtest.socket.name": "/dev/venus/venus.sock",
    "debug.hwui.renderer": "skiavk",
    "debug.renderengine.backend": "skiaglthreaded",
    "debug.sf.nobootanimation": "1",
    "persist.waydroid.use_subsurface": "true",
    "ro.surface_flinger.vsync_event_phase_offset_ns": "0",
    "ro.surface_flinger.vsync_sf_event_phase_offset_ns": "0",
    "ro.surface_flinger.max_frame_buffer_acquired_buffers": "3",
}
refresh = os.environ.get("REFRESH", "")
if refresh:
    props["persist.waydroid.refresh_rate"] = refresh
    if int(refresh) > 240:
        # SF's stock vsync snap window collapses timeslots above ~240 Hz;
        # the patched surfaceflinger reads this override (1 ms).
        props["debug.sf.snap_to_same_vsync_within_ns"] = "1000000"
for k, v in props.items():
    cp.set("properties", k, v)
cp.set("waydroid", "suspend_action", "none")
# Explicit drm_device bypasses upstream gpu.py's nvidia blacklist (our waydroid
# patch trusts it), so the render node gets bind-mounted on NVIDIA-only hosts.
cp.set("waydroid", "drm_device", os.environ["NVNODE"])
# Enables NVIDIA Venus mode in our waydroid patch: strict (non-optional)
# bind-mounts at container start + a socket preflight at session start.
cp.set("waydroid", "nvidia_venus_socket", "/run/waydroid-venus/venus.sock")
# Select the path-preserving dual-ABI mount layout under nv/guest.
cp.set("waydroid", "nvidia_guest_layout", "2")
with open(cfg_path, "w") as f:
    cp.write(f)
print("   properties written")
EOF

echo "== regenerating container config (waydroid upgrade -o)"
waydroid upgrade -o

cat <<'EOF'

Done. To start the stack:
  systemctl enable --now waydroid-container.service         (as root)
  systemctl --user enable --now wd-venus.service            (as your desktop user)
  waydroid session start                                    (or launch from the app menu)

The Venus render server (wd-venus) must be running before the session starts.
Verify GPU acceleration inside the guest: the SurfaceFlinger GLES line in
`waydroid shell dumpsys SurfaceFlinger | grep GLES` should name ANGLE/Venus
on your NVIDIA GPU.
EOF
