#Rollback Edition of SetupDiag Script

try
{
    $tsenv = New-Object -COMObject Microsoft.SMS.TSEnvironment
    $tsBuild = $tsenv.Value("SMSTS_Build")
    $WaaSRegistryPath = "HKLM:\$($tsenv.Value("RegistryPath"))\$tsBuild"
    write-output "WaaS Registry Path = $WaaSRegistryPath"
}
catch
{
                Write-Verbose "Not running in a task sequence."
    $tsbuild = Get-ItemPropertyValue -Path "HKLM:\SOFTWARE\WaaS" -Name "LastOSUpgradeTo"
    $WaaSRegistryPath = "HKLM:\SOFTWARE\WaaS\$tsbuild"
}


function Confirm-TSProgressUISetup()
{
    if ($Script:TaskSequenceProgressUi -eq $null)
    {
        try
        {
            $Script:TaskSequenceProgressUi = New-Object -ComObject Microsoft.SMS.TSProgressUI
        }
        catch
        {
            throw "Unable to connect to the Task Sequence Progress UI! Please verify you are in a running Task Sequence Environment. Please note: TSProgressUI cannot be loaded during a prestart command.`n`nErrorDetails:`n$_"
        }
    }
}

function Confirm-TSEnvironmentSetup()
{

    if ($Script:TaskSequenceEnvironment -eq $null)
    {
        try
        {
            $Script:TaskSequenceEnvironment = New-Object -ComObject Microsoft.SMS.TSEnvironment
        }
        catch
        {
            throw "Unable to connect to the Task Sequence Environment! Please verify you are in a running Task Sequence Environment.`n`nErrorDetails:`n$_"
        }
    }
}
function Show-TSActionProgress()
{

    param(
        [Parameter(Mandatory=$true)]
        [string] $Message,
        [Parameter(Mandatory=$true)]
        [long] $Step,
        [Parameter(Mandatory=$true)]
        [long] $MaxStep
    )

    Confirm-TSProgressUISetup
    Confirm-TSEnvironmentSetup

    $Script:TaskSequenceProgressUi.ShowActionProgress(`
        $Script:TaskSequenceEnvironment.Value("_SMSTSOrgName"),`
        $Script:TaskSequenceEnvironment.Value("_SMSTSPackageName"),`
        $Script:TaskSequenceEnvironment.Value("_SMSTSCustomProgressDialogMessage"),`
        $Script:TaskSequenceEnvironment.Value("_SMSTSCurrentActionName"),`
        [Convert]::ToUInt32($Script:TaskSequenceEnvironment.Value("_SMSTSNextInstructionPointer")),`
        [Convert]::ToUInt32($Script:TaskSequenceEnvironment.Value("_SMSTSInstructionTableSize")),`
        $Message,`
        $Step,`
        $MaxStep)
}

#Create Staging Area
#$Current10Build = (Get-WmiObject Win32_OperatingSystem).BuildNumber #Not needed in Rollback
$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"
    }


#Copy SetupDiag 
if (Test-path -Path "C:\programdata\waas\SetupDiag.exe")
    {
    if (Test-Path -Path $SetupDiagStaging\SetupDiag.exe) {Remove-Item $SetupDiagStaging\SetupDiag.exe -Force}
    Copy-Item "C:\programdata\waas\SetupDiag.exe" $SetupDiagStaging\SetupDiag.exe
    }
else #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 -UseBasicParsing
    }


#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\Rollback\$tsBuild"
    $RegPathPS = "HKLM:\SOFTWARE\WaaS\Rollback\$tsBuild"

    if (Test-Path $RegPathPS){Remove-Item -Path $RegPathPS -Recurse -Force -ErrorAction SilentlyContinue}

    $Arg = "/RegPath:$($RegPath) /Output:$($LocalWaaSPath)\RollBackResults.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
            
            $Step += $Seconds
            Show-TSActionProgress -Message "Running Setup Diag" -Step $Step -MaxStep 400 
            Start-Sleep -Seconds $Seconds
            }
        Else
            {
            Write-Output "  Setup Diag Finished"
            $RuningSetupDiag = $false
            Show-TSActionProgress -Message "Finished Setup Diag" -Step 400 -MaxStep 400 
            }
        } 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 $RegPathPS)
        {
        $RegPathPSItem = get-item -Path $RegPathPS
        $SetupElapsedTime = $RegPathPSItem.GetValue("UpgradeElapsedTime")
        [int]$SetupElapsedTimeHours = $SetupElapsedTime.Split(":")[0]
        [int]$SetupElapsedTimeMinutes = $SetupElapsedTime.Split(":")[1]
        $SetupElapsedTimeTotalMinutes = $SetupElapsedTimeHours * 60 + $SetupElapsedTimeMinutes
        $TargetOSVersion = $RegPathPSItem.GetValue("TargetOSVersion")
        $SetupDiagBuild = ($TargetOSVersion.Split(" ")[0]).split(".")[2]
        }
    }

   
    if ($SetupElapsedTimeTotalMinutes)
        {
        Write-Output "Rollback IPUSetuptime set to value $SetupElapsedTimeTotalMinutes"
        New-ItemProperty -Path $WaaSRegistryPath -Name "IPUSetuptime" -Value $SetupElapsedTimeTotalMinutes -Force
        }

    if ($RegPathPSItem)
        {
        if ($RegPathPSItem.getvalue("ProfileName") -ne "NoMatchFound")
            {
            $SetupDiagErrorCode = ((($RegPathPSItem.getvalue("FailureDetails").split(","))[0]).split("=")[1]).replace(" ","")
            $SetupDiagErrorExt = ((($RegPathPSItem.getvalue("FailureDetails").split(","))[1]).split("=")[1]).replace(" ","")
            $SetupDiagFullError = "$SetupDiagErrorCode-$SetupDiagErrorExt"
            if ($SetupDiagFullError)
                {
                Write-Output "Rollback IPUFailedStepReturnCode set to value $SetupDiagFullError"
                New-ItemProperty -Path $WaaSRegistryPath -Name "IPUFailedStepReturnCode" -Value $SetupDiagFullError -Force
                if ($TSENV)
                    {
                    $tsenv.Value("SetupEngineFailedStepReturnCode") = $SetupDiagFullError
                    }
                }
            }
        
        else
            {
            New-ItemProperty -Path $WaaSRegistryPath -Name "IPUFailedStepReturnCode" -Value "NoMatchFound" -Force
                if ($TSENV)
                {
                $tsenv.Value("SetupEngineFailedStepReturnCode") = "NoMatchFound"
                }
            }
        }
