[CmdletBinding()]
param(
    [switch]$SkipDiskRepair=$true,
    [switch]$PostRestart=$false
)

## Set the file name for this script when copying to $ScriptDirectory
## Note: the .ps1 file extension will be automatically added
$ScriptName = "Fix-WindowsHealth"

## File to log messages into
$LogFile = "${MSPLogs}\${ScriptName}.log"

## Copy script locally and schedule chkdsk at next boot
If ( ($PostRestart -eq $false) -and ($SkipDiskRepair -eq $false) ) {

    ## Copy this script so it can be run post-reboot
    If (! (Test-Path $ScriptDirectory) ) { New-Item -Path $ScriptDirectory -ItemType Directory -Force | Out-Null }
    $ScriptFile = "${ScriptDirectory}\${ScriptName}.ps1"
    If ( $ScriptFile -ne $MyInvocation.MyCommand.Definition ) {
        If ( Test-Path $ScriptFile ) {
            Remove-Item $ScriptFile -Force
        }
        Copy-Item -Path $MyInvocation.MyCommand.Definition -Destination $ScriptFile -Force
        If ( Test-Path $ScriptFile ) {
            Log $LogFile "Successfully copied running script to: ${ScriptFile}"
        } Else { Log $LogFile "Failed to copy running script to: ${ScriptFile}" }
    }

    ## Schedule chkdsk for next reboot
    & cmd.exe /c echo y `| chkdsk.exe /f /r /x
    Log $LogFile "Scheduled chkdsk.exe for next reboot"

    ## Create a RunOnce task to resume this script after the disk has been repaired
    $command = """%SystemRoot%\system32\WindowsPowerShell\v1.0\powershell.exe"" -ExecutionPolicy Bypass -File ""${ScriptFile}"" -PostRestart"
    Set-ItemProperty -Path "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce" -Name $ScriptName -Value $command
    If ( (Get-ItemProperty -Path "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce" -Name $ScriptName) ) {
        Log $LogFile "Successfully scheduled post-reboot task: ${command}"
    } Else { Log $LogFile "Failed to schedule post-reboot task" }
}

## Get the last log entry from chkdsk
If ( $PostRestart ) {

    ## Get the last event chkdsk recorded in the Event Log
    $chkdskmsg = Get-EventLog -LogName 'Application' -Source 'Wininit' -Newest 1 -ErrorAction SilentlyContinue
    If ( $chkdskmsg ) {

        ## Export the chkdsk log to a separate file for review
        $ChkdskLog = "${LogDirectory}\${ScriptName}-Chkdsk.log"
        If ( Test-Path $ChkdskLog ) {
            Remove-Item $ChkdskLog -Force
        }
        New-Item -Path $ChkdskLog -ItemType File -Value $chkdskmsg.Message.ToString()
        If ( Test-Path $ChkdskLog ) {
            Log $LogFile "Successfully exported chkdsk result"
        } Else { Log $LogFile "Failed to export chkdsk result" }

    } Else { Log $LogFile "Last chkdsk result could not be obtained" }
}

## Cleanup online image and restore health
If ( $PostRestart -or $SkipDiskRepair ) {

    ## Run enhanced DISM for Windows 8 and higher assets
    If ( IsWindowsVersion -ge "6.2" ) {

        $BeginDISMRun = $(Get-Date).AddSeconds(-1)
    
        ## Scan the Windows Component Store for problems
        Log $LogFile "Beginning Windows Component Store scan"
        Start-Process "dism.exe" -ArgumentList "/online /cleanup-image /scanhealth" -Wait
        Log $LogFile "Completed Windows Component Store scan"
    
        ## Check to see if the previous command recorded any problems that need repairing
        Log $LogFile "Beginning Windows Component Store health check"
        $CheckHealthOutput = $(dism /online /cleanup-image /checkhealth) | Out-String
        If ( $CheckHealthOutput -match "repairable" ) {
        
            Log $LogFile "Completed Windows Component Store health check: Problems found"
            
            ## Repair problems with the Windows Component Store
            Log $LogFile "Beginning Windows Component Store restoration"
            Start-Process "dism.exe" -ArgumentList "/online /cleanup-image /restorehealth" -Wait
            Log $LogFile "Completed Windows Component Store restoration"
        } Else { Log $LogFile "Completed Windows Component Store health check: No problems found" }

        $EndDISMRun = $(Get-Date).AddSeconds(1)

        ## Parse DISM log file for recent activity
        $DISMLogFile = "${LogDirectory}\${ScriptName}-DISM.log"
        If ( Test-Path $DISMLogFile ) { Remove-Item $DISMLogFile -Force }
        New-Item -Path $DISMLogFile -ItemType File
        Log $LogFile "Beginning Windows Component Store log data collection"
        Get-Content -Path "${Env:WinDir}\Logs\DISM\dism.log" | Select-String -Pattern '^(\d{4})-(\d{2})-(\d{2}) (\d{2}):(\d{2}):(\d{2})' |
            % { If ( ((Get-Date $_.Line.Substring(0,19)) -gt $BeginDISMRun) -and ((Get-Date $_.Line.Substring(0,19)) -le $EndDISMRun) ) {
                Add-Content -Path $DISMLogFile -Value $_.Line
            }
        }

        Log $LogFile "Completed Windows Component Store log data collection"

    } Else { Log $LogFile "Online image cleanup and restoration is only supported on Windows 8 and higher" }
}

## Run SFC
If ( $PostRestart -or $SkipDiskRepair ) {

    ## Start the TrustedInstaller service if it isn't running
    ## Note: this is necessary for the SFC utility to function properly
    If ( (Get-Service -Name TrustedInstaller).Status -ne "Running" ) {
        Restart-Service -Name TrustedInstaller -Force
        Log $LogFile "Started TrustedInstaller service"
    }
    
    ## Start the SFC utility to check validity of Windows system files
    Log $LogFile "Beginning SFC scan and fix"

    $BeginSFCRun = $(Get-Date).AddSeconds(-1)
    Start-Process "sfc" -ArgumentList "/scannow" -Wait
    $EndSFCRun = $(Get-Date).AddSeconds(1)
    Log $LogFile "Completed SFC scan and fix"
    
    ## Parse the CBS.log file and export any SFC related messages
    $SFCLogFile = "${LogDirectory}\${ScriptName}-SFC.log"
    If ( Test-Path $SFCLogFile ) { Remove-Item $SFCLogFile -Force }
    New-Item -Path $SFCLogFile -ItemType File
    Log $LogFile "Beginning SFC log data collection"
    Get-Content -Path "${Env:WinDir}\Logs\CBS\CBS.log" | Select-String -Pattern '^(\d{4})-(\d{2})-(\d{2}) (\d{2}):(\d{2}):(\d{2}).+\[SR\]' |
        % { If ( ((Get-Date $_.Line.Substring(0,19)) -gt $BeginSFCRun) -and ((Get-Date $_.Line.Substring(0,19)) -le $EndSFCRun) ) {
            Add-Content -Path $SFCLogFile -Value $_.Line
        }
    }
    Log $LogFile "Completed SFC log data collection"
}

## Prepare Log File Upload