[CmdletBinding()]
param()

$script:capabilityName = "AzurePS"

function Get-FromModulePath {
    [CmdletBinding()]
    param([string]$ModuleName)

     # Valid ModuleName values are Az.Accounts, AzureRM and Azure
     # We are looking for Az.Accounts module because "Get-Module -Name Az" is not working due to a known PowerShell bug.
    if (($ModuleName -ne "Az.Accounts") -and ($ModuleName -ne "AzureRM") -and ($ModuleName -ne "Azure")) {
        Write-Host "Attempting to find invalid module."
        return $false
    }

    # Attempt to resolve the module.
    Write-Host "Attempting to find the module '$ModuleName' from the module path."
    $module = Get-Module -Name $ModuleName -ListAvailable | Select-Object -First 1
    if (!$module) {
        Write-Host "Not found."
        return $false
    }

    if ($ModuleName -eq "AzureRM") {
        # For AzureRM, validate the AzureRM.profile module can be found as well.
        $profileName = "AzureRM.profile"
        Write-Host "Attempting to find the module $profileName"
        $profileModule = Get-Module -Name $profileName -ListAvailable | Select-Object -First 1
        if (!$profileModule) {
            Write-Host "Not found."
            return $false
        }
    }

    # Add the capability.
    Write-Capability -Name $script:capabilityName -Value $module.Version
    return $true
}

function Get-FromSdkPath {
    [CmdletBinding()]
    param([switch]$Classic)

    if ($Classic) {
        $partialPath = 'Microsoft SDKs\Azure\PowerShell\ServiceManagement\Azure\Azure.psd1'
    } else {
        $partialPath = 'Microsoft SDKs\Azure\PowerShell\ResourceManager\AzureResourceManager\AzureRM.Profile\AzureRM.Profile.psd1'
    }

    foreach ($programFiles in @(${env:ProgramFiles(x86)}, $env:ProgramFiles)) {
        if (!$programFiles) {
            continue
        }

        $path = [System.IO.Path]::Combine($programFiles, $partialPath)
        Write-Host "Checking if path exists: $path"
        if (Test-Path -LiteralPath $path -PathType Leaf) {
            $directory = [System.IO.Path]::GetDirectoryName($path)
            $fileNameOnly = [System.IO.Path]::GetFileNameWithoutExtension($path)

            # Prepend the module path.
            Write-Host "Temporarily adjusting module path."
            $originalPSModulePath = $env:PSModulePath
            if ($env:PSModulePath) {
                $env:PSModulePath = ";$env:PSModulePath"
            }

            $env:PSModulePath = "$directory$env:PSModulePath"
            Write-Host "Env:PSModulePath: '$env:PSModulePath'"
            try {
                # Get the module.
                Write-Host "Get-Module -Name $fileNameOnly -ListAvailable"
                $module = Get-Module -Name $fileNameOnly -ListAvailable | Select-Object -First 1
            } finally {
                # Revert the module path adjustment.
                Write-Host "Reverting module path adjustment."
                $env:PSModulePath = $originalPSModulePath
                Write-Host "Env:PSModulePath: '$env:PSModulePath'"
            }

            # Add the capability.
            Write-Capability -Name $script:capabilityName -Value $module.Version
            return $true
        }
    }

    return $false
}

Write-Host "Env:PSModulePath: '$env:PSModulePath'"
$null = (Get-FromModulePath -ModuleName:"Az.Accounts") -or
    (Get-FromModulePath -ModuleName:"AzureRM") -or
    (Get-FromSdkPath -Classic:$false) -or
    (Get-FromModulePath -ModuleName:"Azure") -or
    (Get-FromSdkPath -Classic:$true)   