Write-Output "Installing NVIDIA GPU driver V100..."
$ProgressPreference = 'SilentlyContinue'
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12

$gpu_driver = "RTXPRO6000.exe"

$gpu_driver_download_url = "https://download.cloudclusters.io/$gpu_driver"
Write-Output("gpu_driver_download_url is: $gpu_driver_download_url")

$download_dest_path = "C:\Users\Administrator\Downloads\nvidia.exe"
Invoke-WebRequest -Uri $gpu_driver_download_url -OutFile $download_dest_path

Start-Sleep 30

trap {
    Write-Error "An unexpected error occurred: $_"
    continue
}

try {
    $process = Start-Process -FilePath $download_dest_path -ArgumentList '/s /noreboot' -PassThru
    $process.ProcessorAffinity = 0x3
    $process.WaitForExit(1200 * 1000)

    if ($process.HasExited) {
        Write-Host "Process exited normally."
    } else {
        Write-Host "Process did not exit within the timeout period. Stopping the process."
        $process.Kill()
    }
} catch {
    Write-Error "error"
    continue
}

$install_successful = $false
$start_time = Get-Date

$possiblePaths = @(
    "nvidia-smi.exe",
    "C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe",
    "C:\Windows\System32\nvidia-smi.exe"
)
Write-Output "Starting GPU driver detection..."
while ((Get-Date) -lt ($start_time.AddMinutes(5))) {
    foreach ($path in $possiblePaths) {
        try {
            Write-Output "Attempting to execute $path ..."
            & $path
            $install_successful = $true
            Write-Output "install-finished"
            break
        } catch {
            Write-Output "$path not found or failed to execute."
        }
    }
    if ($install_successful) {
        break
    }
    Write-Output "Searching the system disk for nvidia-smi.exe..."
    $searchPath = Get-ChildItem -Path "C:\" -Recurse -Filter "nvidia-smi.exe" -ErrorAction SilentlyContinue | Select-Object -First 1
    if ($searchPath) {
        try {
            Write-Output "Found nvidia-smi at $($searchPath.FullName). Attempting to execute..."
            & $searchPath.FullName
            $install_successful = $true
            Write-Output "install-finished"
            break
        } catch {
            Write-Output "Failed to execute $($searchPath.FullName)."
        }
    } else {
        Write-Output "nvidia-smi.exe not found on the system disk."
    }
    if (-not $install_successful) {
        Write-Output "Retrying in 60 seconds..."
        Start-Sleep -Seconds 60
    }
}

# If not successful, attempt to install the driver
if (-not $install_successful) {
    Write-Output "Driver detection failed. Attempting to install the driver..."

    $driverInstallerPath = "C:\Users\Administrator\Downloads\nvidia.exe"

    if (Test-Path $driverInstallerPath) {
        try {
            Write-Output "Executing driver installer at $driverInstallerPath..."
            Start-Process -FilePath $driverInstallerPath -ArgumentList "/silent /noreboot" -Wait
            Write-Output "Driver installer executed. Retrying nvidia-smi detection..."

            # Retry nvidia-smi detection after installation
            foreach ($path in $possiblePaths) {
                try {
                    Write-Output "Attempting to execute $path ..."
                    & $path
                    $install_successful = $true
                    Write-Output "install-finished"
                    break
                } catch {
                    Write-Output "$path not found or failed to execute."
                }
            }
        } catch {
            Write-Output "Driver installation failed."
        }
    } else {
        Write-Output "Driver installer not found at $driverInstallerPath. Please check the path."
    }
}

# Final status
if (-not $install_successful) {
    Write-Output "Driver installation and detection failed."
    return
} else {
    Write-Output "Driver installed and detected successfully."
}

Write-Output "install-finished"
