#!/bin/bash
# Test script for CUDA PTX version fix
# Usage: ./scripts/test_cuda_fix.sh

set -e

echo "================================"
echo "CUDA PTX Version Fix - Test Script"
echo "================================"
echo ""

# Function to test a fix
test_fix() {
    local fix_name="$1"
    local setup_cmd="$2"
    local test_binary="./target/release/examples/hyperopt_mamba2_demo"
    
    echo ">>> Testing Fix: $fix_name"
    echo ""
    
    # Run setup
    echo "Setup:"
    eval "$setup_cmd"
    echo ""
    
    # Test binary
    echo "Testing binary..."
    if timeout 10s $test_binary --help 2>&1 | grep -q "hyperopt\|mamba"; then
        echo "✅ SUCCESS: $fix_name works!"
        return 0
    else
        echo "❌ FAIL: $fix_name did not work"
        return 1
    fi
}

echo "Current System State:"
echo "-------------------"
nvidia-smi | grep "Driver Version"
echo "CUDA_HOME: $CUDA_HOME"
echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
echo ""

echo "Attempting Fix Option D: CUDA Forward Compatibility Package"
echo "============================================================"
echo ""

# Check if cuda-compat package is installed
if dpkg -l | grep -q cuda-compat-12-9; then
    echo "✅ cuda-compat-12-9 is already installed"
else
    echo "⚠️  cuda-compat-12-9 is NOT installed"
    echo ""
    echo "To install, run:"
    echo "  sudo apt install cuda-compat-12-9"
    echo ""
    read -p "Do you want to install it now? (y/N) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        sudo apt install -y cuda-compat-12-9
    else
        echo "Skipping installation. Please install manually."
        exit 1
    fi
fi

echo ""
echo "Testing with CUDA compat library path..."
echo ""

# Test with compat path
export LD_LIBRARY_PATH="/usr/local/cuda-12.9/compat:$LD_LIBRARY_PATH"

if test_fix "CUDA Forward Compatibility" "export LD_LIBRARY_PATH=/usr/local/cuda-12.9/compat:\$LD_LIBRARY_PATH"; then
    echo ""
    echo "================================"
    echo "✅ FIX SUCCESSFUL!"
    echo "================================"
    echo ""
    echo "Add this to your build scripts:"
    echo "  export LD_LIBRARY_PATH=/usr/local/cuda-12.9/compat:\$LD_LIBRARY_PATH"
    echo ""
    echo "Add this to your Dockerfile:"
    echo "  ENV LD_LIBRARY_PATH=/usr/local/cuda-12.9/compat:\$LD_LIBRARY_PATH"
    echo ""
    exit 0
else
    echo ""
    echo "================================"
    echo "❌ Fix Option D failed"
    echo "================================"
    echo ""
    echo "Next steps:"
    echo "1. Try Option A: Downgrade driver to 575.x"
    echo "   sudo ubuntu-drivers install nvidia:575"
    echo "   sudo reboot"
    echo ""
    echo "2. See full analysis in:"
    echo "   CUDA_PTX_VERSION_DEEP_INVESTIGATION.md"
    echo ""
    exit 1
fi