param(
    [Parameter(Mandatory=$true)]
    [string]$VsConfigPath = "C:\TEMP\.vsconfig",
    
    [Parameter(Mandatory=$false)]
    [string]$LogsPath = "C:\vslogs.zip",
    
    [Parameter(Mandatory=$true)]
    [string]$InstallPath = "C:\BuildTools",
    
    [Parameter(Mandatory=$true)]
    [string]$BuildToolsURI
)

$ErrorActionPreference = 'Stop'

Write-Host "Installing Visual Studio Build Tools..." -ForegroundColor Green

$InstallerPath = "C:\TEMP\vs_buildtools.exe"
$vsCachePath = "C:\vs_cache"

# We _have to_ remove the cache (it takes too much space in the image),
# so the commands below are supposed to fail on errors like "file is being used by another process".
# BUT we have to leave SOME directories, otherwise `vswhere` won't be able
# to find the BuildTools installation, see https://developercommunity.visualstudio.com/t/bug-remove-download-cache-after-installed/362186
# So it's technically not just a cache, but Microsoft be Microsoft...
function Clean-Cache {
    Write-Host "Cleaning up Visual Studio cache at $vsCachePath..." -ForegroundColor Yellow
    try {
        # Remove everything except the '_Instances' directory
        Get-ChildItem -Path $vsCachePath | Where-Object { $_.Name -ne '_Instances' } | ForEach-Object {
            Remove-Item $_.FullName -Recurse -Force -ErrorAction Stop
        }
        Write-Host "Cache cleaned up successfully." -ForegroundColor Green
    } catch {
        Write-Warning "Failed to clean up cache: $_"
    }
}

function Install-VsWhere {
    $vsWhereUrl = "https://github.com/microsoft/vswhere/releases/download/3.1.7/vswhere.exe"
    $vsWherePath = "$InstallPath\vswhere.exe"

    Write-Host "Downloading vswhere from $vsWhereUrl..." -ForegroundColor Yellow
    Invoke-WebRequest -Uri $vsWhereUrl -OutFile $vsWherePath -UseBasicParsing
    Write-Host "vswhere downloaded to $vsWherePath" -ForegroundColor Green

    [Environment]::SetEnvironmentVariable('PATH', "$installPath;" + `
        [Environment]::GetEnvironmentVariable('PATH', 'Machine'), 'Machine')
}

try {
    # Download Build Tools installer
    Write-Host "Downloading Visual Studio Build Tools from $BuildToolsURI..." -ForegroundColor Yellow
    Invoke-WebRequest -Uri $BuildToolsURI -OutFile $InstallerPath -UseBasicParsing
    Write-Host "Download completed." -ForegroundColor Green

    # Run Visual Studio Build Tools installer
    Write-Host "Running Visual Studio Build Tools installer..." -ForegroundColor Yellow
    # Don't use `--nocache` - installation takes too long
    $process = Start-Process -FilePath $InstallerPath -ArgumentList @(
        "--quiet",
        "--wait", 
        "--norestart",
        "--path", "cache=$vsCachePath",
        "--installPath", $InstallPath,
        "--config", $VsConfigPath
    ) -Wait -PassThru -NoNewWindow
    
    $exitCode = $process.ExitCode
    
    # 3010 means "restart required" - treat as success
    if ($exitCode -eq 3010 -or $exitCode -eq 0) {
        Write-Host "Build Tools installed successfully" -ForegroundColor Green

        # vswhere.exe is usually installed in "C:\Program Files (x86)\Microsoft Visual Studio\Installer",
        # but we don't want to keep the whole installed in the image, so install it separately
        Install-VsWhere
        
        # Cleanup
        Remove-Item $InstallerPath -Force -ErrorAction Stop
        Remove-Item "C:\Program Files (x86)\Microsoft Visual Studio\Installer" -Recurse -Force -ErrorAction Stop
        Clean-Cache

        exit 0
    } else {
        # Installation failed, and we don't care about the image size - it should not be used
        Write-Warning "Build Tools installation failed with exit code: $exitCode"
        
        # Try to collect diagnostic logs
        Write-Host "Downloading collect.exe for diagnostic logs..." -ForegroundColor Yellow
        try {
            Invoke-WebRequest -Uri "https://aka.ms/vscollect.exe" -OutFile "C:\TEMP\collect.exe" -UseBasicParsing
            Write-Host "Running collect.exe..." -ForegroundColor Yellow
            Start-Process -FilePath "C:\TEMP\collect.exe" -ArgumentList @("-zip:$LogsPath") -Wait
            Write-Host "Installation logs saved to $LogsPath" -ForegroundColor Yellow
        } catch {
            Write-Warning "Failed to download or run collect.exe: $_"
        }

        throw "Installation failed with exit code: $exitCode"
    }
} catch {
    Write-Error "Failed to run Build Tools installer: $_"
   
    exit 1
}