#!/bin/bash
# benchmark_infrastructure.sh - Benchmark runner for VM vs bare metal comparison

set -euo pipefail

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

RESULTS_DIR="/tmp/gsp_emu_benchmarks/$(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_warn() { echo -e "\033[1;33m[WARN]\033[0m $1"; }

run_bare_metal() {
    local test_name=$1
    shift
    log_info "Running $test_name on bare metal..."
    local output="$RESULTS_DIR/bare_metal_${test_name}.json"
    "$@" --benchmark --output "$output" 2>&1 | tee "$RESULTS_DIR/bare_metal_${test_name}.log"
}

run_in_vm() {
    local test_name=$1
    shift
    log_info "Running $test_name in VM..."
    local output="$RESULTS_DIR/vm_${test_name}.json"
    # This would run inside the guest via guest_test_agent
    # For now, placeholder
    echo '{"test": "'$test_name'", "status": "not_implemented"}' > "$output"
}

compare_results() {
    local test_name=$1
    local bare_file="$RESULTS_DIR/bare_metal_${test_name}.json"
    local vm_file="$RESULTS_DIR/vm_${test_name}.json"
    
    if [[ -f "$bare_file" && -f "$vm_file" ]]; then
        python3 -c "
import json, sys
with open('$bare_file') as f: bare = json.load(f)
with open('$vm_file') as f: vm = json.load(f)
print(f'Test: $test_name')
for key in bare:
    if key in vm and isinstance(bare[key], (int, float)) and isinstance(vm[key], (int, float)):
        ratio = vm[key] / bare[key] * 100
        print(f'  {key}: bare={bare[key]:.2f} vm={vm[key]:.2f} ratio={ratio:.1f}%')
"
    fi
}

# CUDA benchmarks
benchmark_cuda() {
    log_info "=== CUDA Benchmarks ==="
    
    # vectorAdd
    run_bare_metal "vectorAdd" /usr/local/cuda/samples/1_Utilities/vectorAdd/vectorAdd
    run_in_vm "vectorAdd"
    compare_results "vectorAdd"
    
    # matrixMul
    run_bare_metal "matrixMul" /usr/local/cuda/samples/0_Simple/matrixMul/matrixMul
    run_in_vm "matrixMul"
    compare_results "matrixMul"
    
    # deviceQuery (bandwidth)
    run_bare_metal "deviceQuery" /usr/local/cuda/samples/1_Utilities/deviceQuery/deviceQuery
    run_in_vm "deviceQuery"
    compare_results "deviceQuery"
    
    # bandwidthTest
    run_bare_metal "bandwidthTest" /usr/local/cuda/samples/1_Utilities/bandwidthTest/bandwidthTest
    run_in_vm "bandwidthTest"
    compare_results "bandwidthTest"
}

# cuBLAS benchmarks
benchmark_cublas() {
    log_info "=== cuBLAS Benchmarks ==="
    run_bare_metal "cublas_sgemm" /usr/local/cuda/samples/7_CUDALibraries/simpleCUBLAS/simpleCUBLAS
    run_in_vm "cublas_sgemm"
    compare_results "cublas_sgemm"
}

# Graphics benchmarks
benchmark_graphics() {
    log_info "=== Graphics Benchmarks ==="
    
    # vkcube (frame time)
    run_bare_metal "vkcube" vkcube --duration 10 --benchmark
    run_in_vm "vkcube"
    compare_results "vkcube"
    
    # glmark2
    run_bare_metal "glmark2" glmark2 --benchmark
    run_in_vm "glmark2"
    compare_results "glmark2"
}

# Video benchmarks
benchmark_video() {
    log_info "=== Video Benchmarks ==="
    
    # NVENC encode
    run_bare_metal "nvenc_encode" ffmpeg -i test_1080p.yuv -c:v h264_nvenc -f null -
    run_in_vm "nvenc_encode"
    compare_results "nvenc_encode"
    
    # NVDEC decode
    run_bare_metal "nvdec_decode" ffmpeg -i test_1080p.h264 -c:v h264_nvdec -f null -
    run_in_vm "nvdec_decode"
    compare_results "nvdec_decode"
}

# PyTorch benchmarks
benchmark_pytorch() {
    log_info "=== PyTorch Benchmarks ==="
    
    cat > /tmp/bench_pytorch.py <<'EOF'
import torch, time, json, sys
device = 'cuda' if torch.cuda.is_available() else 'cpu'
results = {}

# Matmul benchmark
for size in [512, 1024, 2048, 4096]:
    x = torch.randn(size, size, device=device)
    y = torch.randn(size, size, device=device)
    torch.cuda.synchronize() if device == 'cuda' else None
    start = time.time()
    for _ in range(10):
        z = torch.mm(x, y)
    torch.cuda.synchronize() if device == 'cuda' else None
    elapsed = (time.time() - start) / 10
    gflops = 2 * size**3 / elapsed / 1e9
    results[f'matmul_{size}'] = {'time': elapsed, 'gflops': gflops}
    print(f"Matmul {size}x{size}: {elapsed:.4f}s ({gflops:.1f} GFLOPS)")

# ResNet-50 forward
import torchvision.models as models
model = models.resnet50().to(device)
model.eval()
x = torch.randn(32, 3, 224, 224).to(device)
torch.cuda.synchronize() if device == 'cuda' else None
start = time.time()
for _ in range(10):
    with torch.no_grad():
        _ = model(x)
torch.cuda.synchronize() if device == 'cuda' else None
elapsed = (time.time() - start) / 10
results['resnet50_forward'] = {'time': elapsed, 'fps': 32/elapsed}
print(f"ResNet-50 forward: {elapsed:.4f}s ({32/elapsed:.1f} FPS)")

with open(sys.argv[1], 'w') as f:
    json.dump(results, f, indent=2)
EOF
    
    run_bare_metal "pytorch" python3 /tmp/bench_pytorch.py
    run_in_vm "pytorch"
    compare_results "pytorch"
}

# Generate summary report
generate_summary() {
    log_info "Generating summary report..."
    cat > "$RESULTS_DIR/summary.md" <<EOF
# GSP Emulation Benchmark Summary
Date: $(date)
Host: $(hostname)
Kernel: $(uname -r)

## Results

| Test | Bare Metal | VM | Ratio | Status |
|------|------------|-----|-------|--------|
EOF
    
    for f in "$RESULTS_DIR"/bare_metal_*.json; do
        test_name=$(basename "$f" .json | sed 's/bare_metal_//')
        vm_f="$RESULTS_DIR/vm_${test_name}.json"
        if [[ -f "$vm_f" ]]; then
            python3 -c "
import json
with open('$f') as bf: bare = json.load(bf)
with open('$vm_f') as vf: vm = json.load(vf)
for k in bare:
    if k in vm and isinstance(bare[k], (int, float)):
        ratio = vm[k] / bare[k] * 100
        status = '✅' if ratio > 85 else '⚠️' if ratio > 70 else '❌'
        print(f'| {test_name} / {k} | {bare[k]:.2f} | {vm[k]:.2f} | {ratio:.1f}% | {status} |')
" >> "$RESULTS_DIR/summary.md"
        fi
    done
    
    log_ok "Summary: $RESULTS_DIR/summary.md"
}

main() {
    case "${1:-all}" in
        cuda) benchmark_cuda ;;
        cublas) benchmark_cublas ;;
        graphics) benchmark_graphics ;;
        video) benchmark_video ;;
        pytorch) benchmark_pytorch ;;
        all)
            benchmark_cuda
            benchmark_cublas
            benchmark_graphics
            benchmark_video
            benchmark_pytorch
            ;;
        *) echo "Usage: $0 {cuda|cublas|graphics|video|pytorch|all}"; exit 1 ;;
    esac
    generate_summary
}

main "$@"