#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# GSP vGPU-IO Integration Test Runner

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
KERNEL_TEST_DIR="$PROJECT_ROOT/kernel/tests"
DRIVER_DIR="$PROJECT_ROOT/driver"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        exit 1
    fi
}

# Build kernel test module
build_test_module() {
    log_info "Building kernel test module..."
    cd "$KERNEL_TEST_DIR"
    make clean
    make
    if [[ $? -ne 0 ]]; then
        log_error "Failed to build test module"
        exit 1
    fi
    log_info "Test module built successfully"
}

# Load test module and capture output
run_test_module() {
    log_info "Loading test module..."
    cd "$KERNEL_TEST_DIR"
    
    # Clear dmesg
    dmesg -C
    
    # Load module
    insmod gsp_vgpu_comprehensive_test.ko 2>&1 || {
        log_error "Failed to load test module"
        dmesg -T | tail -30
        return 1
    }
    
    # Wait for tests to complete
    sleep 2
    
    # Get test results from dmesg
    dmesg -T | grep -E "(PASS|FAIL|Test Results|TEST SUITE)" | tail -100
    
    # Check if tests passed
    if dmesg | grep -q "TEST SUITE PASSED"; then
        log_info "All tests PASSED!"
        return 0
    elif dmesg | grep -q "TEST SUITE FAILED"; then
        log_error "Tests FAILED!"
        return 1
    else
        log_warn "Could not determine test result"
        return 1
    fi
}

# Unload test module
unload_test_module() {
    log_info "Unloading test module..."
    rmmod gsp_vgpu_comprehensive_test 2>/dev/null || true
    sleep 1
}

# Build driver
build_driver() {
    log_info "Building vGPU-IO driver..."
    cd "$DRIVER_DIR"
    make clean
    make
    if [[ $? -ne 0 ]]; then
        log_error "Failed to build driver"
        exit 1
    fi
    log_info "Driver built successfully"
}

# Load driver
load_driver() {
    log_info "Loading vGPU-IO driver..."
    cd "$DRIVER_DIR"
    insmod gsp-vgpu.ko firmware_name=gsp-570.144.fwimage.bin 2>&1 || {
        log_error "Failed to load driver"
        dmesg -T | tail -30
        return 1
    }
    
    # Check if device created
    sleep 2
    if ls /dev/gsp-vgpu* 2>/dev/null; then
        log_info "Driver loaded, device created"
        return 0
    else
        log_warn "Driver loaded but no device node found"
        dmesg -T | tail -20
    fi
}

# Unload driver
unload_driver() {
    log_info "Unloading vGPU-IO driver..."
    rmmod gsp-vgpu 2>/dev/null || true
    sleep 1
}

# Run VF creation test
test_vf_creation() {
    log_info "Testing VF creation via debugfs..."
    
    # Check debugfs
    if [[ -d /sys/kernel/debug/gsp-vgpu ]]; then
        log_info "Debugfs interface available"
        cat /sys/kernel/debug/gsp-vgpu/vfs
    else
        log_warn "Debugfs not mounted or driver not loaded"
        return 1
    fi
}

# Run RPC test via debugfs
test_rpc_stats() {
    log_info "Checking RPC statistics..."
    if [[ -f /sys/kernel/debug/gsp-vgpu/rpc ]]; then
        cat /sys/kernel/debug/gsp-vgpu/rpc
    fi
}

# Main test runner
main() {
    local test_type="${1:-all}"
    
    check_root
    
    case "$test_type" in
        "unit")
            build_test_module
            run_test_module
            unload_test_module
            ;;
        "driver")
            build_driver
            load_driver
            test_vf_creation
            test_rpc_stats
            unload_driver
            ;;
        "all"|*)
            log_info "Running ALL tests..."
            build_test_module
            run_test_module
            unload_test_module
            
            build_driver
            load_driver
            test_vf_creation
            test_rpc_stats
            unload_driver
            ;;
    esac
    
    log_info "Test run completed"
}

main "$@"