[CmdletBinding()]
param()

function Add-TestCapability {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name,

        [Parameter(Mandatory = $true)]
        $ShellPath,

        [Parameter(Mandatory = $true)]
        [ref]$Value)

    $directory = [System.IO.Path]::Combine($ShellPath, 'Common7\IDE\CommonExtensions\Microsoft\TestWindow')
    if (!(Test-Container -LiteralPath $directory)) {
        return
    }

    [string]$file = [System.IO.Path]::Combine($directory, 'vstest.console.exe')
    if (!(Test-Leaf -LiteralPath $file)) {
        return
    }

    Write-Capability -Name $Name -Value $directory
    $Value.Value = $directory
}

function Get-VSCapabilities {
    param (
        [Parameter(Mandatory = $true)]
        [ValidateSet(15, 16, 17)]
        [int]$MajorVersion,

        [Parameter(Mandatory = $true)]
        [string]$keyName
    )
    $vs = Get-VisualStudio -MajorVersion $MajorVersion
    if ($vs -and $vs.installationPath) {
        # Add VisualStudio_$($MajorVersion).0.
        # End with "\" for consistency with old ShellFolder values.
        $shellFolder = $vs.installationPath.TrimEnd('\'[0]) + "\"
        Write-Capability -Name "VisualStudio_$($MajorVersion).0" -Value $shellFolder
        $latestVS = $shellFolder
        # Add VisualStudio_IDE_$($MajorVersion).0.
        # End with "\" for consistency with old InstallDir values.
        $installDir = ([System.IO.Path]::Combine($shellFolder, 'Common7', 'IDE')) + '\'
        if ((Test-Container -LiteralPath $installDir)) {
            Write-Capability -Name "VisualStudio_IDE_$($MajorVersion).0" -Value $installDir
            $latestIde = $installDir
        }
    
        # Add VSTest_$($MajorVersion).0.
        $testWindowDir = [System.IO.Path]::Combine($installDir, 'CommonExtensions\Microsoft\TestWindow')
        $vstestConsole = [System.IO.Path]::Combine($testWindowDir, 'vstest.console.exe')
        if ((Test-Leaf -LiteralPath $vstestConsole)) {
            Write-Capability -Name "VSTest_$($MajorVersion).0" -Value $testWindowDir
            $latestTest = $testWindowDir
        }
    }
    else {
        if ((Add-CapabilityFromRegistry -Name "VisualStudio_$($MajorVersion).0" -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName -ValueName 'ShellFolder' -Value ([ref]$latestVS))) {
            $null = Add-CapabilityFromRegistry -Name "VisualStudio_IDE_$($MajorVersion).0" -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName -ValueName 'InstallDir' -Value ([ref]$latestIde)
            Add-TestCapability -Name "VSTest_$($MajorVersion).0" -ShellPath $latestVS -Value ([ref]$latestTest)
        }
    }

    if ($latestVS) {
        Write-Capability -Name 'VisualStudio' -Value $latestVS
    }

    if ($latestIde) {
        Write-Capability -Name 'VisualStudio_IDE' -Value $latestIde
    }

    if ($latestTest) {
        Write-Capability -Name 'VSTest' -Value $latestTest
    }
}

# Define the key names.
$keyName10 = 'Software\Microsoft\VisualStudio\10.0'
$keyName11 = 'Software\Microsoft\VisualStudio\11.0'
$keyName12 = 'Software\Microsoft\VisualStudio\12.0'
$keyName14 = 'Software\Microsoft\VisualStudio\14.0'
$keyName15 = 'Software\Microsoft\VisualStudio\15.0'
$keyName16 = 'Software\Microsoft\VisualStudio\16.0'
$keyName17 = 'Software\Microsoft\VisualStudio\17.0'

# Add the capabilities.
$latestVS = $null
$latestIde = $null
$latestTest = $null
$null = Add-CapabilityFromRegistry -Name 'VisualStudio_10.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName10 -ValueName 'ShellFolder' -Value ([ref]$latestVS)
$null = Add-CapabilityFromRegistry -Name 'VisualStudio_IDE_10.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName10 -ValueName 'InstallDir' -Value ([ref]$latestIde)
$null = Add-CapabilityFromRegistry -Name 'VisualStudio_11.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName11 -ValueName 'ShellFolder' -Value ([ref]$latestVS)
$null = Add-CapabilityFromRegistry -Name 'VisualStudio_IDE_11.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName11 -ValueName 'InstallDir' -Value ([ref]$latestIde)
if ((Add-CapabilityFromRegistry -Name 'VisualStudio_12.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName12 -ValueName 'ShellFolder' -Value ([ref]$latestVS))) {
    $null = Add-CapabilityFromRegistry -Name 'VisualStudio_IDE_12.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName12 -ValueName 'InstallDir' -Value ([ref]$latestIde)
    Add-TestCapability -Name 'VSTest_12.0' -ShellPath $latestVS -Value ([ref]$latestTest)
}

if ((Add-CapabilityFromRegistry -Name 'VisualStudio_14.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName14 -ValueName 'ShellFolder' -Value ([ref]$latestVS))) {
    $null = Add-CapabilityFromRegistry -Name 'VisualStudio_IDE_14.0' -Hive 'LocalMachine' -View 'Registry32' -KeyName $keyName14 -ValueName 'InstallDir' -Value ([ref]$latestIde)
    Add-TestCapability -Name 'VSTest_14.0' -ShellPath $latestVS -Value ([ref]$latestTest)
}

Get-VSCapabilities -MajorVersion 15 -keyName $keyName15

Get-VSCapabilities -MajorVersion 16 -keyName $keyName16

Get-VSCapabilities -MajorVersion 17 -keyName $keyName17
