﻿
Param( 
    [parameter(Mandatory=$false)][string]$AxeCoreNet,
    [parameter(Mandatory=$false)][string]$ResultFile,
    [parameter(Mandatory=$false)][string]$TempFile,
    [parameter(Mandatory=$false)][int32]$RestartCount
)


# This try/finally block wraps the whole script so we can always properly dispose of objects.
#
try {


    # Need to know where our script is so we can use full path to all our dependent files.
    #
    $folderScript = Split-Path -Path $MyInvocation.MyCommand.Path -Parent

    # Try to load up Axe because it is cool.
    #
    $boolRunningUnderAxe = $false
    if ( $AxeCoreNet ) {
        $axeAssembly = [Reflection.Assembly]::LoadFrom( $AxeCoreNet )
        if ( $axeAssembly ) {
            $axeSupport = [Microsoft.Assessments.Runtime.Support]::Initialize()
            if ( $axeSupport ) {

                $boolRunningUnderAxe = $axeSupport.EngineRunning

                $axeLogger = $axeSupport.CreateLogger()
            }
        }
    }

    # This function will write what is piped to it as a string to the default output
    # (usually the host/console) and also to the log file if running under axe.
    #
    function Log-Info {
        PROCESS {
            if ( $_ ) {
                "$_" | Out-Default
                if ( $axeLogger ) {
                    $axeLogger.LogMessage( "$_" )
                }
           }
        }
    }

    # This function will write what is piped to it as a string as an error
    # and also to the log file if running under axe.
    #
    function Log-Error([int32]$ErrorCode=1) {
        PROCESS {
            if ( $_ ) {
                "$_" | Write-Error
                if ( $axeLogger ) {
                    $axeLogger.LogError( $ErrorCode, "$_" )
                }
           }
        }
    }


    #
    # This is where the real 'meat' of the assessment begins:
    #
    
    "Restart Count=$RestartCount" | Log-Info

    # We can also only do restart if we are running in axe because it is what will launch us again.
    #
    if ( -not $boolRunningUnderAxe ) {
        "This script only works when run from the Axe framework." | Log-Error
        Exit 1
    }

    # We only have to do anything special if this is the first time the script was launched.
    #
    if ( $RestartCount -eq 0 ) {
        if ( $TempFile ) {
            (Get-Date).ToString("o") | Out-File "$TempFile"
        }
        $axeSupport.ReportProgress( [Microsoft.Assessments.AxeProgressType]::OnOff, 0, "Running restart assessment." )
        Restart-Computer -Force
        Exit 0
    }

    # Otherwise just calculate our one metric and are done otherwise.
    #
    if ( $TempFile ) {
        $timespanRestart = New-TimeSpan (Get-Content "$TempFile") (Get-Date)
        "Restart took a total of $($timespanRestart.TotalSeconds) seconds." | Log-Info
        if ( $ResultFile ) {

            $xmlOut = [System.Xml.XmlWriter]::Create( $ResultFile )
            $xmlOut.WriteStartElement("AssessmentResult")
            $xmlOut.WriteStartElement("Iterations")
            $xmlOut.WriteStartElement("Iteration")
            $xmlOut.WriteStartElement("MetricValues")
            $xmlOut.WriteStartElement("MetricValue")
            $xmlOut.WriteElementString("ProgrammaticName", "DurationSeconds")
            $xmlOut.WriteElementString("Value", $timespanRestart.TotalSeconds.ToString())
            $xmlOut.Close()

        }
    }


# This try/finally block wraps the whole script so we can always properly dispose of objects.
#
} finally {

    # We should always properly dispose of this guy.
    #
    if ( $axeSupport ) {
        $axeSupport.Dispose()
    }

}
