#!/bin/bash
# stack_manager.sh - Complete vGPU stack management (load, unload, test, debug)
# Usage: ./stack_manager.sh <command> [options]

set -euo pipefail

BASE_DIR="${OPEN_IOV_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
NVIDIA_KO="$BASE_DIR/nvidia-open-gpu-kernel-modules/kernel-open/nvidia.ko"
VFIO_KO="$BASE_DIR/kernel/open-iov-vfio.ko"
FIRMWARE_PATH="$BASE_DIR/firmware/57X-GSP/570.144/gsp-570.144.fwimage.bin"
FIRMWARE_DEST="/lib/firmware/gsp-570.144.fwimage.bin"

MDEV_SCRIPT="$BASE_DIR/scripts/create_mdev.sh"
QEMU_SCRIPT="$BASE_DIR/scripts/launch_qemu_vgpu.sh"
TEST_SCRIPT="$BASE_DIR/scripts/test_vgpu_stack.sh"

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_info() { echo -e "${color_blue}[INFO]${color_nc} $*"; }
log_ok() { echo -e "${color_green}[OK]${color_nc} $*"; }
log_warn() { echo -e "${color_yellow}[WARN]${color_nc} $*"; }
log_err() { echo -e "${color_red}[ERROR]${color_nc} $*"; }

# ============================================================
# BUILD COMMANDS (correct order matters!)
# ============================================================

cmd_build_kernel() {
    log_info "Building open-iov-vfio.ko (kernel module)..."
    cd "$BASE_DIR"
    make kernel
    log_ok "open-iov-vfio.ko built"
}

cmd_build_nvidia() {
    log_info "Building NVIDIA OpenRM with vGPU symbols (KBUILD_EXTRA_SYMBOLS)..."
    cd "$BASE_DIR"
    
    # Step 1: Build nvidia kernel module (nv-kernel.o)
    log_info "Building nvidia kernel module (nv-kernel.o)..."
    cd "$BASE_DIR/nvidia-open-gpu-kernel-modules/src/nvidia"
    make -j$(nproc)
    log_ok "nvidia kernel module built"
    
    # Step 2: Build nvidia-modeset kernel module (nv-modeset-kernel.o)
    log_info "Building nvidia-modeset kernel module (nv-modeset-kernel.o)..."
    cd "$BASE_DIR/nvidia-open-gpu-kernel-modules/src/nvidia-modeset"
    make -j$(nproc)
    log_ok "nvidia-modeset kernel module built"
    
    # Step 3: Build kernel-open modules with KBUILD_EXTRA_SYMBOLS
    log_info "Building kernel-open modules with KBUILD_EXTRA_SYMBOLS..."
    cd "$BASE_DIR/nvidia-open-gpu-kernel-modules/kernel-open"
    KBUILD_EXTRA_SYMBOLS="$BASE_DIR/kernel/Module.symvers" make -j$(nproc) modules
    log_ok "NVIDIA OpenRM built"
}

cmd_build_userspace() {
    log_info "Building userspace (libiov.so + examples)..."
    cd "$BASE_DIR"
    make userspace
    make examples
    log_ok "Userspace built"
}

cmd_build_all() {
    log_info "=== BUILD ALL (correct order) ==="
    cmd_build_kernel
    cmd_build_nvidia
    cmd_build_userspace
    log_ok "=== ALL BUILDS COMPLETE ==="
}

# ============================================================
# LOAD COMMANDS (correct order matters!)
# ============================================================

cmd_load_prereqs() {
    log_info "Loading kernel prerequisites (vfio, mdev)..."
    sudo modprobe vfio vfio_iommu_type1 vfio_pci mdev vfio_mdev 2>/dev/null || true
    lsmod | grep -E "vfio|mdev"
}

cmd_load_open_iov() {
    if [[ ! -f "$VFIO_KO" ]]; then
        log_err "Missing $VFIO_KO - run 'build' first"
        return 1
    fi
    if ! lsmod | grep -q "^open_iov_vfio "; then
        log_info "Loading open-iov-vfio.ko..."
        sudo insmod "$VFIO_KO" || { log_err "Failed to load open-iov-vfio.ko"; return 1; }
        log_ok "open-iov-vfio.ko loaded"
    else
        log_ok "open-iov-vfio.ko already loaded"
    fi
}

cmd_load_nvidia() {
    if [[ ! -f "$NVIDIA_KO" ]]; then
        log_err "Missing $NVIDIA_KO - run 'build-nvidia' first"
        return 1
    fi
    if ! lsmod | grep -q "^nvidia "; then
        log_info "Loading nvidia.ko (OpenRM)..."
        sudo insmod "$NVIDIA_KO" || { log_err "Failed to load nvidia.ko"; return 1; }
        log_ok "nvidia.ko loaded"
    else
        log_ok "nvidia.ko already loaded"
    fi
}

cmd_install_firmware() {
    if [[ ! -f "$FIRMWARE_PATH" ]]; then
        log_err "Missing firmware at $FIRMWARE_PATH"
        return 1
    fi
    if [[ ! -f "$FIRMWARE_DEST" ]] || ! cmp -s "$FIRMWARE_PATH" "$FIRMWARE_DEST"; then
        log_info "Installing firmware to $FIRMWARE_DEST..."
        sudo cp "$FIRMWARE_PATH" "$FIRMWARE_DEST"
        log_ok "Firmware installed"
    else
        log_ok "Firmware already installed"
    fi
}

cmd_load_stack() {
    log_info "=== LOAD STACK (correct order) ==="
    cmd_load_prereqs
    cmd_load_open_iov
    cmd_load_nvidia
    cmd_install_firmware
    log_ok "=== DRIVER STACK LOADED ==="
    lsmod | grep -E "nvidia|open_iov|vfio|mdev"
}

# Full cycle: build all + load stack + create mdev
cmd_up() {
    log_info "=== FULL STARTUP: build + load + mdev ==="
    cmd_build_all
    cmd_load_stack
    cmd_create_mdev
    log_ok "=== SYSTEM READY FOR QEMU ==="
    echo "Run: ./stack_manager.sh qemu"
}

# ============================================================
# UNLOAD COMMANDS
# ============================================================

cmd_unload_stack() {
    log_info "=== UNLOAD STACK ==="

    if [[ -e /sys/bus/mdev/devices/* ]]; then
        log_warn "MDEV devices exist - removing them first..."
        for mdev in /sys/bus/mdev/devices/*; do
            uuid=$(basename "$mdev")
            log_info "Removing mdev $uuid..."
            echo 1 | sudo tee "$mdev/remove" >/dev/null
        done
        sleep 2
    fi

    # CRITICAL: nvidia.ko holds a reference to open_iov_vfio, so unload nvidia FIRST
    for mod in nvidia_uvm nvidia_modeset nvidia_drm nvidia_peermem nvidia; do
        if lsmod | grep -qE "^$mod\>"; then
            log_info "Unloading $mod..."
            sudo rmmod "$mod" || { log_err "Failed to unload $mod"; return 1; }
            log_ok "$mod unloaded"
        fi
    done

    if lsmod | grep -qE "^open_iov_vfio\>"; then
        log_info "Unloading open-iov-vfio.ko..."
        sudo rmmod open_iov_vfio || { log_err "Failed to unload open-iov-vfio.ko"; return 1; }
        log_ok "open-iov-vfio.ko unloaded"
    fi

    sudo modprobe -r vfio_mdev vfio_pci vfio_iommu_type1 vfio mdev 2>/dev/null || true
    log_ok "=== DRIVER STACK UNLOADED ==="
}

# ============================================================
# MDEV COMMANDS
# ============================================================

cmd_create_mdev() {
    log_info "Creating mdev device..."
    if [[ ! -f "$MDEV_SCRIPT" ]]; then
        log_err "create_mdev.sh not found"
        return 1
    fi
    "$MDEV_SCRIPT" "$@"
}

cmd_list_mdev() {
    log_info "MDEV devices:"
    ls -la /sys/bus/mdev/devices/ 2>/dev/null || log_warn "No mdev devices"
}

# ============================================================
# QEMU COMMANDS
# ============================================================

cmd_launch_qemu() {
    log_info "Launching QEMU with vGPU..."
    if [[ ! -f "$QEMU_SCRIPT" ]]; then
        log_err "launch_qemu_vgpu.sh not found"
        return 1
    fi
    "$QEMU_SCRIPT" "$@"
}

# ============================================================
# TEST COMMANDS
# ============================================================

cmd_test_stack() {
    log_info "Running comprehensive test suite..."
    if [[ ! -f "$TEST_SCRIPT" ]]; then
        log_err "test_vgpu_stack.sh not found"
        return 1
    fi
    "$TEST_SCRIPT" "$@"
}

# ============================================================
# STATUS / DEBUG
# ============================================================

cmd_status() {
    echo "=== Kernel Modules ==="
    lsmod | grep -E "nvidia|open_iov|vfio|mdev" || echo "  (none)"
    echo ""
    echo "=== MDEV Devices ==="
    ls -la /sys/bus/mdev/devices/ 2>/dev/null || echo "  (none)"
    echo ""
    echo "=== VFIO Devices ==="
    ls -la /dev/vfio/ 2>/dev/null || echo "  (none)"
    echo ""
    echo "=== Firmware ==="
    ls -la "$FIRMWARE_DEST" 2>/dev/null || echo "  Not installed"
    echo ""
    echo "=== dmesg (last 20) ==="
    sudo dmesg -T | grep -i -e nvidia -e gsp -e vgpu -e open-iov -e mdev | tail -20
}

cmd_dmesg() {
    sudo dmesg -T -w | grep -i -e nvidia -e gsp -e vgpu -e open-iov -e mdev
}

cmd_debug_qemu() {
    log_info "Starting debug QEMU session..."
    if [[ ! -f "$BASE_DIR/scripts/debug_guest_init.sh" ]]; then
        log_err "debug_guest_init.sh not found"
        return 1
    fi
    "$BASE_DIR/scripts/debug_guest_init.sh" "$@"
}

cmd_cleanup_crash() {
    log_info "Cleaning up after crash..."
    cmd_unload_stack
    log_info "Clearing any stuck mdev devices..."
    for mdev in /sys/bus/mdev/devices/*; do
        [[ -e "$mdev" ]] || continue
        uuid=$(basename "$mdev")
        echo 1 | sudo tee "$mdev/remove" >/dev/null 2>&1 || true
    done
    log_ok "Cleanup complete"
}

# ============================================================
# HELP
# ============================================================

cmd_help() {
    cat <<EOF
stack_manager.sh - Complete vGPU Stack Management

USAGE:
    ./stack_manager.sh <command> [options]

BUILD COMMANDS (run in order):
    build-kernel      Build open-iov-vfio.ko (generates Module.symvers)
    build-nvidia      Build NVIDIA OpenRM with KBUILD_EXTRA_SYMBOLS
    build-userspace   Build libiov.so + examples
    build-all         Run all three build steps in correct order

LOAD COMMANDS (run in order):
    load-prereqs      Load vfio, vfio_iommu_type1, vfio_pci, mdev, vfio_mdev
    load-open-iov     Load open-iov-vfio.ko
    load-nvidia       Load nvidia.ko (OpenRM)
    install-firmware  Copy GSP firmware to /lib/firmware/
    load-stack        Run all four load steps in correct order

FULL CYCLE:
    up                Build all + load stack + create mdev (one command!)

MDEV COMMANDS:
    mdev-create       Create mdev device
    mdev-list         List mdev devices

QEMU COMMANDS:
    qemu              Launch QEMU with vGPU
    qemu-debug        Launch QEMU with GDB + guest monitor

TEST COMMANDS:
    test              Run comprehensive test suite

STATUS / DEBUG:
    status            Show current stack status
    dmesg             Follow filtered dmesg output
    cleanup-crash     Emergency cleanup after crash/hang
    help              Show this help

EXAMPLES:
    # Full startup from clean state:
    ./stack_manager.sh up

    # Step by step:
    ./stack_manager.sh build-all
    ./stack_manager.sh load-stack
    ./stack_manager.sh mdev-create
    ./stack_manager.sh qemu

    # Or just load (if already built):
    ./stack_manager.sh load-stack
    ./stack_manager.sh mdev-create
    ./stack_manager.sh qemu

    # Testing:
    ./stack_manager.sh test

    # Cleanup:
    ./stack_manager.sh unload
    # Or force cleanup after crash:
    ./stack_manager.sh cleanup-crash

NOTE:
    - Requires root for module operations
    - BUILD ORDER: open-iov-vfio.ko FIRST (generates Module.symvers),
                   THEN nvidia.ko with KBUILD_EXTRA_SYMBOLS
    - LOAD ORDER: vfio/mdev prerequisites -> open-iov-vfio.ko -> nvidia.ko
    - UNLOAD ORDER: remove mdevs -> open-iov-vfio.ko -> nvidia.ko
    - NEVER REBOOT without explicit approval
EOF
}

main() {
    case "${1:-help}" in
        # Build
        build-kernel) cmd_build_kernel ;;
        build-nvidia) cmd_build_nvidia ;;
        build-userspace) cmd_build_userspace ;;
        build-all) cmd_build_all ;;

        # Load
        load-prereqs) cmd_load_prereqs ;;
        load-open-iov) cmd_load_open_iov ;;
        load-nvidia) cmd_load_nvidia ;;
        install-firmware) cmd_install_firmware ;;
        load-stack) cmd_load_stack ;;

        # Full cycle
        up) cmd_up ;;

        # Unload
        unload) cmd_unload_stack ;;

        # MDEV
        mdev-create) shift; cmd_create_mdev "$@" ;;
        mdev-list) cmd_list_mdev ;;

        # QEMU
        qemu) shift; cmd_launch_qemu "$@" ;;
        qemu-debug) shift; cmd_debug_qemu "$@" ;;

        # Test
        test) shift; cmd_test_stack "$@" ;;

        # Status/Debug
        status) cmd_status ;;
        dmesg) cmd_dmesg ;;
        cleanup-crash) cmd_cleanup_crash ;;
        help|*) cmd_help ;;
    esac
}

main "$@"