#!/bin/bash
# sunshine-gpu-helper — runtime GPU passthrough bind helper
# Runtime host-to-VM bind only. AMD return-to-host recovery still requires the
# established host shutdown, confirmed power-off, cooldown and Gaming startup.
# SUNSHINE_GPU_PREFLIGHT_PROTOCOL=2

set -e

# Always talk to the system libvirt daemon. Without this, a non-root
# invocation resolves to qemu:///session, whose daemon sees a different
# (stripped) set of domains and cannot use KVM.
export LIBVIRT_DEFAULT_URI="qemu:///system"

MODE="${1:-}"
LOGFILE="/var/log/sunshine-gpu-helper.log"

log() { echo "$(date '+%Y-%m-%d %H:%M:%S') [$1] $2" >> "$LOGFILE"; }

detect_gpu_pci() {
    # Pick the best AMD discrete GPU for passthrough. Skips:
    #  - non-AMD vendors
    #  - AMD Raphael onboard iGPU (1002:164e) — never use for VM passthrough
    #  - Display bridge subclass (0x80) — only accept VGA (0x00) or 3D (0x02)
    # Prefers devices with an audio sibling (discrete cards have one).
    local best="" best_score=0
    for ddir in /sys/bus/pci/devices/*/; do
        local v dev cls sub base pci
        read -r v < "$ddir/vendor" 2>/dev/null || continue
        [ "$v" = "0x1002" ] || continue
        read -r dev < "$ddir/device" 2>/dev/null || continue
        # Explicitly skip AMD Raphael onboard iGPU (1002:164e)
        [ "$dev" = "0x164e" ] && { log "INFO" "Skipping AMD Raphael iGPU $(basename "$ddir") (1002:164e)"; continue; }
        read -r cls < "$ddir/class" 2>/dev/null || continue
        base=$(( (cls >> 16) & 0xFF ))
        sub=$(( (cls >> 8) & 0xFF ))
        [ "$base" = 3 ] || continue
        # Accept only VGA (0x00) or 3D (0x02). Skip Display bridge (0x80).
        [ "$sub" = 0 -o "$sub" = 2 ] || continue
        pci=$(basename "$ddir")
        # Prefer devices with an audio sibling (discrete cards have HDMI/DP audio)
        local prefix="${pci%%.*}."
        local has_audio=0
        for sib in /sys/bus/pci/devices/${prefix}*/; do
            [ "$(basename "$sib")" = "$pci" ] && continue
            local sc
            read -r sc < "$sib/class" 2>/dev/null || continue
            [ $(( (sc >> 16) & 0xFF )) = 4 ] && { has_audio=1; break; }
        done
        local score=$(( (sub == 2 ? 2 : 1) + has_audio * 2 ))
        if [ "$score" -gt "$best_score" ]; then
            best="$pci"; best_score="$score"
        fi
    done
    [ -n "$best" ] && { echo "$best"; return 0; }
    return 1
}

detect_audio_pci() {
    local gpu_pci="$1"
    local prefix=""
    [ -n "$gpu_pci" ] && {
        local dot_pos="${gpu_pci##*.}"
        prefix="${gpu_pci%%.*}."
    }
    for ddir in /sys/bus/pci/devices/${prefix}*/; do
        local v cls
        read -r v < "$ddir/vendor" 2>/dev/null || continue
        [ "$v" = "0x1002" ] || continue
        read -r cls < "$ddir/class" 2>/dev/null || cls=""
        [ $(( (cls >> 16) & 0xFF )) = 4 ] || continue
        basename "$ddir"; return 0
    done
    return 1
}

gpu_driver() {
    local link
    link=$(readlink "/sys/bus/pci/devices/$1/driver" 2>/dev/null) || { echo "none"; return; }
    basename "$link"
}

if [ "$MODE" = passthrough ]; then
    # The service instance pins VM UUID, physical GPU and audio function.
    # Validate and, if needed, repair only the stopped VM's ROM reference
    # before display/console release or any sysfs writes.
    IFS=_ read -r VM_UUID EXPECTED_GPU EXPECTED_AUDIO EXTRA <<< "${2:-}"
    [ -n "$VM_UUID" ] && [ -n "$EXPECTED_GPU" ] && [ -n "$EXPECTED_AUDIO" ] && [ -z "$EXTRA" ] || {
        echo "ERROR: Validated VM/GPU service instance required" >&2; exit 1;
    }
    PREFLIGHT="$(dirname "$(readlink -f "$0")")/gpu_preflight.py"
    PLAN=$(/usr/bin/python3 -I "$PREFLIGHT" --vm-uuid "$VM_UUID" --gpu "$EXPECTED_GPU" --audio "$EXPECTED_AUDIO" --helper-plan --repair-rom) || exit 1
    mapfile -t CHECKED <<< "$PLAN"
    [ "${#CHECKED[@]}" -eq 5 ] || { echo "ERROR: Invalid preflight result" >&2; exit 1; }
    VM_NAME="${CHECKED[0]}"
    GPU_PCI="${CHECKED[1]}"
    AUDIO_PCI="${CHECKED[2]}"
    PREFLIGHT_ROM="${CHECKED[3]}"
    [ "${CHECKED[4]}" = 1 ] && { echo "PASSTHROUGH_OK VM_ALREADY_RUNNING"; exit 0; }
else
    GPU_PCI=$(detect_gpu_pci)
    AUDIO_PCI=$(detect_audio_pci "$GPU_PCI")
fi

[ -z "$GPU_PCI" ] && { echo "ERROR: No AMD GPU found"; exit 1; }

case "$MODE" in
    passthrough)
        log "INFO" "Runtime passthrough bind: GPU=$GPU_PCI Audio=$AUDIO_PCI VM=$VM_NAME"

        # Stop display manager
        for dm in lightdm gdm3 sddm xdm lxdm; do
            if systemctl is-active --quiet "$dm" 2>/dev/null; then
                systemctl stop "$dm" 2>/dev/null && log "INFO" "Stopped $dm"
                echo "$dm" > /var/lib/sunshine/dm_name
                break
            fi
        done
        sleep 1

        # Release console
        for vtcon in /sys/class/vtconsole/vtcon*/bind; do
            echo 0 > "$vtcon" 2>/dev/null && log "INFO" "Unbound $(dirname "$vtcon")" || true
        done

        # Load vfio-pci
        modprobe vfio-pci 2>/dev/null || true

        # Register GPU IDs with vfio-pci (auto-detect from sysfs)
        GPU_VID=$(cat /sys/bus/pci/devices/$GPU_PCI/vendor 2>/dev/null | sed 's/0x//')
        GPU_DID=$(cat /sys/bus/pci/devices/$GPU_PCI/device 2>/dev/null | sed 's/0x//')
        echo "$GPU_VID $GPU_DID" > /sys/bus/pci/drivers/vfio-pci/new_id 2>/dev/null || true
        log "INFO" "Registered GPU ID: $GPU_VID $GPU_DID"
        if [ -n "$AUDIO_PCI" ]; then
            AUD_VID=$(cat /sys/bus/pci/devices/$AUDIO_PCI/vendor 2>/dev/null | sed 's/0x//')
            AUD_DID=$(cat /sys/bus/pci/devices/$AUDIO_PCI/device 2>/dev/null | sed 's/0x//')
            echo "$AUD_VID $AUD_DID" > /sys/bus/pci/drivers/vfio-pci/new_id 2>/dev/null || true
            log "INFO" "Registered Audio ID: $AUD_VID $AUD_DID"
        fi

        # Unbind from amdgpu (idempotent: nothing to do if already on vfio-pci,
        # e.g. the boot-time rebind running after a manual bind)
        gpu_drv=$(gpu_driver "$GPU_PCI")
        if [ "$gpu_drv" = "vfio-pci" ]; then
            log "INFO" "GPU already bound to vfio-pci; bind skipped"
            mkdir -p /var/lib/sunshine
            echo passthrough > /var/lib/sunshine/gpu_mode
        elif [ "$gpu_drv" = "amdgpu" ]; then
            echo "$GPU_PCI" > /sys/bus/pci/drivers/amdgpu/unbind 2>/dev/null && \
                log "INFO" "Unbound GPU from amdgpu" || \
                { log "ERROR" "GPU unbind failed"; exit 1; }
        fi
        if [ -n "$AUDIO_PCI" ]; then
            ad=$(gpu_driver "$AUDIO_PCI")
            [ "$ad" = "snd_hda_intel" ] && echo "$AUDIO_PCI" > /sys/bus/pci/drivers/snd_hda_intel/unbind 2>/dev/null
        fi

        sleep 1

        # Bind to vfio-pci (skip if already bound)
        gpu_drv=$(gpu_driver "$GPU_PCI")
        if [ "$gpu_drv" != "vfio-pci" ]; then
            echo "$GPU_PCI" > /sys/bus/pci/drivers/vfio-pci/bind 2>/dev/null && \
                log "INFO" "GPU bound to vfio-pci" || \
                { log "ERROR" "GPU bind failed"; exit 1; }
        fi
        [ -n "$AUDIO_PCI" ] && echo "$AUDIO_PCI" > /sys/bus/pci/drivers/vfio-pci/bind 2>/dev/null

        mkdir -p /var/lib/sunshine
        echo passthrough > /var/lib/sunshine/gpu_mode
        log "INFO" "Passthrough complete"

        # Start VM if name was provided
        if [ -n "$VM_NAME" ]; then
            log "INFO" "Starting VM: $VM_NAME"
            # Ensure libvirt default network is in a clean state before VM start.
            # After a host reboot, libvirtd reports the "default" network as
            # "active" but the virbr0 bridge interface may be missing, causing
            # VM start to fail with "Network not found: no network with
            # matching name 'virbr0'". Only cycle if there's a mismatch
            # (network active but bridge missing) to avoid disrupting
            # already-running VMs on the same network.
            NET_ACTIVE=$(virsh net-info default 2>/dev/null | grep 'Active:' | awk '{print $2}')
            BRIDGE_EXISTS=no
            [ -e /sys/devices/virtual/net/virbr0 ] && BRIDGE_EXISTS=yes
            if [ "$NET_ACTIVE" = "yes" ] && [ "$BRIDGE_EXISTS" = "no" ]; then
                log "INFO" "Cycling libvirt default network (active but virbr0 missing)"
                virsh net-destroy default 2>/dev/null || true
                sleep 1
                virsh net-start default 2>/dev/null || true
            elif [ "$NET_ACTIVE" != "yes" ]; then
                log "INFO" "Starting libvirt default network (was inactive)"
                virsh net-start default 2>/dev/null || true
            fi
            virsh net-autostart default 2>/dev/null || true
            # Sync the PulseAudio endpoint (pid + cookie) QEMU's pa backend
            # needs. pipewire-pulse restarts change these; a stale pid or
            # cookie leaves the VM audio device permanently in "Bad state".
            if [ -x /usr/libexec/sunshine/vm-audio-setup ]; then
                SYNC_ONLY=1 /usr/libexec/sunshine/vm-audio-setup >>"$LOGFILE" 2>&1 || \
                    log "WARNING" "vm-audio-setup sync failed (audio may not work)"
            fi

            # Attach the GPU (and its audio function) to the VM for streaming.
            # Revalidate again at use time. A changed/invalid ROM must never be
            # attached even if it changed after the pre-disruption check.
            PLAN=$(/usr/bin/python3 -I "$PREFLIGHT" --vm-uuid "$VM_NAME" --gpu "$GPU_PCI" --audio "$AUDIO_PCI" --helper-plan) || exit 1
            mapfile -t CHECKED <<< "$PLAN"
            PREFLIGHT_ROM="${CHECKED[3]}"
            # Install/console boots use plain virtio video, so the hostdev is
            # only attached here, at passthrough time (and only once).
            #
            # Stale-hostdev guard: when the host is re-provisioned with a
            # different GPU (possibly at a different PCI address), a VM that
            # still carries its old hostdev fails `virsh start` with
            # "no such device" on every attempt. Drop hostdevs whose
            # <source> address no longer matches the current GPU/audio so
            # they re-attach with fresh addresses below.
            if virsh dumpxml "$VM_NAME" --inactive 2>/dev/null | grep -q "<hostdev" \
                && command -v python3 >/dev/null 2>&1; then
                GPU_PCI="$GPU_PCI" AUDIO_PCI="${AUDIO_PCI:-}" python3 - "$VM_NAME" >>"$LOGFILE" 2>&1 << 'PYEOF' || true
import os, re, subprocess, sys, tempfile

name = sys.argv[1]

def parse_pci(pci):
    # "0000:03:00.1" -> ("0x03", "0x00", "0x1")
    rest = pci[len("0000:"):]
    bus, rest = rest.split(":", 1)
    slot, func = rest.rsplit(".", 1)
    return ("0x" + bus, "0x" + slot, "0x" + func)

wanted = {parse_pci(p) for p in
          (os.environ.get("GPU_PCI", ""), os.environ.get("AUDIO_PCI", "")) if p}

xml = subprocess.run(["virsh", "dumpxml", name, "--inactive"],
                     capture_output=True, text=True).stdout

def hostdev_source_addr(block):
    # Host address lives inside <source>; the top-level <address> is the
    # GUEST-side slot and must not be compared.
    src = re.search(r"<source>(.*?)</source>", block, re.S)
    if not src:
        return None
    d = dict(re.findall(r"(bus|slot|function)='(0x[0-9a-fA-F]+)'", src.group(1)))
    return (d.get("bus"), d.get("slot"), d.get("function"))

out, pos, stale = [], 0, False
while True:
    s = xml.find("<hostdev", pos)
    if s == -1:
        out.append(xml[pos:])
        break
    e = xml.find("</hostdev>", s)
    if e == -1:
        out.append(xml[pos:])
        break
    out.append(xml[pos:s])  # text before the block is always kept
    if hostdev_source_addr(xml[s:e]) in wanted:
        out.append(xml[s:e + len("</hostdev>")])
    else:
        stale = True  # GPU/audio moved, or a device this host no longer has
    pos = e + len("</hostdev>")

if stale:
    cleaned = "".join(out)
    with tempfile.NamedTemporaryFile("w", suffix=".xml", delete=False) as f:
        f.write(cleaned)
        tmp = f.name
    r = subprocess.run(["virsh", "define", tmp], capture_output=True, text=True)
    os.unlink(tmp)
    if r.returncode == 0:
        print("stripped stale GPU hostdev(s) from %s; re-attaching with current addresses" % name)
    else:
        print("WARNING: failed to strip stale hostdev(s): %s" % r.stderr.strip())
PYEOF
            fi
            if ! virsh dumpxml "$VM_NAME" --inactive 2>/dev/null | grep -q "<hostdev"; then
                pci_addr() {
                    # 0000:03:00.0 → bus/slot/function as hex strings
                    local rest b s f
                    rest="${1#0000:}"
                    b="${rest%%:*}"
                    rest="${rest#*:}"
                    s="${rest%%.*}"
                    f="${rest##*.}"
                    printf "bus='0x%s' slot='0x%s' function='0x%s'" "$b" "$s" "$f"
                }
                ROM_LINE=""
                if [ -n "$PREFLIGHT_ROM" ]; then
                    ROM_LINE="    <rom file='$PREFLIGHT_ROM'/>"
                fi
                cat > /tmp/sunshine-gpu-hostdev.xml <<EOF
<hostdev mode='subsystem' type='pci' managed='yes'>
  <source><address domain='0x0000' $(pci_addr "$GPU_PCI")/></source>
${ROM_LINE}
</hostdev>
EOF
                if virsh attach-device "$VM_NAME" /tmp/sunshine-gpu-hostdev.xml --config 2>>"$LOGFILE"; then
                    log "INFO" "Attached GPU $GPU_PCI to $VM_NAME"
                else
                    log "WARNING" "Failed to attach GPU hostdev to $VM_NAME"
                fi
                if [ -n "$AUDIO_PCI" ]; then
                    cat > /tmp/sunshine-audio-hostdev.xml <<EOF
<hostdev mode='subsystem' type='pci' managed='yes'>
  <source><address domain='0x0000' $(pci_addr "$AUDIO_PCI")/></source>
</hostdev>
EOF
                    if virsh attach-device "$VM_NAME" /tmp/sunshine-audio-hostdev.xml --config 2>>"$LOGFILE"; then
                        log "INFO" "Attached GPU audio $AUDIO_PCI to $VM_NAME"
                    else
                        log "WARNING" "Failed to attach GPU audio hostdev to $VM_NAME"
                    fi
                fi
                rm -f /tmp/sunshine-gpu-hostdev.xml /tmp/sunshine-audio-hostdev.xml
            fi

            if virsh start "$VM_NAME" 2>/dev/null; then
                log "INFO" "VM $VM_NAME started successfully"
                echo "PASSTHROUGH_OK VM_STARTED"
            else
                log "ERROR" "Failed to start VM: $VM_NAME"
                echo "PASSTHROUGH_OK VM_START_FAILED"
                exit 1
            fi
        else
            echo "PASSTHROUGH_OK"
        fi
        ;;

    status)
        echo "GPU: $GPU_PCI driver=$(gpu_driver "$GPU_PCI") mode_file=$(cat /var/lib/sunshine/gpu_mode 2>/dev/null || echo unknown)"
        [ -n "$AUDIO_PCI" ] && echo "AUDIO: $AUDIO_PCI driver=$(gpu_driver "$AUDIO_PCI")"
        ;;

    dump-rom)
        # Dump GPU ROM to /var/lib/libvirt/images/gpu.rom
        # Must run while GPU is on amdgpu (before passthrough bind).
        # Called at boot by sunshine-gpu-rom-dump.service.
        ROM_PATH="/var/lib/libvirt/images/gpu.rom"
        mkdir -p /var/lib/libvirt/images

        gpu_drv=$(gpu_driver "$GPU_PCI")
        if [ "$gpu_drv" != "amdgpu" ]; then
            log "WARNING" "dump-rom: GPU not on amdgpu (driver=$gpu_drv), skipping"
            echo "SKIP: GPU not on amdgpu"
            exit 0
        fi

        # Validate a dumped ROM: size floor, 55AA signature, and a PCIR entry
        # whose vendor/device match the physical card. A flaky sysfs read at
        # boot (GPU mid-init or power-state transition) otherwise produces a
        # corrupt image that the guest firmware then uses to enumerate the
        # WRONG card — breaking the guest display driver and the stream.
        validate_rom() {
            local f="$1"
            [ "$(stat -c%s "$f" 2>/dev/null || echo 0)" -lt 32768 ] && return 1
            [ "$(dd if="$f" bs=2 count=1 2>/dev/null | od -An -tx1 | tr -d ' \n')" != "55aa" ] && return 1
            grep -aq "PCIR" "$f" 2>/dev/null || return 1
            local want
            want=$(cat "/sys/bus/pci/devices/$GPU_PCI/device" 2>/dev/null)
            want=${want#0x}
            if command -v python3 >/dev/null 2>&1 && [ -n "$want" ]; then
                python3 - "$f" "$want" << 'PYEOF'
import sys
d = open(sys.argv[1], 'rb').read()
want = int(sys.argv[2], 16)
i = d.find(b'PCIR')
while i != -1:
    ven = int.from_bytes(d[i+4:i+6], 'little')
    dev = int.from_bytes(d[i+6:i+8], 'little')
    if ven == 0x1002 and dev == want:
        sys.exit(0)
    i = d.find(b'PCIR', i + 1)
sys.exit(1)
PYEOF
            fi
        }

        TMP_ROM="$ROM_PATH.tmp.$$"
        ok=0
        attempt=1
        while [ $attempt -le 5 ]; do
            echo 1 > "/sys/bus/pci/devices/$GPU_PCI/rom" 2>/dev/null || true
            cat "/sys/bus/pci/devices/$GPU_PCI/rom" > "$TMP_ROM" 2>/dev/null
            echo 0 > "/sys/bus/pci/devices/$GPU_PCI/rom" 2>/dev/null || true
            if validate_rom "$TMP_ROM"; then
                ok=1
                break
            fi
            log "WARNING" "dump-rom: attempt $attempt produced an invalid ROM, retrying"
            sleep 2
            attempt=$((attempt + 1))
        done

        if [ $ok -eq 1 ]; then
            mv -f "$TMP_ROM" "$ROM_PATH"
            chmod 644 "$ROM_PATH"
            ROM_SIZE=$(stat -c%s "$ROM_PATH" 2>/dev/null || echo 0)
            log "INFO" "dump-rom: Dumped $ROM_SIZE bytes to $ROM_PATH"
            echo "ROM_OK: $ROM_PATH ($ROM_SIZE bytes)"
        else
            rm -f "$TMP_ROM"
            if [ -s "$ROM_PATH" ] && validate_rom "$ROM_PATH"; then
                log "WARNING" "dump-rom: fresh dump invalid; keeping previous known-good ROM"
                echo "ROM_KEEP_PREVIOUS: $ROM_PATH"
                exit 0
            fi
            log "ERROR" "dump-rom: could not obtain a valid ROM after 5 attempts"
            echo "ERROR: ROM dump failed"
            exit 1
        fi
        ;;

    reconcile)
        # Boot-time reconcile: record the ACTUAL driver state in gpu_mode.
        # Never binds the GPU — bare metal must always be available after a
        # cold boot, and passthrough is engaged only when a VM is started
        # from the Maximizer client (any VM — with several, the user picks).
        mkdir -p /var/lib/sunshine
        drv=$(gpu_driver "$GPU_PCI")
        if [ "$drv" = "vfio-pci" ]; then
            echo passthrough > /var/lib/sunshine/gpu_mode
        else
            echo baremetal > /var/lib/sunshine/gpu_mode
        fi
        log "INFO" "reconcile: GPU on $drv -> gpu_mode recorded"
        echo "RECONCILE_OK: $drv"
        ;;

    *)
        echo "Usage: $0 {passthrough|status|dump-rom|reconcile}"; exit 1 ;;
esac
