<#
    LXS - System Restore Point (Windows)
    Description: Enable System Protection, create checkpoints, list them and
                 roll back. The safety net for harden.ps1 and debloat.ps1.
    Repo: https://git.hyko.cx/hykocx/lxs
#>

# Load LXS common library (colors, UI helpers, spinner, loggers, guards).
# Prefers the sibling ..\lib\common.ps1 (repo checkout or installed layout) and
# only hits the network when this script is run standalone.
$LxsRawBase = if ($env:LXS_RAW_BASE) { $env:LXS_RAW_BASE } else { 'https://git.hyko.cx/hykocx/lxs/raw/branch/main' }
if (-not $env:LXS_RAW_PLATFORM_BASE) { $env:LXS_RAW_PLATFORM_BASE = "$LxsRawBase/windows" }
$LxsLibPath = if ($PSScriptRoot) { Join-Path $PSScriptRoot '..\lib\common.ps1' } else { $null }
if ($LxsLibPath -and (Test-Path $LxsLibPath)) {
    . $LxsLibPath
} else {
    try {
        $LxsLibSource = Invoke-RestMethod -Uri "$env:LXS_RAW_PLATFORM_BASE/lib/common.ps1" -UseBasicParsing -ErrorAction Stop
    } catch {
        Write-Error 'Failed to fetch lib/common.ps1'
        exit 1
    }
    . ([scriptblock]::Create($LxsLibSource))
}

$env:LXS_LOG_FILE = Join-Path (Get-LxsTempDir) 'lxs_restore_point.log'

if (-not (Assert-LxsWindows)) { exit 1 }
Assert-LxsAdmin -ScriptPath $PSCommandPath -Arguments $args

function Get-LxsRestorePoints {
    try {
        return @(Get-ComputerRestorePoint -ErrorAction Stop)
    } catch {
        return @()
    }
}

function Show-LxsRestorePoints {
    Clear-Host
    Show-LxsBoxTop -Title 'RESTORE POINTS'
    Write-Host ''
    $points = Get-LxsRestorePoints
    if ($points.Count -eq 0) {
        Write-LxsWarn 'No restore points found (System Protection may be disabled).'
        return
    }
    $points | Sort-Object SequenceNumber -Descending |
        Format-Table -AutoSize `
            @{ Name = 'Seq'; Expression = { $_.SequenceNumber } },
            @{ Name = 'Created'; Expression = { $_.ConvertToDateTime($_.CreationTime).ToString('yyyy-MM-dd HH:mm') } },
            @{ Name = 'Type'; Expression = { $_.RestorePointType } },
            @{ Name = 'Description'; Expression = { $_.Description } } |
        Out-String -Width 160 | Write-Host
}

function Show-LxsProtectionStatus {
    Clear-Host
    Show-LxsBoxTop -Title 'SYSTEM PROTECTION'
    Write-Host ''
    $drive = "$env:SystemDrive\"
    try {
        # Volumes listed here have shadow storage configured, i.e. protection on.
        $shadow = @(Get-CimInstance Win32_ShadowStorage -ErrorAction Stop)
        if ($shadow.Count -gt 0) {
            Write-LxsOk "System Protection appears to be ON for $($shadow.Count) volume(s)"
            foreach ($s in $shadow) {
                $maxGB = [math]::Round($s.MaxSpace / 1GB, 1)
                $usedGB = [math]::Round($s.UsedSpace / 1GB, 2)
                Write-Host "  Shadow storage: $usedGB GB used of $maxGB GB allocated"
            }
        } else {
            Write-LxsWarn "System Protection looks OFF for $drive — no restore points can be created."
        }
    } catch {
        Write-LxsWarn 'Could not read the shadow storage configuration.'
    }
    Write-Host ''
    Write-Host "$($script:Gray)Restore points found: $((Get-LxsRestorePoints).Count)$($script:NC)"
    Write-Host ''
}

function Enable-LxsProtection {
    Clear-Host
    Show-LxsBoxTop -Title 'ENABLE SYSTEM PROTECTION'
    Write-Host ''
    $drive = "$env:SystemDrive\"
    try {
        Enable-ComputerRestore -Drive $drive -ErrorAction Stop
        Write-LxsOk "System Protection enabled on $drive"
    } catch {
        Write-LxsErr "Failed to enable System Protection: $($_.Exception.Message)"
        return
    }
    # Default shadow storage can be as low as 1%; 5% leaves room for a few points.
    try {
        & vssadmin.exe Resize ShadowStorage /For=$drive /On=$drive /MaxSize=5% | Out-String | Write-Host
    } catch {
        Write-LxsWarn 'Could not resize the shadow storage allocation.'
    }
}

function New-LxsCheckpoint {
    Clear-Host
    Show-LxsBoxTop -Title 'CREATE RESTORE POINT'
    Write-Host ''
    $desc = Read-Host -Prompt 'Description (default: LXS manual checkpoint)'
    if ([string]::IsNullOrWhiteSpace($desc)) { $desc = 'LXS manual checkpoint' }
    Write-Host ''
    Write-LxsInfo 'Creating the restore point (this can take a minute)...'
    New-LxsRestorePoint -Description $desc | Out-Null
}

function Restore-LxsCheckpoint {
    Clear-Host
    Show-LxsBoxTop -Title 'ROLL BACK' -Right 'REBOOTS'
    Write-Host ''
    $points = Get-LxsRestorePoints
    if ($points.Count -eq 0) {
        Write-LxsWarn 'No restore points to roll back to.'
        return
    }
    Show-LxsRestorePoints
    Write-Host "$($script:Red)Restoring reboots the machine immediately and reverts system files,$($script:NC)"
    Write-Host "$($script:Red)drivers, registry and installed programs to that point in time.$($script:NC)"
    Write-Host "$($script:Gray)Your documents are not touched.$($script:NC)"
    Write-Host ''
    $seq = Read-Host -Prompt 'Sequence number to restore (empty to cancel)'
    if ([string]::IsNullOrWhiteSpace($seq)) { Write-LxsInfo 'Cancelled.'; return }
    if ($seq -notmatch '^\d+$') { Write-LxsErr 'Sequence number must be numeric.'; return }
    if (-not ($points | Where-Object { $_.SequenceNumber -eq [int]$seq })) {
        Write-LxsErr "No restore point with sequence number $seq."
        return
    }
    if (-not (Confirm-LxsAction -Question "Restore to point $seq and reboot NOW?")) {
        Write-LxsInfo 'Cancelled.'
        return
    }
    try {
        Restore-Computer -RestorePoint ([int]$seq) -Confirm:$false -ErrorAction Stop
    } catch {
        Write-LxsErr "Restore failed: $($_.Exception.Message)"
        Write-Host "$($script:Gray)     You can also use: rstrui.exe$($script:NC)"
    }
}

function Show-LxsRestoreMenu {
    while ($true) {
        Clear-Host
        Show-LxsBoxTop -Title 'RESTORE POINT' -Right 'ADMIN'
        Write-Host ''
        Show-LxsMenuItem '1' 'Show protection status'
        Show-LxsMenuItem '2' 'Enable System Protection'
        Show-LxsMenuItem '3' 'Create a restore point'
        Show-LxsMenuItem '4' 'List restore points'
        Show-LxsMenuItem '5' 'Roll back to a restore point' 'reboots the machine'
        Show-LxsMenuItem '0' 'Back' '' -Exit
        Write-Host ''
        Show-LxsBoxBottom
        Write-Host ''
        $choice = Read-LxsChoice
        Write-Host ''

        switch ($choice) {
            '1' { Show-LxsProtectionStatus; Read-LxsEnter }
            '2' { Enable-LxsProtection;     Read-LxsEnter }
            '3' { New-LxsCheckpoint;        Read-LxsEnter }
            '4' { Show-LxsRestorePoints;    Read-LxsEnter }
            '5' { Restore-LxsCheckpoint;    Read-LxsEnter }
            '0' { return }
            default { Write-LxsErr 'Invalid protocol. Select 0-5.'; Start-Sleep -Seconds 1 }
        }
    }
}

Show-LxsRestoreMenu
exit 75