#!/bin/bash
# ============================================================================
# Open-IOV install.sh - Reproduce the complete vGPU stack on a fresh Ubuntu
#
# This script takes a plain Ubuntu 22.04/24.04 machine with a supported
# NVIDIA GPU (Turing TU102/TU104 validated) to a running vGPU stack:
#
#   1. deps      Install build/runtime dependencies (apt)
#   2. nvidia    Clone NVIDIA open-gpu-kernel-modules @ pinned tag + patch
#   3. firmware  Fetch GSP firmware from linux-firmware.git + extract fwimage
#   4. build     Build vgsp.ko -> nvidia.ko -> userspace (in order)
#   5. system    Host config: VFIO modules, hugepages, IOMMU/blacklist check
#   6. load      Load the driver stack (root)
#   7. mdev      Create a vGPU mdev device (root)
#
# After install.sh completes, attach the mdev to ANY QEMU wrapper you like
# (libvirt <hostdev>, or raw QEMU -device vfio-pci,sysfsdev=...).
# A reference raw-QEMU launcher is provided: scripts/launch_qemu_vgpu.sh
#
# Read first (progressive disclosure):
#   docs/02-getting-started/01-install-and-first-run.md  - this flow, explained
#   docs/README.md                                       - project front door
#   docs/03-architecture/01-system-overview.md           - how it works
#
# License: BSD-3-Clause (project) / GPL-2.0 (kernel modules). See LICENSE.
#
# ----------------------------------------------------------------------------
# FUTURE: DKMS PACKAGE
# ----------------------------------------------------------------------------
# This script is the "bring up a development machine" path. A future release
# will additionally distribute the kernel modules as a **DKMS package**, so
# users can install Open-IOV on a plain Ubuntu system (or any DKMS-capable
# distribution) with:
#
#     sudo apt install open-iov-dkms        # (planned; not yet available)
#
# and get automatic module rebuilds on kernel upgrades - no source checkout,
# no build tooling, no manual insmod ordering. The DKMS package will carry the
# same configuration profiles (config/) and the same firmware-fetch logic that
# this script performs in the `firmware` stage. Until then, install.sh is the
# supported path.
#
# Usage:
#   ./install.sh                 # full install (all stages)
#   ./install.sh deps nvidia     # run only selected stages
#   ./install.sh --no-load       # everything except load+mdev (no root needed
#                                #  until the 'system' stage)
#
# Environment overrides:
#   NVIDIA_TAG          NVIDIA open-gpu-kernel-modules tag (default 610.57.04)
#   GSP_VERSION         GSP firmware version        (default 570.144)
#   FW_MD5              Expected fwimage MD5        (default: see below)
#   SKIP_FIRMWARE=1     Skip firmware stage
# ============================================================================

set -euo pipefail

OPEN_IOV_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PATCH_DIR="$OPEN_IOV_ROOT/legacy-patches/nvidia"
NVIDIA_DIR="$OPEN_IOV_ROOT/nvidia-open-gpu-kernel-modules"
NVIDIA_REPO="https://github.com/NVIDIA/open-gpu-kernel-modules.git"
NVIDIA_TAG="${NVIDIA_TAG:-610.57.04}"

GSP_VERSION="${GSP_VERSION:-570.144}"
# linux-firmware.git: nvidia/tu104/gsp is a symlink to nvidia/tu102/gsp
FW_GIT_REPO="https://github.com/NVIDIA/linux-firmware.git"
FW_GIT_COMMIT="2237b99e29d09bebffdb706ae05005d9026fd7e4"   # known-good pin
FW_RAW_BASE="https://raw.githubusercontent.com/NVIDIA/linux-firmware/${FW_GIT_COMMIT}/nvidia/tu102/gsp"
FW_DIR="$OPEN_IOV_ROOT/firmware/57X-GSP/570.144"
FW_MD5="${FW_MD5:-13ce12a392a31eeb79a2e6ebd3a4a93a}"

NR_HUGEPAGES="${NR_HUGEPAGES:-8192}"

color_red=$'\033[0;31m'; color_green=$'\033[0;32m'
color_yellow=$'\033[1;33m'; color_blue=$'\033[0;34m'; color_nc=$'\033[0m'
log()  { echo "${color_blue}[install]${color_nc} $*"; }
ok()   { echo "${color_green}[  ok  ]${color_nc} $*"; }
warn() { echo "${color_yellow}[ warn ]${color_nc} $*"; }
die()  { echo "${color_red}[ FAIL ]${color_nc} $*" >&2; exit 1; }

need_sudo() {
    if [[ $EUID -ne 0 ]]; then
        command -v sudo >/dev/null || die "stage requires root and sudo is missing"
        echo sudo
    fi
}

# ----------------------------------------------------------------------------
stage_deps() {
    log "Installing dependencies (apt)..."
    local SUDO; SUDO=$(need_sudo)
    $SUDO apt-get update
    $SUDO apt-get install -y --no-install-recommends \
        build-essential gcc make bc flex bison libssl-dev libelf-dev \
        linux-headers-$(uname -r) \
        git ca-certificates curl wget \
        python3 python3-pip \
        binutils uuid-runtime pkg-config \
        qemu-system-x86 qemu-utils ovmf \
        pciutils kmod gdb
    ok "dependencies installed"
}

# ----------------------------------------------------------------------------
stage_nvidia() {
    if [[ -d "$NVIDIA_DIR/.git" ]]; then
        ok "NVIDIA tree already present ($NVIDIA_DIR)"
    else
        log "Cloning NVIDIA open-gpu-kernel-modules @ $NVIDIA_TAG ..."
        git clone --depth 1 --branch "$NVIDIA_TAG" "$NVIDIA_REPO" "$NVIDIA_DIR"
    fi

    log "Applying Open-IOV patches from $PATCH_DIR ..."
    cd "$NVIDIA_DIR"
    local applied=0 p
    for p in "$PATCH_DIR"/*.patch; do
        [[ -e "$p" ]] || die "no patches found in $PATCH_DIR"
        if git apply --check "$p" 2>/dev/null; then
            git apply "$p"
            applied=$((applied+1))
            log "  applied $(basename "$p")"
        elif git apply --reverse --check "$p" 2>/dev/null; then
            ok "  $(basename "$p") already applied"
        else
            die "patch does not apply cleanly: $p (expected pristine tag $NVIDIA_TAG)"
        fi
    done
    ok "NVIDIA tree patched ($NVIDIA_TAG + $applied patch(es))"
}

# ----------------------------------------------------------------------------
stage_firmware() {
    [[ "${SKIP_FIRMWARE:-0}" == "1" ]] && { warn "firmware stage skipped"; return; }
    mkdir -p "$FW_DIR"
    local raw_bin="$FW_DIR/gsp-${GSP_VERSION}.bin"
    local fwimage="$FW_DIR/gsp-${GSP_VERSION}.fwimage.bin"

    if [[ -f "$fwimage" ]] && echo "$FW_MD5  $fwimage" | md5sum -c --status 2>/dev/null; then
        ok "fwimage already present and verified"
    else
        log "Fetching GSP $GSP_VERSION firmware (linux-firmware, nvidia/tu102/gsp)..."
        if ! curl -fSL --retry 3 -o "$raw_bin" "$FW_RAW_BASE/gsp-${GSP_VERSION}.bin"; then
            warn "raw download failed; falling back to sparse git clone"
            local tmp; tmp=$(mktemp -d)
            git clone --depth 1 --filter=blob:none --sparse "$FW_GIT_REPO" "$tmp/lf"
            git -C "$tmp/lf" sparse-checkout set nvidia/tu102/gsp
            cp "$tmp/lf/nvidia/tu102/gsp/gsp-${GSP_VERSION}.bin" "$raw_bin"
            rm -rf "$tmp"
        fi

        log "Extracting .fwimage section -> $(basename "$fwimage")"
        objcopy -O binary --only-section=.fwimage "$raw_bin" "$fwimage"
        echo "$FW_MD5  $fwimage" | md5sum -c --status \
            || die "fwimage MD5 mismatch (expected $FW_MD5) - firmware corrupt or wrong version"
        ok "fwimage extracted and verified (md5 $FW_MD5)"
    fi

    # Install the *raw* GSP package for the host nvidia.ko physical-GSP boot
    log "Installing GSP firmware to /lib/firmware/nvidia/{tu102,tu104}/gsp ..."
    local SUDO; SUDO=$(need_sudo)
    local f
    for f in gsp-${GSP_VERSION}.bin bootloader-${GSP_VERSION}.bin \
             booter_load-${GSP_VERSION}.bin booter_unload-${GSP_VERSION}.bin; do
        [[ -f "$FW_DIR/$f" ]] || curl -fSL --retry 3 -o "$FW_DIR/$f" "$FW_RAW_BASE/$f" \
            || warn "could not fetch $f (host GSP boot may fail)"
    done
    for arch in tu102 tu104; do
        $SUDO mkdir -p "/lib/firmware/nvidia/$arch/gsp"
        for f in gsp-${GSP_VERSION}.bin bootloader-${GSP_VERSION}.bin \
                 booter_load-${GSP_VERSION}.bin booter_unload-${GSP_VERSION}.bin; do
            [[ -f "$FW_DIR/$f" ]] && $SUDO cp -f "$FW_DIR/$f" "/lib/firmware/nvidia/$arch/gsp/$f"
        done
    done
    # fwimage location used by scripts/stack_manager.sh install-firmware
    $SUDO cp -f "$fwimage" "/lib/firmware/gsp-${GSP_VERSION}.fwimage.bin"
    ok "firmware installed"
}

# ----------------------------------------------------------------------------
stage_build() {
    log "Building vgsp.ko (generates Module.symvers)..."
    make -C "$OPEN_IOV_ROOT" kernel

    log "Building NVIDIA OpenRM (3-step, KBUILD_EXTRA_SYMBOLS)..."
    make -C "$NVIDIA_DIR/src/nvidia" -j"$(nproc)"
    make -C "$NVIDIA_DIR/src/nvidia-modeset" -j"$(nproc)"
    make -C "$NVIDIA_DIR/kernel-open" -j"$(nproc)" \
        KBUILD_EXTRA_SYMBOLS="$OPEN_IOV_ROOT/kernel/Module.symvers" modules

    log "Verifying vGPU VFIO symbols resolved (modinfo depends)..."
    modinfo "$NVIDIA_DIR/kernel-open/nvidia.ko" | grep -q "^depends:.*vgsp" \
        || die "nvidia.ko missing 'depends: vgsp' - KBUILD_EXTRA_SYMBOLS resolution failed"

    log "Building userspace (libiov.so + examples)..."
    make -C "$OPEN_IOV_ROOT" userspace examples
    ok "build complete"
}

# ----------------------------------------------------------------------------
stage_system() {
    log "Configuring host system..."
    local SUDO; SUDO=$(need_sudo)

    # VFIO/mdev modules at boot
    if [[ ! -f /etc/modules-load.d/vgsp.conf ]]; then
        printf 'vfio\nvfio_iommu_type1\nvfio_pci\nmdev\nvfio_mdev\n' \
            | $SUDO tee /etc/modules-load.d/vgsp.conf >/dev/null
        ok "wrote /etc/modules-load.d/vgsp.conf"
    else
        ok "modules-load.d already configured"
    fi
    $SUDO modprobe vfio vfio_iommu_type1 vfio_pci mdev vfio_mdev 2>/dev/null || true

    # Hugepages (used by the reference QEMU launcher; harmless otherwise)
    if [[ "$(cat /proc/sys/vm/nr_hugepages)" -lt "$NR_HUGEPAGES" ]]; then
        echo "$NR_HUGEPAGES" | $SUDO tee /proc/sys/vm/nr_hugepages >/dev/null
        ok "nr_hugepages=$NR_HUGEPAGES"
    fi
    if ! mountpoint -q /dev/hugepages 2>/dev/null; then
        $SUDO mkdir -p /dev/hugepages
        $SUDO mount -t hugetlbfs hugetlbfs /dev/hugepages 2>/dev/null || true
    fi

    # IOMMU + blacklist sanity checks (informational, never fatal)
    local cmdline; cmdline=$(cat /proc/cmdline)
    if ! grep -qE 'intel_iommu=on|amd_iommu=on' <<<"$cmdline"; then
        warn "IOMMU does not appear enabled in kernel cmdline."
        warn "Add 'intel_iommu=on' (Intel) or 'amd_iommu=on' (AMD) to"
        warn "GRUB_CMDLINE_LINUX_DEFAULT in /etc/default/grub and run update-grub."
    fi
    if grep -q 'module_blacklist=.*nvidia' <<<"$cmdline"; then
        warn "kernel cmdline blacklists nvidia modules - remove from GRUB and reboot."
    fi
    ok "system stage complete"
}

# ----------------------------------------------------------------------------
stage_load() {
    log "Loading driver stack..."
    "$OPEN_IOV_ROOT/scripts/stack_manager.sh" load-stack
}

stage_mdev() {
    log "Creating mdev device..."
    "$OPEN_IOV_ROOT/scripts/stack_manager.sh" mdev-create
}

# ----------------------------------------------------------------------------
main() {
    local stages=()
    local do_load=1
    for a in "$@"; do
        case "$a" in
            --no-load) do_load=0 ;;
            -h|--help)
                sed -n '2,54p' "$0"; exit 0 ;;
            *) stages+=("$a") ;;
        esac
    done
    if [[ ${#stages[@]} -eq 0 ]]; then
        stages=(deps nvidia firmware build system)
        [[ $do_load -eq 1 ]] && stages+=(load mdev)
    fi

    log "Open-IOV install - stages: ${stages[*]}"
    log "repo root: $OPEN_IOV_ROOT"
    local s
    for s in "${stages[@]}"; do
        case "$s" in
            deps|nvidia|firmware|build|system|load|mdev) "stage_$s" ;;
            *) die "unknown stage: $s (valid: deps nvidia firmware build system load mdev)" ;;
        esac
    done

    echo
    ok "=== Open-IOV install complete ==="
    echo "Next: attach the vGPU to a VM with the wrapper of your choice:"
    echo "  docs     : docs/02-getting-started/01-install-and-first-run.md"
    echo "  raw QEMU : scripts/launch_qemu_vgpu.sh"
    echo "  libvirt  : <hostdev mode='subsystem' type='mdev' ...> with the printed UUID"
    echo "  status   : scripts/stack_manager.sh status"
    echo
    echo "Note: a future release will also ship this stack as a DKMS package"
    echo "      (open-iov-dkms) for installs without a source checkout."
}

main "$@"