# ------- Helpers -------
function Invoke-Remote {
    param(
        [Parameter(Mandatory)][string]$Computer,
        [Parameter(Mandatory)][scriptblock]$ScriptBlock,
        [Parameter()][object[]]$ArgumentList
    )
    try {
        return Invoke-Command -ComputerName $Computer -ScriptBlock $ScriptBlock -ArgumentList $ArgumentList -ErrorAction Stop
    } catch {
        $currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name
        Write-Warning "Implicit auth to '$Computer' failed: $($_.Exception.Message)"
        $cred = Get-Credential -UserName $currentUser -Message "Enter the password for $currentUser"
        return Invoke-Command -ComputerName $Computer -Credential $cred -ScriptBlock $ScriptBlock -ArgumentList $ArgumentList -ErrorAction Stop
    }
}

function Remove-ProvisionedAndExisting {
    <#
      .SYNOPSIS
        Preview + remove provisioned and installed Appx packages, locally or remotely.
      .PARAMETER Pattern
        -Like pattern, e.g. 'Microsoft.MSPaint*'
      .PARAMETER Computer
        If provided, runs on remote via Invoke-Remote; otherwise runs locally.
    #>
    param(
        [Parameter(Mandatory)][string]$Pattern,
        [string]$Computer
    )

    # Scriptblocks reused for both local and remote
    $sbPreview = {
        param($p)
        $prov = Get-AppxProvisionedPackage -Online |
                Where-Object DisplayName -Like $p |
                Select-Object @{n='Type';e={'Provisioned'}}, DisplayName, PackageName

        $inst = Get-AppxPackage -AllUsers |
                Where-Object Name -Like $p |
                Select-Object @{n='Type';e={'Installed'}}, Name, PackageFullName

        [PSCustomObject]@{ Provisioned = $prov; Installed = $inst }
    }

    $sbRemove = {
        param($p)
        $errs = @()

        # Remove from provisioned image (new profiles)
        try {
            Get-AppxProvisionedPackage -Online |
                Where-Object DisplayName -Like $p |
                Remove-AppxProvisionedPackage -Online -ErrorAction Stop | Out-Null
            Write-Host "Removed provisioned packages matching '$p'."
        } catch { $errs += "Provisioned removal: $($_.Exception.Message)" }

        # Remove from existing users
        try {
            $pkgs = Get-AppxPackage -AllUsers | Where-Object Name -Like $p
            foreach ($pkg in $pkgs) {
                try {
                    Remove-AppxPackage -Package $pkg.PackageFullName -AllUsers -ErrorAction Stop
                    Write-Host "Removed installed package: $($pkg.PackageFullName)"
                } catch {
                    # Fallback on older systems without -AllUsers support
                    try {
                        Get-AppxPackage -AllUsers | Where-Object Name -eq $pkg.Name |
                          ForEach-Object { Remove-AppxPackage -Package $_.PackageFullName -ErrorAction Stop }
                        Write-Host "Removed per-user instances of: $($pkg.Name)"
                    } catch {
                        $errs += "Installed removal ($($pkg.Name)): $($_.Exception.Message)"
                    }
                }
            }
        } catch { $errs += "Installed discovery: $($_.Exception.Message)" }

        if ($errs.Count) { Write-Warning ("Errors:`n - " + ($errs -join "`n - ")) }
        else { Write-Host "Removal complete." -ForegroundColor Green }
    }

    $isLocal = [string]::IsNullOrWhiteSpace($Computer) -or $Computer -in @('localhost','127.0.0.1','.')

    # --- Preview ---
    if ($isLocal) {
        $preview = & $sbPreview $Pattern
    } else {
        $preview = Invoke-Remote -Computer $Computer -ScriptBlock $sbPreview -ArgumentList $Pattern
    }

    $prov = $preview.Provisioned
    $inst = $preview.Installed

    if (-not $prov -and -not $inst) {
        if ($isLocal) { $scope = "local machine" } else { $scope = $Computer }
        Write-Host "No matches for '$Pattern' on $scope." -ForegroundColor Yellow
        return
    }

    if ($prov) {
        Write-Host "`nProvisioned packages matching '$Pattern':" -ForegroundColor Cyan
        $prov | Format-Table -AutoSize
    } else {
        Write-Host "`nNo provisioned packages matched '$Pattern'."
    }

    if ($inst) {
        Write-Host "`nInstalled packages (existing profiles) matching '$Pattern':" -ForegroundColor Cyan
        $inst | Select-Object Name, PackageFullName | Format-Table -AutoSize
    } else {
        Write-Host "`nNo installed packages (existing profiles) matched '$Pattern'."
    }

    $confirm = Read-Host "`nProceed to remove from provisioned image AND existing profiles? (Y/N)"
    if ($confirm -notin @('Y','y')) { Write-Host "Cancelled."; return }

    # --- Remove ---
    if ($isLocal) {
        & $sbRemove $Pattern
    } else {
        Invoke-Remote -Computer $Computer -ScriptBlock $sbRemove -ArgumentList $Pattern | Out-Host
    }
}

function Invoke-WindowsUpdateRepair {
    param([string]$Computer)

    $sb = {
        function Test-Admin {
            $id = [Security.Principal.WindowsIdentity]::GetCurrent()
            $p  = New-Object Security.Principal.WindowsPrincipal($id)
            $p.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
        }

        if (-not (Test-Admin)) { throw "You must run this elevated (Administrator)." }

        # Local helper for progress (only available when running locally)
        function Invoke-DismWithProgress {
            param([string[]]$Arguments)
            $ts   = Get-Date -Format 'yyyyMMdd_HHmmss'
            $outf = Join-Path $env:TEMP "dism_out_$ts.txt"
            $errf = Join-Path $env:TEMP "dism_err_$ts.txt"
            if (-not ($Arguments -match '/LogPath:')) {
                $logFile = Join-Path $env:TEMP "dism_log_$ts.log"
                $Arguments += "/LogPath:$logFile"
            }
            $p = Start-Process dism.exe -ArgumentList ($Arguments -join ' ') -NoNewWindow -RedirectStandardOutput $outf -RedirectStandardError $errf -PassThru
            $lastPct = -1
            while (-not $p.HasExited) {
                if (Test-Path $outf) {
                    $content = Get-Content $outf -Raw -ErrorAction SilentlyContinue
                    if ($content) {
                        $m = [regex]::Matches($content, '(\d{1,3}(?:\.\d+)?)%')
                        if ($m.Count -gt 0) {
                            $pct = [int][math]::Min(100, [double]$m[$m.Count-1].Value)
                            if ($pct -ne $lastPct) {
                                Write-Progress -Activity "DISM $($Arguments -join ' ')" -Status "$pct%" -PercentComplete $pct
                                $lastPct = $pct
                            }
                        } else {
                            Write-Progress -Activity "DISM $($Arguments -join ' ')" -Status "Working..." -PercentComplete 0
                        }
                    }
                }
                Start-Sleep -Milliseconds 400
            }
            Write-Progress -Activity "DISM $($Arguments -join ' ')" -Completed
            if ($p.ExitCode -ne 0) {
                Write-Warning "DISM exited with code $($p.ExitCode)"
                if (Test-Path $errf) { Get-Content $errf -Tail 25 | Write-Warning }
            }
            return $p.ExitCode
        }

        Write-Host "=== Step 1: DISM (ScanHealth) ===" -ForegroundColor Cyan
        $e1 = Invoke-DismWithProgress -Arguments @('/Online','/Cleanup-Image','/ScanHealth')
        if ($e1 -ne 0) { throw "DISM ScanHealth failed with code $e1" }

        Write-Host "`n=== Step 1b: DISM (RestoreHealth) ===" -ForegroundColor Cyan
        $e2 = Invoke-DismWithProgress -Arguments @('/Online','/Cleanup-Image','/RestoreHealth','/NoRestart')
        if ($e2 -ne 0) { throw "DISM RestoreHealth failed with code $e2" }

        Write-Host "`n=== Step 2: SFC ===" -ForegroundColor Cyan
        & sfc.exe /scannow | Out-Host

        Write-Host "`n=== Step 3: Reset Windows Update components ===" -ForegroundColor Cyan
        $services = 'wuauserv','bits','cryptsvc','msiserver'
        foreach ($s in $services) { try { Stop-Service -Name $s -Force -ErrorAction Stop } catch { } }

        $sd  = Join-Path $env:SystemRoot 'SoftwareDistribution'
        $cat = Join-Path $env:SystemRoot 'System32\catroot2'
        $ts  = Get-Date -Format 'yyyyMMdd_HHmmss'
        if (Test-Path $sd)  { try { Rename-Item $sd  ("SoftwareDistribution.bak_$ts") -ErrorAction Stop } catch { } }
        if (Test-Path $cat) { try { Rename-Item $cat ("catroot2.bak_$ts")             -ErrorAction Stop } catch { } }
        foreach ($s in $services) { try { Start-Service -Name $s -ErrorAction Stop } catch { } }

        Write-Host "`n=== Step 4: Trigger update scan ===" -ForegroundColor Cyan
        foreach ($cmd in @(
            { & usoclient.exe StartScan },
            { & usoclient.exe StartDownload },
            { & usoclient.exe StartInstall },
            { & wuauclt.exe /detectnow },
            { & wuauclt.exe /reportnow }
        )) { try { & $cmd | Out-Null } catch { } }

        Write-Host "`nWindows Update repair sequence completed." -ForegroundColor Green
    }

    if ([string]::IsNullOrWhiteSpace($Computer) -or $Computer -in @('localhost','127.0.0.1','.')) {
        & $sb
    } else {
        # Over remoting, Write-Progress doesn't render well—still run the repair,
        # but you'll see normal text output. (If you want, we can add a remote-side
        # log tail that writes periodic progress messages.)
        Invoke-Remote -Computer $Computer -ScriptBlock $sb
    }
}

function Invoke-DismWithProgress {
    <#
      Runs DISM with a real progress bar.
      Returns: DISM exit code (0 = success)
      Note: Progress display is local-only (remoting doesn't surface Write-Progress well).
    #>
    param(
        [Parameter(Mandatory)][string[]]$Arguments
    )

    $ts   = Get-Date -Format 'yyyyMMdd_HHmmss'
    $outf = Join-Path $env:TEMP "dism_out_$ts.txt"
    $errf = Join-Path $env:TEMP "dism_err_$ts.txt"

    # Ensure we get a log file too (helps debugging)
    if (-not ($Arguments -match '/LogPath:')) {
        $logFile = Join-Path $env:TEMP "dism_log_$ts.log"
        $Arguments += "/LogPath:$logFile"
    }

    # Launch DISM and redirect output so we can parse percentages safely
    $p = Start-Process -FilePath dism.exe `
                       -ArgumentList ($Arguments -join ' ') `
                       -NoNewWindow `
                       -RedirectStandardOutput $outf `
                       -RedirectStandardError  $errf `
                       -PassThru

    $lastPct = -1
    while (-not $p.HasExited) {
        if (Test-Path $outf) {
            # Read current output and grab the latest percentage like "63.8%" or "63%"
            $content = Get-Content -Path $outf -Raw -ErrorAction SilentlyContinue
            if ($content) {
                $m = [regex]::Matches($content, '(\d{1,3}(?:\.\d+)?)%')
                if ($m.Count -gt 0) {
                    $pct = [int][math]::Min(100, [double]$m[$m.Count-1].Groups[1].Value)
                    if ($pct -ne $lastPct) {
                        Write-Progress -Activity "DISM $($Arguments -join ' ')" -Status "$pct%" -PercentComplete $pct
                        $lastPct = $pct
                    }
                } else {
                    # No explicit percent yet—show a spinner
                    Write-Progress -Activity "DISM $($Arguments -join ' ')" -Status "Working..." -PercentComplete 0
                }
            }
        }
        Start-Sleep -Milliseconds 400
    }

    # Complete the bar cleanly
    Write-Progress -Activity "DISM $($Arguments -join ' ')" -Completed

    $exit = $p.ExitCode

    # Show a quick summary if it failed
    if ($exit -ne 0) {
        Write-Warning "DISM exited with code $exit"
        if (Test-Path $errf) { Get-Content $errf -Tail 25 | Write-Warning }
    }

    return $exit
}

# -------  task functions (menu sections) -------

function Do-ConnectRemote {
    $computer = Read-Host "Enter computer name (e.g. USERS-LAPTOP)"
    if ([string]::IsNullOrWhiteSpace($computer)) { return }
    try {
        Enter-PSSession -ComputerName $computer
    } catch {
        Write-Warning "Implicit auth to '$computer' failed: $($_.Exception.Message)"
        $currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name
        $cred = Get-Credential -UserName $currentUser -Message "Enter the password for $currentUser"
        Enter-PSSession -ComputerName $computer -Credential $cred
    }
}

function Do-RemoveAppxMenu {
    # Choose Local or Remote first
    $targetType = $null
    $computer   = $null
    do {
        Clear-Host
        Write-Host "=== Remove Windows Apps ===" -ForegroundColor Cyan
        Write-Host " 1) Local machine"
        Write-Host " 2) Remote machine"
        Write-Host " 0) Back`n"
        $t = Read-Host "Choose target"
        switch ($t) {
            '1' { $targetType = 'Local' }
            '2' { $targetType = 'Remote'; $computer = Read-Host "Enter remote computer name"; if ([string]::IsNullOrWhiteSpace($computer)) { $targetType = $null } }
            '0' { return }
            default { Write-Host "Invalid choice." -ForegroundColor Yellow; Start-Sleep -Milliseconds 800 }
        }
    } until ($targetType)

    if ($targetType -eq 'Local') { $scopeLabel = 'local machine' } else { $scopeLabel = $computer }

    # Presets
    $presets = @(
        @{ Label = 'Paint 3D';                 Pattern = 'Microsoft.MSPaint*' }
        @{ Label = '3D Viewer';                Pattern = 'Microsoft.Microsoft3DViewer*' }
        @{ Label = 'Mail & Calendar';          Pattern = 'microsoft.windowscommunicationsapps*' }
        @{ Label = 'Sticky Notes';             Pattern = 'Microsoft.MicrosoftStickyNotes*' }
        @{ Label = 'Solitaire Collection';     Pattern = 'Microsoft.MicrosoftSolitaireCollection*' }
        @{ Label = 'Xbox (Gaming App)';        Pattern = 'Microsoft.GamingApp*' }
        @{ Label = 'Xbox Overlay/Identity';    Pattern = 'Microsoft.Xbox*' }
        @{ Label = 'Clipchamp';                Pattern = 'Microsoft.Clipchamp*' }
        @{ Label = 'Weather';                  Pattern = 'Microsoft.BingWeather*' }
        @{ Label = 'News';                     Pattern = 'Microsoft.BingNews*' }
    )

    do {
        Clear-Host
        Write-Host "=== Remove Apps on $scopeLabel ===" -ForegroundColor Cyan
        $i = 1
        foreach ($p in $presets) { Write-Host (" {0,2}) {1}  (pattern: {2})" -f $i, $p.Label, $p.Pattern); $i++ }
        Write-Host (" {0,2}) {1}" -f $i, 'Custom pattern...')
        Write-Host "  0) Back`n"
        $sub = Read-Host "Choose an option"

        if ($sub -eq '0') { break }

        if ($sub -as [int]) {
            $idx = [int]$sub
            if ($idx -ge 1 -and $idx -le $presets.Count) {
                $pat = $presets[$idx-1].Pattern
                if ($targetType -eq 'Local') { Remove-ProvisionedAndExisting -Pattern $pat }
                else                         { Remove-ProvisionedAndExisting -Pattern $pat -Computer $computer }
                Read-Host "`nPress Enter to continue" | Out-Null
                continue
            } elseif ($idx -eq ($presets.Count + 1)) {
                $pat = Read-Host "Enter -Like pattern (use * wildcards, e.g. Microsoft.Todos*)"
                if ($pat) {
                    if ($targetType -eq 'Local') { Remove-ProvisionedAndExisting -Pattern $pat }
                    else                         { Remove-ProvisionedAndExisting -Pattern $pat -Computer $computer }
                }
                Read-Host "`nPress Enter to continue" | Out-Null
                continue
            }
        }

        Write-Host "Invalid choice." -ForegroundColor Yellow
        Start-Sleep -Milliseconds 800
    } while ($true)
}

function Do-RestartServiceMenu {
    do {
        Clear-Host
        Write-Host "=== Restart a Service ===" -ForegroundColor Cyan
        Write-Host " 1) Local machine"
        Write-Host " 2) Remote machine"
        Write-Host " 0) Back`n"
        $opt = Read-Host "Choose an option"

        switch ($opt) {
            '1' {
                $svc = Read-Host "Enter service name (e.g., Spooler)"
                if ([string]::IsNullOrWhiteSpace($svc)) { break }
                try {
                    Restart-Service -Name $svc -Force -ErrorAction Stop
                    Get-Service -Name $svc | Select-Object Status, Name, DisplayName | Format-Table -AutoSize
                    Write-Host "Local service restarted." -ForegroundColor Green
                } catch {
                    Write-Warning "Local restart failed: $($_.Exception.Message)"
                }
                Read-Host "`nPress Enter to continue" | Out-Null
            }
            '2' {
                $computer = Read-Host "Enter remote computer name"
                if ([string]::IsNullOrWhiteSpace($computer)) { break }
                $svc = Read-Host "Enter service name on $computer"
                if ([string]::IsNullOrWhiteSpace($svc)) { break }

                try {
                    $res = Invoke-Remote -Computer $computer -ScriptBlock {
                        param($name)
                        $null = Get-Service -Name $name -ErrorAction Stop
                        Restart-Service -Name $name -Force -ErrorAction Stop
                        Get-Service -Name $name | Select-Object Status, Name, DisplayName
                    } -ArgumentList $svc

                    if ($res) { $res | Format-Table -AutoSize }
                    Write-Host "Remote service restarted on $computer." -ForegroundColor Green
                } catch {
                    Write-Warning "Remote restart failed on $($computer): $($_.Exception.Message)"
                }
                Read-Host "`nPress Enter to continue" | Out-Null
            }
            '0' { return }
            default {
                Write-Host "Invalid choice." -ForegroundColor Yellow
                Start-Sleep -Milliseconds 800
            }
        }
    } while ($true)
}

function do-wingetupgrades {
    winget upgrade --all --silent --accept-source-agreements --accept-package-agreements
}

function do-winupdatefix {
    do {
        Clear-Host
        Write-Host "=== Windows Update Repair ===" -ForegroundColor Cyan
        Write-Host " 1) Local machine"
        Write-Host " 2) Remote machine"
        Write-Host " 0) Back`n"
        $t = Read-Host "Choose target"

        switch ($t) {
            '1' {
                $go = Read-Host "This will run DISM/SFC and reset WU components on THIS machine. Proceed? (Y/N)"
                if ($go -in @('Y','y')) {
                    try { Invoke-WindowsUpdateRepair } catch { Write-Warning $_ }
                }
                Read-Host "`nPress Enter to continue" | Out-Null
            }
            '2' {
                $comp = Read-Host "Enter remote computer name"
                if ([string]::IsNullOrWhiteSpace($comp)) { break }
                $go = Read-Host "Run repair on '$comp'? This requires admin on that machine. Proceed? (Y/N)"
                if ($go -in @('Y','y')) {
                    try { Invoke-WindowsUpdateRepair -Computer $comp } catch { Write-Warning $_ }
                }
                Read-Host "`nPress Enter to continue" | Out-Null
            }
            '0' { return }  # leave submenu back to main menu
            default {
                Write-Host "Invalid choice." -ForegroundColor Yellow
                Start-Sleep -Milliseconds 800
            }
        }
    } while ($true)
}

# ------- Main Menu -------
do {
    Clear-Host
    Write-Host "=== SRG Admin Menu ===" -ForegroundColor Cyan
    Write-Host " 1) Connect to a machine (Enter-PSSession as current user)"
    Write-Host " 2) Remove Windows apps (local or remote)"
    Write-Host " 3) Restart a service (submenu)"
    write-Host " 4) Run Winget Upgrades" 
    write-Host " 5) Run Windows Update Repair"
    Write-Host " 0) Quit`n"
    $choice = Read-Host "Choose an option"

    switch ($choice) {
        '1' { Do-ConnectRemote }
        '2' { Do-RemoveAppxMenu }
        '3' { Do-RestartServiceMenu }
        '4' { do-wingetupgrades }
        '5' { do-winupdatefix }
        '0' { exit }
        default {
            Write-Host "Invalid choice." -ForegroundColor Yellow
            Start-Sleep -Milliseconds 800
        }
    }
} while ($true)