#!/usr/bin/env pwsh

param (
  [switch]$DisableVisualStudioFeatures = $false,
  [switch]$DisableSilentMode = $false
)

Import-Module -Name $PSScriptRoot/utils.psm1 -Force

$url = "https://developer.download.nvidia.com/compute/cuda/${cuda_version_full}/network_installers/cuda_${cuda_version_full}_windows_network.exe"

$CudaFeatures = " nvcc_${cuda_version_short} " + `
  " cublas_${cuda_version_short} " + `
  " cublas_dev_${cuda_version_short} " + `
  " cuda_profiler_api_${cuda_version_short} " + `
  " cudart_${cuda_version_short} " + `
  " cufft_${cuda_version_short} " + `
  " cufft_dev_${cuda_version_short} " + `
  " cuobjdump_${cuda_version_short} " + `
  " cupti_${cuda_version_short} " + `
  " curand_${cuda_version_short} " + `
  " curand_dev_${cuda_version_short} " + `
  " cusolver_${cuda_version_short} " + `
  " cusolver_dev_${cuda_version_short} " + `
  " cusparse_${cuda_version_short} " + `
  " cusparse_dev_${cuda_version_short} " + `
  " cuxxfilt_${cuda_version_short} " + `
  " npp_${cuda_version_short} " + `
  " npp_dev_${cuda_version_short} " + `
  " nsight_compute_${cuda_version_short} " + `
  " nsight_systems_${cuda_version_short} " + `
  " nsight_vse_${cuda_version_short} " + `
  " nvdisasm_${cuda_version_short} " + `
  " nvjitlink_${cuda_version_short} " + `
  " nvjpeg_${cuda_version_short} " + `
  " nvjpeg_dev_${cuda_version_short} " + `
  " nvml_dev_${cuda_version_short} " + `
  " nvprof_${cuda_version_short} " + `
  " nvprune_${cuda_version_short} " + `
  " nvrtc_${cuda_version_short} " + `
  " nvrtc_dev_${cuda_version_short} " + `
  " nvtx_${cuda_version_short} " + `
  " occupancy_calculator_${cuda_version_short} " + `
  " opencl_${cuda_version_short} " + `
  " sanitizer_${cuda_version_short} " + `
  " thrust_${cuda_version_short} " + `
  " visual_profiler_${cuda_version_short} "

if (-Not $DisableVisualStudioFeatures) {
  $CudaFeatures = $CudaFeatures + "visual_studio_integration_${cuda_version_short} visual_profiler_${cuda_version_short}  "
}

if ($DisableSilentMode) {
  $SilentFlag = ' '
}
else {
  $SilentFlag = '-s '
}

try {
  Push-Location $PSScriptRoot
  Write-Host "Downloading CUDA from $url..."
  curl.exe -L -o cuda_${cuda_version_full}_windows_network.exe -s -S $url
  Write-Host 'Installing CUDA...'
  $proc = Start-Process -PassThru -FilePath "./cuda_${cuda_version_full}_windows_network.exe" -ArgumentList @($SilentFlag + $CudaFeatures)
  $proc.WaitForExit()
  $exitCode = $proc.ExitCode
  Pop-Location
  if ($exitCode -eq 0) {
    Write-Host 'Installation successful!'
  }
  else {
    Throw "Installation failed! Exited with $exitCode."
  }
}
catch {
  Throw "Failed to install CUDA! $($_.Exception.Message)"
}