try
{
    $tsenv = New-Object -COMObject Microsoft.SMS.TSEnvironment
}
catch
{
	Write-Verbose "Not running in a task sequence."
    $WaaSRegistryPath = "HKLM:\SOFTWARE\WaaS\1909"
}

if ($tsenv)
    {
    $tsBuild = $tsenv.Value("SMSTS_Build")
    #SMSTS_Build is set in the Task Sequence
    $WaaSRegistryPath = "HKLM:\$($tsenv.Value("RegistryPath"))\$($tsenv.Value("SMSTS_Build"))"
    }

#Only Run if OS upgraded.
if ((Get-ItemPropertyValue 'HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion\' 'ReleaseID') -eq $tsBuild)
    {
    #Create Staging Area
    $Current10Build = (Get-WmiObject Win32_OperatingSystem).BuildNumber
    $LocalWaaSPath = "C:\ProgramData\WaaS"
    $SetupDiagStaging = "$($LocalWaaSPath)\SetupDiagStaging"
    $SetupDiagStagingLogPath = "$($SetupDiagStaging)\Logs"
    if (!(Test-Path -Path $SetupDiagStagingLogPath))
        {
        New-Item -Path $SetupDiagStagingLogPath -ItemType Directory
        Write-Output "Created Staging Folder"
        }

    #Download Latest Version of SetupDiag
    #if (Test-Path -Path $SetupDiagStaging\SetupDiag.exe) {Remove-Item $SetupDiagStaging\SetupDiag.exe -Force}
    #Invoke-WebRequest -Uri "https://go.microsoft.com/fwlink/?linkid=870142" -OutFile $SetupDiagStaging\SetupDiag.exe

    #Copy SetupDiag
    if (Test-Path -Path $SetupDiagStaging\SetupDiag.exe) {Remove-Item $SetupDiagStaging\SetupDiag.exe -Force}
    if (Test-path -Path "C:\programdata\waas\SetupDiag.exe"){Copy-Item "C:\programdata\waas\SetupDiag.exe" $SetupDiagStaging\SetupDiag.exe}


    #Confirm Download & Do Stuff if SetupDiag was there.
    if (!(Test-Path -Path $SetupDiagStaging\SetupDiag.exe))
        {Write-Output "SetupDiag.exe Not Downloaded"}
    Else
        {
        $Process = "$SetupDiagStaging\SetupDiag.exe"
        Write-Output "SetupDiag Downloaded with version: $(((get-item -Path $Process).VersionInfo).FileVersion)"
        $RegPath = "HKEY_LOCAL_MACHINE\SOFTWARE\WaaS\SetupDiag\LatestRun"
        $Arg = "/RegPath:$($RegPath) /Output:$($LocalWaaSPath)\Results.xml /Format:xml"
    
        Write-Output "Starting $Process  $Arg"
        Start-Process $Process -ArgumentList $Arg
    
        #Monitor SetupDiag
        $Seconds = 20
        Write-Output "Monitoring SetupDiag Process every $Seconds seconds"
        DO
            {
             if(Get-Process "SetupDiag" -ErrorAction SilentlyContinue)
                {
                Write-Output "    Setup Diag Running"
                $RuningSetupDiag = $true
                Start-Sleep -Seconds $Seconds
                }
            Else
                {
                Write-Output "  Setup Diag Finished"
                $RuningSetupDiag = $false
                }
            } Until ($RuningSetupDiag -eq $false)

        Write-Output "SetupDiag Process Complete"
        #Rename Setup Diag Key to Build Number for Historical Reasons.
        Write-Output "Rename Registry Export to Build Number"
        if (Test-Path -Path "HKLM:\Software\WaaS\SetupDiag\LatestRun")
            {
            $SetupDiagKey = get-item -Path "HKLM:\Software\WaaS\SetupDiag\LatestRun"
            $SetupElapsedTime = $SetupDiagKey.GetValue("UpgradeElapsedTime")
            [int]$SetupElapsedTimeHours = $SetupElapsedTime.Split(":")[0]
            [int]$SetupElapsedTimeMinutes = $SetupElapsedTime.Split(":")[1]
            $SetupElapsedTimeTotalMinutes = $SetupElapsedTimeHours * 60 + $SetupElapsedTimeMinutes
            $TargetOSVersion = $SetupDiagKey.GetValue("TargetOSVersion")
            $SetupDiagBuild = ($TargetOSVersion.Split(" ")[0]).split(".")[2]
            if (test-path "HKLM:\Software\WaaS\SetupDiag\$SetupDiagBuild"){Remove-Item -path "HKLM:\Software\WaaS\SetupDiag\$SetupDiagBuild"}
            Rename-Item -Path "HKLM:\Software\WaaS\SetupDiag\LatestRun" -NewName $SetupDiagBuild -Force
            Write-Output "Registry Path: HKLM:\Software\WaaS\SetupDiag\$SetupDiagBuild"
            }
        }

    if (test-path "C:\Windows\ccm\logs\SMSTSlog\smsts.log") 
        {
        $smstslog = get-content "C:\Windows\ccm\logs\SMSTSlog\smsts.log"
        Write-Output "Found Log here: C:\Windows\ccm\logs\SMSTSlog\smsts.log"
        }
    else 
        {
        $smstslog = get-content "C:\Windows\ccm\logs\smsts.log"
        Write-Output "Found Log here: C:\Windows\ccm\logs\smsts.log"
        }
    if ($smstslog | Where-Object -FilterScript { $_ -match "ServicePack Build :"})
        {
        Write-Output "Found Build information in SMSTS Log"
        $stringBuild = $smstslog | Where-Object -FilterScript { $_ -match "Version : 10.0."}
        $stringUBR = $smstslog | Where-Object -FilterScript { $_ -match "ServicePack Build :"}
        $SMSTSBuild = ($stringBuild.Split(".")[2]).trim()
        $SMSTSBuild
        $SMSTSUBR = ($stringUBR.Split(":")[1]).trim()
        $SMSTSUBR
        $SMSTSFullBuild = "$($SMSTSbuild).$($SMSTSUBR)"
        $SMSTSFullBuild
        }
    else {Write-Output "Did NOT find Build information in SMSTS Log"}


    if ($SMSTSFullBuild)
        {
        Write-Output "$SMSTSFullBuild Recorded in SMSTS.LOG"
        if ($Current10Build -eq $SMSTSBuild)
           {
           Write-Output "IPUBuild set to Value $SMSTSFullBuild"
           New-ItemProperty -Path $WaaSRegistryPath -Name "IPUBuild" -Value $SMSTSFullBuild -Force
           Write-Output "IPUSetuptime set to value $SetupElapsedTimeTotalMinutes"
           New-ItemProperty -Path $WaaSRegistryPath -Name "IPUSetuptime" -Value $SetupElapsedTimeTotalMinutes -Force
           }
        }
    Else {Write-Output "SMSTSFullBuild Not Created"}
    }