<#
.SYNOPSIS
    Simple remote maintenance toolkit for Windows PCs via WinRM.

.DESCRIPTION
    Executes maintenance tasks on remote Windows PCs using WinRM.
    Reads target computers from a text file (one hostname/IP per line).

.PARAMETER HostsFile
    Path to text file containing computer names/IPs (one per line).
    Lines starting with # are treated as comments.
    Default: .\hosts.txt

.PARAMETER ComputerName
    Single computer name or IP address (alternative to HostsFile).

.PARAMETER Task
    Maintenance task to execute. Available tasks:
      - RestartSpooler   : Restart Print Spooler service
      - FlushDNS         : Clear DNS resolver cache
      - ClearTempFiles   : Clear Windows temp files
      - DiskCleanup      : Run Windows Disk Cleanup
      - OptimizeDisk     : TRIM (SSD) or Defrag (HDD)
      - SyncTime         : Force time sync with domain controller
      - RestartService   : Restart a specific Windows service
      - RunCommand       : Run a custom command
      - RestartComputer  : Restart the remote PC (requires confirmation)

.PARAMETER ServiceName
    Service name for RestartService task.

.PARAMETER Command
    Custom command for RunCommand task.

.PARAMETER Credential
    PSCredential for remote authentication. Prompts if not provided.

.PARAMETER DnsSuffix
    DNS suffix to append to hostnames (if not already FQDN).
    Default: logon.ds.ge.com

.PARAMETER ThrottleLimit
    Maximum number of concurrent remote connections.
    Default: 10

.PARAMETER LogResults
    Save results to a timestamped log file in the script directory.
    Log files are saved as: RemoteTask_YYYYMMDD_HHMMSS.log

.EXAMPLE
    # Restart print spooler on all hosts in hosts.txt
    .\Invoke-RemoteTask.ps1 -Task RestartSpooler

.EXAMPLE
    # Flush DNS on a single computer
    .\Invoke-RemoteTask.ps1 -ComputerName "PC001" -Task FlushDNS

.EXAMPLE
    # Run disk cleanup on hosts from custom file
    .\Invoke-RemoteTask.ps1 -HostsFile ".\shopfloor-pcs.txt" -Task DiskCleanup

.EXAMPLE
    # Restart a specific service
    .\Invoke-RemoteTask.ps1 -Task RestartService -ServiceName "Spooler"

.EXAMPLE
    # Run custom command
    .\Invoke-RemoteTask.ps1 -Task RunCommand -Command "Get-Process | Select -First 5"

.NOTES
    Author: Shop Floor Tools
    Requirements: PowerShell 5.1+, WinRM enabled on targets
#>

[CmdletBinding(DefaultParameterSetName='ByFile')]
param(
    [Parameter(ParameterSetName='ByFile')]
    [string]$HostsFile = ".\hosts.txt",

    [Parameter(ParameterSetName='ByName')]
    [string[]]$ComputerName,

    [Parameter(Mandatory=$true)]
    [ValidateSet(
        'RestartSpooler', 'FlushDNS', 'ClearTempFiles', 'DiskCleanup',
        'OptimizeDisk', 'SyncTime', 'RestartService', 'RunCommand',
        'GetDiskSpace', 'GetUptime', 'TestConnection', 'RestartComputer'
    )]
    [string]$Task,

    [Parameter()]
    [string]$ServiceName,

    [Parameter()]
    [string]$Command,

    [Parameter()]
    [PSCredential]$Credential,

    [Parameter()]
    [string]$DnsSuffix = "logon.ds.ge.com",

    [Parameter()]
    [int]$ThrottleLimit = 10,

    [Parameter()]
    [switch]$LogResults
)

# =============================================================================
# Helper Functions
# =============================================================================

function Write-Log {
    param([string]$Message, [string]$Level = "INFO")
    $timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    $color = switch ($Level) {
        "ERROR"   { "Red" }
        "WARNING" { "Yellow" }
        "SUCCESS" { "Green" }
        "TASK"    { "Cyan" }
        default   { "White" }
    }
    Write-Host "[$timestamp] [$Level] $Message" -ForegroundColor $color
}

# =============================================================================
# Task Scriptblocks
# =============================================================================

$TaskScripts = @{

    'RestartSpooler' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            Stop-Service -Name Spooler -Force -ErrorAction Stop
            $queuePath = "$env:SystemRoot\System32\spool\PRINTERS"
            if (Test-Path $queuePath) { Remove-Item "$queuePath\*" -Force -ErrorAction SilentlyContinue }
            Start-Service -Name Spooler -ErrorAction Stop
            $status = (Get-Service -Name Spooler).Status
            $result.Success = ($status -eq 'Running')
            $result.Output = "Print Spooler restarted. Status: $status"
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'FlushDNS' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            $flushResult = & ipconfig /flushdns 2>&1
            $result.Output = ($flushResult -join " ").Trim()
            $result.Success = $true
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'ClearTempFiles' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null; FilesDeleted = 0 }
        try {
            $tempPaths = @("$env:TEMP", "$env:SystemRoot\Temp")
            foreach ($path in $tempPaths) {
                if (Test-Path $path) {
                    $files = Get-ChildItem -Path $path -Recurse -Force -ErrorAction SilentlyContinue
                    foreach ($file in $files) {
                        try { Remove-Item $file.FullName -Force -Recurse -ErrorAction SilentlyContinue; $result.FilesDeleted++ } catch { }
                    }
                }
            }
            $result.Success = $true
            $result.Output = "Deleted $($result.FilesDeleted) temp files"
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'DiskCleanup' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null; SpaceFreedMB = 0 }
        try {
            $initialFree = (Get-PSDrive C).Free
            $cleanupPath = "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches"
            $categories = @("Temporary Files", "Temporary Setup Files", "Old ChkDsk Files", "Windows Update Cleanup", "Recycle Bin")
            foreach ($cat in $categories) {
                $catPath = Join-Path $cleanupPath $cat
                if (Test-Path $catPath) { Set-ItemProperty -Path $catPath -Name "StateFlags0100" -Value 2 -ErrorAction SilentlyContinue }
            }
            Start-Process -FilePath "cleanmgr.exe" -ArgumentList "/sagerun:100" -Wait -WindowStyle Hidden
            Start-Sleep -Seconds 2
            $finalFree = (Get-PSDrive C).Free
            $result.SpaceFreedMB = [math]::Round(($finalFree - $initialFree) / 1MB, 0)
            $result.Success = $true
            $result.Output = "Disk cleanup completed. Space freed: $($result.SpaceFreedMB) MB"
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'OptimizeDisk' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            $volumes = Get-Volume | Where-Object { $_.DriveType -eq 'Fixed' -and $_.DriveLetter }
            $optimized = @()
            foreach ($vol in $volumes) {
                $driveLetter = $vol.DriveLetter
                $physicalDisk = Get-PhysicalDisk | Where-Object { $_.DeviceId -eq (Get-Partition -DriveLetter $driveLetter -ErrorAction SilentlyContinue).DiskNumber }
                $mediaType = if ($physicalDisk) { $physicalDisk.MediaType } else { "Unknown" }
                try {
                    if ($mediaType -eq 'SSD') { Optimize-Volume -DriveLetter $driveLetter -ReTrim -ErrorAction Stop; $action = "TRIM" }
                    else { Optimize-Volume -DriveLetter $driveLetter -Defrag -ErrorAction Stop; $action = "Defrag" }
                    $optimized += "${driveLetter}:($action)"
                } catch { $optimized += "${driveLetter}:(Failed)" }
            }
            $result.Success = $true
            $result.Output = "Optimized: $($optimized -join ', ')"
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'SyncTime' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            $source = (& w32tm /query /source 2>&1) -join " "
            $syncResult = & w32tm /resync /force 2>&1
            $result.Success = ($syncResult -match "success" -or $LASTEXITCODE -eq 0)
            $result.Output = "Time synced with $source. Current: $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')"
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'RestartService' = {
        param($ServiceName)
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        if (-not $ServiceName) { $result.Error = "ServiceName parameter required"; return $result }
        try {
            Restart-Service -Name $ServiceName -Force -ErrorAction Stop
            $status = (Get-Service -Name $ServiceName).Status
            $result.Success = ($status -eq 'Running')
            $result.Output = "Service '$ServiceName' restarted. Status: $status"
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'RunCommand' = {
        param($Command)
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        if (-not $Command) { $result.Error = "Command parameter required"; return $result }
        try {
            $output = Invoke-Expression $Command 2>&1
            $result.Output = ($output | Out-String).Trim()
            $result.Success = $true
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'GetDiskSpace' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            $drives = Get-PSDrive -PSProvider FileSystem | Where-Object { $_.Used -ne $null }
            $info = foreach ($drive in $drives) {
                $freeGB = [math]::Round($drive.Free / 1GB, 1)
                $usedGB = [math]::Round($drive.Used / 1GB, 1)
                $totalGB = $freeGB + $usedGB
                $pctFree = if ($totalGB -gt 0) { [math]::Round(($freeGB / $totalGB) * 100, 0) } else { 0 }
                "$($drive.Name): $freeGB GB free ($pctFree%)"
            }
            $result.Output = $info -join ", "
            $result.Success = $true
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'GetUptime' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            $os = Get-CimInstance -ClassName Win32_OperatingSystem
            $uptime = (Get-Date) - $os.LastBootUpTime
            $result.Output = "Up $([math]::Floor($uptime.TotalDays))d $($uptime.Hours)h $($uptime.Minutes)m (Last boot: $($os.LastBootUpTime.ToString('yyyy-MM-dd HH:mm')))"
            $result.Success = $true
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }

    'TestConnection' = {
        $result = @{ Success = $true; Hostname = $env:COMPUTERNAME; Output = "Connection successful"; Error = $null }
        return $result
    }

    'RestartComputer' = {
        $result = @{ Success = $false; Hostname = $env:COMPUTERNAME; Output = ""; Error = $null }
        try {
            $result.Output = "Restart initiated"
            $result.Success = $true
            # Schedule restart in 5 seconds to allow response to return
            Start-Process -FilePath "shutdown.exe" -ArgumentList "/r /t 5 /c `"Remote restart initiated via WinRM`"" -NoNewWindow
        } catch { $result.Error = $_.Exception.Message }
        return $result
    }
}

# =============================================================================
# Main Execution
# =============================================================================

Write-Host ""
Write-Host ("=" * 60) -ForegroundColor Cyan
Write-Host "  Remote Task Executor - Task: $Task" -ForegroundColor Cyan
Write-Host ("=" * 60) -ForegroundColor Cyan
Write-Host ""

# Validate task-specific parameters
if ($Task -eq 'RestartService' -and -not $ServiceName) {
    Write-Log "ServiceName parameter is required for RestartService task" -Level "ERROR"
    exit 1
}
if ($Task -eq 'RunCommand' -and -not $Command) {
    Write-Log "Command parameter is required for RunCommand task" -Level "ERROR"
    exit 1
}
if ($Task -eq 'RestartComputer') {
    Write-Host ""
    Write-Host "WARNING: This will restart the target computer(s)!" -ForegroundColor Yellow
    $confirm = Read-Host "Type 'YES' to confirm"
    if ($confirm -ne 'YES') {
        Write-Log "Restart cancelled by user" -Level "WARNING"
        exit 0
    }
}

# Get credentials
if (-not $Credential) {
    Write-Log "Enter credentials for remote PCs:" -Level "INFO"
    $Credential = Get-Credential -Message "Enter admin credentials for remote PCs"
    if (-not $Credential) {
        Write-Log "Credentials required. Exiting." -Level "ERROR"
        exit 1
    }
}

# Build computer list
$computers = @()

if ($ComputerName) {
    $computers = $ComputerName
} else {
    if (-not (Test-Path $HostsFile)) {
        Write-Log "Hosts file not found: $HostsFile" -Level "ERROR"
        Write-Log "Create a text file with one hostname or IP per line." -Level "INFO"
        exit 1
    }
    $computers = Get-Content $HostsFile | Where-Object { $_.Trim() -and -not $_.StartsWith("#") }
}

if ($computers.Count -eq 0) {
    Write-Log "No computers specified." -Level "ERROR"
    exit 1
}

Write-Log "Target computers: $($computers.Count)" -Level "INFO"
Write-Host ""

# Build FQDNs if DNS suffix provided
$targets = $computers | ForEach-Object {
    $name = $_.Trim()
    if ($DnsSuffix -and $name -notlike "*.*") {
        "$name.$DnsSuffix"
    } else {
        $name
    }
}

# Get the scriptblock
$scriptBlock = $TaskScripts[$Task]

# Create session options
$sessionOption = New-PSSessionOption -OpenTimeout 30000 -OperationTimeout 300000 -NoMachineProfile

Write-Log "Executing on $($targets.Count) computer(s) in parallel (ThrottleLimit: $ThrottleLimit)..." -Level "INFO"
Write-Host ""

# Build arguments for tasks that need them
$taskArgs = @()
if ($Task -eq 'RestartService') { $taskArgs = @($ServiceName) }
if ($Task -eq 'RunCommand') { $taskArgs = @($Command) }

# Show progress indicator
$stopwatch = [System.Diagnostics.Stopwatch]::StartNew()
Write-Host "  [" -NoNewline
Write-Host "Running..." -ForegroundColor Yellow -NoNewline
Write-Host "] Please wait..." -NoNewline

# Execute on all remote computers in parallel using Invoke-Command
$results = @()
$connectionErrors = @()

try {
    if ($taskArgs.Count -gt 0) {
        $results = Invoke-Command -ComputerName $targets -ScriptBlock $scriptBlock -ArgumentList $taskArgs `
            -Credential $Credential -SessionOption $sessionOption -Authentication Negotiate `
            -ThrottleLimit $ThrottleLimit -ErrorAction SilentlyContinue -ErrorVariable connectionErrors
    } else {
        $results = Invoke-Command -ComputerName $targets -ScriptBlock $scriptBlock `
            -Credential $Credential -SessionOption $sessionOption -Authentication Negotiate `
            -ThrottleLimit $ThrottleLimit -ErrorAction SilentlyContinue -ErrorVariable connectionErrors
    }
} catch {
    Write-Host ""
    Write-Log "Execution error: $($_.Exception.Message)" -Level "ERROR"
}

$stopwatch.Stop()
$elapsed = $stopwatch.Elapsed.TotalSeconds

# Clear the progress line
Write-Host "`r" -NoNewline
Write-Host (" " * 60) -NoNewline
Write-Host "`r" -NoNewline
Write-Log "Completed in $([math]::Round($elapsed, 1)) seconds" -Level "INFO"
Write-Host ""

# Collect all results (successes, failures, connection errors)
$allResults = @()
$successCount = 0
$failCount = 0

foreach ($result in $results) {
    $status = if ($result.Success) { "OK"; $successCount++ } else { "FAIL"; $failCount++ }
    $message = if ($result.Success) { $result.Output } else { $result.Error }
    $allResults += [PSCustomObject]@{
        Status = $status
        Computer = $result.Hostname
        Message = $message
    }
}

# Process connection errors
foreach ($err in $connectionErrors) {
    $targetName = if ($err.TargetObject) { $err.TargetObject } else { "Unknown" }
    # Extract just the computer name from FQDN for display
    $shortName = ($targetName -split '\.')[0]
    $errorMsg = $err.Exception.Message -replace '\r?\n', ' '
    # Truncate long error messages
    if ($errorMsg.Length -gt 60) { $errorMsg = $errorMsg.Substring(0, 57) + "..." }
    $allResults += [PSCustomObject]@{
        Status = "FAIL"
        Computer = $shortName
        Message = $errorMsg
    }
    $failCount++
}

# Sort results: failures first, then successes
$allResults = $allResults | Sort-Object @{Expression={$_.Status}; Descending=$true}, Computer

# Display results in a formatted table
Write-Host "  STATUS   COMPUTER             RESULT" -ForegroundColor Cyan
Write-Host "  ------   --------             ------" -ForegroundColor Cyan

foreach ($r in $allResults) {
    $statusColor = if ($r.Status -eq "OK") { "Green" } else { "Red" }
    $statusIcon = if ($r.Status -eq "OK") { "[OK]  " } else { "[FAIL]" }

    # Pad/truncate computer name to 20 chars
    $compName = $r.Computer
    if ($compName.Length -gt 18) { $compName = $compName.Substring(0, 15) + "..." }
    $compName = $compName.PadRight(20)

    # Truncate message if too long
    $msg = $r.Message
    if ($msg.Length -gt 50) { $msg = $msg.Substring(0, 47) + "..." }

    Write-Host "  " -NoNewline
    Write-Host $statusIcon -ForegroundColor $statusColor -NoNewline
    Write-Host " $compName " -NoNewline
    Write-Host $msg -ForegroundColor $(if ($r.Status -eq "OK") { "White" } else { "Yellow" })
}

Write-Host ""

# Summary
Write-Host ""
Write-Host ("=" * 60) -ForegroundColor Cyan
Write-Host "  SUMMARY" -ForegroundColor Cyan
Write-Host ("=" * 60) -ForegroundColor Cyan
Write-Host "  Task:       $Task" -ForegroundColor White
Write-Host "  Total:      $($computers.Count)" -ForegroundColor White
Write-Host "  Successful: $successCount" -ForegroundColor Green
Write-Host "  Failed:     $failCount" -ForegroundColor $(if ($failCount -gt 0) { "Red" } else { "White" })
Write-Host ("=" * 60) -ForegroundColor Cyan
Write-Host ""

# Save to log file if requested
if ($LogResults) {
    $scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
    if (-not $scriptDir) { $scriptDir = Get-Location }
    $logTimestamp = Get-Date -Format "yyyyMMdd_HHmmss"
    $logFile = Join-Path $scriptDir "RemoteTask_$logTimestamp.log"

    $logContent = @()
    $logContent += "=" * 60
    $logContent += "Remote Task Execution Log"
    $logContent += "=" * 60
    $logContent += "Date/Time:    $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')"
    $logContent += "Task:         $Task"
    $logContent += "Targets:      $($computers.Count)"
    $logContent += "ThrottleLimit: $ThrottleLimit"
    $logContent += "Elapsed:      $([math]::Round($elapsed, 1)) seconds"
    if ($ServiceName) { $logContent += "ServiceName:  $ServiceName" }
    if ($Command) { $logContent += "Command:      $Command" }
    $logContent += ""
    $logContent += "=" * 60
    $logContent += "RESULTS"
    $logContent += "=" * 60
    $logContent += ""
    $logContent += "STATUS     COMPUTER                     RESULT"
    $logContent += "------     --------                     ------"

    foreach ($r in $allResults) {
        $statusText = if ($r.Status -eq "OK") { "[OK]    " } else { "[FAIL]  " }
        $compText = $r.Computer.PadRight(28)
        $logContent += "$statusText $compText $($r.Message)"
    }

    $logContent += ""
    $logContent += "=" * 60
    $logContent += "SUMMARY"
    $logContent += "=" * 60
    $logContent += "Total:      $($computers.Count)"
    $logContent += "Successful: $successCount"
    $logContent += "Failed:     $failCount"
    $logContent += "=" * 60

    $logContent | Out-File -FilePath $logFile -Encoding UTF8
    Write-Log "Results saved to: $logFile" -Level "SUCCESS"
    Write-Host ""
}

# Results are displayed above and optionally saved to log file
# To capture results programmatically, use: $results = Invoke-Command ... directly