#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# windows_guest_automation.sh - Windows guest test automation via WinRM/SSH

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BASE_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"

RESULTS_DIR="/tmp/gsp_emu_windows/$(date +%Y%m%d_%H%M%S)"
mkdir -p "$RESULTS_DIR"

log_info() { echo -e "\033[0;34m[INFO]\033[0m $1"; }
log_ok() { echo -e "\033[0;32m[OK]\033[0m $1"; }
log_error() { echo -e "\033[0;31m[ERROR]\033[0m $1"; }

# Configuration
WIN_USER="${WIN_USER:-Administrator}"
WIN_IP="${WIN_IP:-192.168.122.100}"
WIN_PASSWORD="${WIN_PASSWORD:-}"
WINRM_PORT="${WINRM_PORT:-5985}"
SSH_PORT="${SSH_PORT:-22}"

check_winrm() {
    log_info "Checking WinRM connectivity to $WIN_IP..."
    if command -v winrm &>/dev/null; then
        winrm identify -r:"http://$WIN_IP:$WINRM_PORT/wsman" -u:"$WIN_USER" -p:"$WIN_PASSWORD" -auth:basic 2>/dev/null
    elif command -v python3 &>/dev/null; then
        python3 -c "
import winrm
session = winrm.Session('$WIN_IP:$WINRM_PORT', auth=('$WIN_USER', '$WIN_PASSWORD'), transport='ntlm')
result = session.run_cmd('echo test')
print('WinRM OK' if result.status_code == 0 else 'WinRM FAILED')
"
    else
        log_error "No WinRM client found. Install winrm or python3 with pywinrm"
        return 1
    fi
}

run_winrm_cmd() {
    local cmd="$1"
    local output_file="$2"
    
    if command -v winrm &>/dev/null; then
        winrm invoke -r:"http://$WIN_IP:$WINRM_PORT/wsman" -u:"$WIN_USER" -p:"$WIN_PASSWORD" \
            -auth:basic Create "http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Command" \
            "Command=cmd.exe" "Arguments=/c $cmd" > "$output_file" 2>&1
    else
        python3 -c "
import winrm, sys
session = winrm.Session('$WIN_IP:$WINRM_PORT', auth=('$WIN_USER', '$WIN_PASSWORD'), transport='ntlm')
result = session.run_cmd('$cmd')
with open('$output_file', 'w') as f:
    f.write(result.std_out.decode() + result.std_err.decode())
sys.exit(result.status_code)
"
    fi
}

run_ssh_cmd() {
    local cmd="$1"
    local output_file="$2"
    sshpass -p "$WIN_PASSWORD" ssh -o StrictHostKeyChecking=no -p "$SSH_PORT" "$WIN_USER@$WIN_IP" "$cmd" > "$output_file" 2>&1
}

test_nvidia_smi() {
    log_info "Testing nvidia-smi..."
    run_winrm_cmd "nvidia-smi" "$RESULTS_DIR/nvidia_smi.txt"
    if grep -q "NVIDIA" "$RESULTS_DIR/nvidia_smi.txt"; then
        log_ok "nvidia-smi: PASS"
        return 0
    else
        log_error "nvidia-smi: FAIL"
        cat "$RESULTS_DIR/nvidia_smi.txt"
        return 1
    fi
}

test_cuda_vectoradd() {
    log_info "Testing CUDA vectorAdd..."
    run_winrm_cmd "cd C:\\cuda_samples\\1_Utilities\\vectorAdd && .\\vectorAdd.exe" "$RESULTS_DIR/cuda_vectoradd.txt"
    if grep -q "Test passed" "$RESULTS_DIR/cuda_vectoradd.txt"; then
        log_ok "CUDA vectorAdd: PASS"
        return 0
    else
        log_error "CUDA vectorAdd: FAIL"
        cat "$RESULTS_DIR/cuda_vectoradd.txt"
        return 1
    fi
}

test_cuda_devquery() {
    log_info "Testing CUDA deviceQuery..."
    run_winrm_cmd "cd C:\\cuda_samples\\1_Utilities\\deviceQuery && .\\deviceQuery.exe" "$RESULTS_DIR/cuda_devquery.txt"
    if grep -q "Result = PASS" "$RESULTS_DIR/cuda_devquery.txt"; then
        log_ok "CUDA deviceQuery: PASS"
        return 0
    else
        log_error "CUDA deviceQuery: FAIL"
        cat "$RESULTS_DIR/cuda_devquery.txt"
        return 1
    fi
}

test_vulkan() {
    log_info "Testing Vulkan..."
    run_winrm_cmd "vulkaninfo.exe" "$RESULTS_DIR/vulkaninfo.txt"
    if grep -q "GPU id" "$RESULTS_DIR/vulkaninfo.txt"; then
        log_ok "Vulkan: PASS"
        return 0
    else
        log_error "Vulkan: FAIL"
        cat "$RESULTS_DIR/vulkaninfo.txt"
        return 1
    fi
}

test_dx12() {
    log_info "Testing DirectX 12..."
    run_winrm_cmd "dxdiag /t $RESULTS_DIR/dxdiag.txt" "$RESULTS_DIR/dxdiag_cmd.txt"
    # Parse dxdiag output
    if grep -q "Feature Level" "$RESULTS_DIR/dxdiag.txt"; then
        log_ok "DirectX 12: PASS"
        return 0
    else
        log_error "DirectX 12: FAIL"
        return 1
    fi
}

test_ffmpeg_nvenc() {
    log_info "Testing FFmpeg NVENC..."
    run_winrm_cmd "ffmpeg -encoders" "$RESULTS_DIR/ffmpeg_encoders.txt"
    if grep -q "h264_nvenc\|hevc_nvenc" "$RESULTS_DIR/ffmpeg_encoders.txt"; then
        log_ok "FFmpeg NVENC: PASS"
        return 0
    else
        log_error "FFmpeg NVENC: FAIL"
        return 1
    fi
}

test_pytorch() {
    log_info "Testing PyTorch..."
    run_winrm_cmd "python -c \"import torch; print('CUDA:', torch.cuda.is_available()); print('Device:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A')\"" "$RESULTS_DIR/pytorch.txt"
    if grep -q "CUDA: True" "$RESULTS_DIR/pytorch.txt"; then
        log_ok "PyTorch: PASS"
        return 0
    else
        log_error "PyTorch: FAIL"
        cat "$RESULTS_DIR/pytorch.txt"
        return 1
    fi
}

test_windows_driver_signature() {
    log_info "Checking driver signature status..."
    run_winrm_cmd "pnputil /enum-drivers | findstr nvidia" "$RESULTS_DIR/driver_signature.txt"
    cat "$RESULTS_DIR/driver_signature.txt"
}

collect_event_logs() {
    log_info "Collecting Windows event logs..."
    run_winrm_cmd "wevtutil qe System /q:\"*[System[Provider[@Name='NVIDIA']]]\" /f:text /c:50" "$RESULTS_DIR/nvidia_events.txt"
    run_winrm_cmd "wevtutil qe System /q:\"*[System[Provider[@Name='Display']]]\" /f:text /c:50" "$RESULTS_DIR/display_events.txt"
}

run_full_suite() {
    log_info "=== Windows Guest Test Suite ==="
    local failed=0
    
    test_nvidia_smi || ((failed++))
    test_cuda_devquery || ((failed++))
    test_cuda_vectoradd || ((failed++))
    test_vulkan || ((failed++))
    test_dx12 || ((failed++))
    test_ffmpeg_nvenc || ((failed++))
    test_pytorch || ((failed++))
    test_windows_driver_signature
    collect_event_logs
    
    log_info "=== Results ==="
    log_info "Failed tests: $failed"
    
    # Generate summary
    cat > "$RESULTS_DIR/summary.json" <<EOF
{
    "timestamp": "$(date -Iseconds)",
    "guest_ip": "$WIN_IP",
    "failed_tests": $failed,
    "tests": {
        "nvidia_smi": $(test -f "$RESULTS_DIR/nvidia_smi.txt" && grep -q "NVIDIA" "$RESULTS_DIR/nvidia_smi.txt" && echo true || echo false),
        "cuda_devquery": $(test -f "$RESULTS_DIR/cuda_devquery.txt" && grep -q "Result = PASS" "$RESULTS_DIR/cuda_devquery.txt" && echo true || echo false),
        "cuda_vectoradd": $(test -f "$RESULTS_DIR/cuda_vectoradd.txt" && grep -q "Test passed" "$RESULTS_DIR/cuda_vectoradd.txt" && echo true || echo false),
        "vulkan": $(test -f "$RESULTS_DIR/vulkaninfo.txt" && grep -q "GPU id" "$RESULTS_DIR/vulkaninfo.txt" && echo true || echo false),
        "dx12": $(test -f "$RESULTS_DIR/dxdiag.txt" && grep -q "Feature Level" "$RESULTS_DIR/dxdiag.txt" && echo true || echo false),
        "ffmpeg_nvenc": $(test -f "$RESULTS_DIR/ffmpeg_encoders.txt" && grep -q "h264_nvenc\|hevc_nvenc" "$RESULTS_DIR/ffmpeg_encoders.txt" && echo true || echo false),
        "pytorch": $(test -f "$RESULTS_DIR/pytorch.txt" && grep -q "CUDA: True" "$RESULTS_DIR/pytorch.txt" && echo true || echo false)
    }
}
EOF
    
    log_ok "Results saved to $RESULTS_DIR"
    cat "$RESULTS_DIR/summary.json"
}

main() {
    case "${1:-full}" in
        check) check_winrm ;;
        nvidia-smi) test_nvidia_smi ;;
        cuda) test_cuda_devquery && test_cuda_vectoradd ;;
        graphics) test_vulkan && test_dx12 ;;
        video) test_ffmpeg_nvenc ;;
        pytorch) test_pytorch ;;
        full) run_full_suite ;;
        *) echo "Usage: $0 {check|nvidia-smi|cuda|graphics|video|pytorch|full}"; exit 1 ;;
    esac
}

main "$@"