#!/bin/bash
# ============================================================
#  fix-gpu-vm.sh  —  Run this INSIDE the VM (your app server)
#  Run AFTER fix-gpu-host.sh has completed and host has rebooted
# ============================================================
set -euo pipefail

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'

log()  { echo -e "${GREEN}[✓]${NC} $1"; }
warn() { echo -e "${YELLOW}[!]${NC} $1"; }
fail() { echo -e "${RED}[✗]${NC} $1"; exit 1; }
info() { echo -e "${BLUE}[i]${NC} $1"; }

echo ""
echo "======================================================"
echo "  NVIDIA Driver Fix — VM Userspace"
echo "======================================================"
echo ""

[[ $EUID -ne 0 ]] && fail "Please run as root: sudo bash fix-gpu-vm.sh"

FFMPEG="/usr/lib/jellyfin-ffmpeg/ffmpeg"
TARGET="580.65.06"
APP_DIR="/var/www/videoplatform"

# ── 1. Check Proxmox host kernel module ──────────────────────────────────────
HOST_VER=$(cat /proc/driver/nvidia/version 2>/dev/null | grep -oP '\d+\.\d+\.\d+' | head -1 || echo "unknown")
info "Proxmox host kernel module: ${HOST_VER}"

if [[ "$HOST_VER" != "580.65.06" ]]; then
    warn "Host kernel module is ${HOST_VER}, not 580.65.06."
    warn "Please run fix-gpu-host.sh on the Proxmox host first, then reboot the host."
    read -p "Continue anyway? (yes/no): " FORCE
    [[ "$FORCE" != "yes" ]] && { warn "Aborted."; exit 0; }
fi

# ── 2. Check jellyfin-ffmpeg is installed ────────────────────────────────────
[[ -x "$FFMPEG" ]] || fail "jellyfin-ffmpeg not found at ${FFMPEG}. Install it first: dpkg -i /tmp/jellyfin-ffmpeg7.deb"
log "jellyfin-ffmpeg found: $($FFMPEG -version 2>&1 | head -1)"

# ── 3. Downgrade NVIDIA userspace packages ───────────────────────────────────
CURRENT_COMPUTE=$(dpkg -l libnvidia-compute-580 2>/dev/null | awk '/^ii/{print $3}' | cut -d- -f1 || echo "none")
info "Current libnvidia-compute-580 version: ${CURRENT_COMPUTE}"

if [[ "$CURRENT_COMPUTE" == "$TARGET" ]]; then
    log "NVIDIA userspace already at ${TARGET}. Skipping downgrade."
else
    warn "Downgrading NVIDIA userspace from ${CURRENT_COMPUTE} → ${TARGET}..."
    echo ""

    UBUNTU_VER="${TARGET}-0ubuntu0.22.04.1"

    apt-get install -y --allow-downgrades \
        "libnvidia-compute-580=${UBUNTU_VER}" \
        "libnvidia-encode-580=${UBUNTU_VER}"  \
        "libnvidia-decode-580=${UBUNTU_VER}"  \
        "nvidia-utils-580=${UBUNTU_VER}"      || \
        fail "Package downgrade failed. Try: apt-get update && then re-run this script."

    log "Packages downgraded successfully."

    # Pin to prevent apt from auto-upgrading and breaking things again
    for pkg in libnvidia-compute-580 libnvidia-encode-580 libnvidia-decode-580 nvidia-utils-580; do
        apt-mark hold "$pkg"
    done
    log "Packages pinned at ${TARGET} to prevent accidental upgrades."
fi

# ── 4. Test NVENC ─────────────────────────────────────────────────────────────
echo ""
info "Testing NVENC encoding with new setup..."
TMP=$(mktemp /tmp/nvenc_test_XXXX.mp4)

"$FFMPEG" -f lavfi -i color=c=black:s=128x72:r=25 -frames:v 1 \
    -c:v h264_nvenc -y "$TMP" > /tmp/nvenc_test.log 2>&1
EXIT=$?

if [[ $EXIT -eq 0 && -s "$TMP" ]]; then
    log "NVENC test PASSED! GPU encoding is working."
    rm -f "$TMP"
    NVENC_OK=true
else
    warn "NVENC test failed (exit ${EXIT}). Falling back to CPU encoding."
    cat /tmp/nvenc_test.log | tail -5
    rm -f "$TMP"
    NVENC_OK=false
fi

# ── 5. Update app settings ────────────────────────────────────────────────────
echo ""
info "Updating application settings..."

if [[ -d "$APP_DIR" ]]; then
    # Set ffmpeg binary path
    php "$APP_DIR/artisan" tinker --execute="
        App\Models\Setting::set('ffmpeg_binary', '$FFMPEG');
        echo 'ffmpeg_binary set to: ' . App\Models\Setting::ffmpegBinary();
    " 2>/dev/null && log "App FFmpeg binary path set to: ${FFMPEG}"

    if [[ "$NVENC_OK" == "true" ]]; then
        php "$APP_DIR/artisan" tinker --execute="
            App\Models\Setting::set('gpu_enabled', 'true');
            echo 'gpu_enabled: ' . App\Models\Setting::get('gpu_enabled');
        " 2>/dev/null && log "GPU encoding enabled in app settings."
    else
        php "$APP_DIR/artisan" tinker --execute="
            App\Models\Setting::set('gpu_enabled', 'false');
            echo 'gpu_enabled set to false (CPU fallback active)';
        " 2>/dev/null && warn "GPU disabled in app settings — CPU fallback active."
    fi
else
    warn "App directory not found at ${APP_DIR}, skipping settings update."
fi

# ── 6. Summary ────────────────────────────────────────────────────────────────
echo ""
echo "======================================================"
if [[ "$NVENC_OK" == "true" ]]; then
    echo -e "${GREEN}  ALL DONE — GPU encoding is active!${NC}"
    echo ""
    echo "  FFmpeg binary : ${FFMPEG}"
    echo "  GPU encoder   : h264_nvenc (NVIDIA RTX 3060)"
    echo "  Driver match  : kernel=${HOST_VER} / userspace=${TARGET}"
else
    echo -e "${YELLOW}  DONE — Running with CPU fallback (libx264)${NC}"
    echo ""
    echo "  Videos will still encode correctly, just slightly slower."
    echo "  Check /tmp/nvenc_test.log for NVENC error details."
fi
echo "======================================================"
echo ""