<#
.SYNOPSIS
    WSUS Compliance Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for WSUS compliance - server status, computer targets,
    update approval states, sync status, database and content store sizes,
    and per-group compliance percentages. Exports metrics as
    Prometheus-compatible text format.
.PARAMETER Mode
    Output mode: 'stdout' (default), 'textfile', or 'http'
.PARAMETER Port
    HTTP port for http mode (default: 9516)
.PARAMETER TextfileDir
    Directory for textfile collector output (default: C:\ProgramData\node_exporter)
.PARAMETER InstallScheduledTask
    Switch to create a scheduled task for auto-start on system boot
.PARAMETER TaskIntervalMinutes
    Interval in minutes for the scheduled task (default: 2)
.NOTES
    Author:  Phil Connor
    Contact: contact@mylinux.work
    Website: https://mylinux.work
    License: MIT
    Version: 1.0

    Metrics Exported:
      Core Status:
        - wsus_up
        - wsus_exporter_info{version}

      Computers:
        - wsus_computers_total
        - wsus_computers_needing_updates
        - wsus_computers_with_errors
        - wsus_computers_not_contacted

      Updates:
        - wsus_updates_approved_total
        - wsus_updates_declined_total
        - wsus_updates_not_approved_total
        - wsus_updates_needed_total
        - wsus_updates_installed_total
        - wsus_updates_failed_total

      Sync:
        - wsus_last_sync_timestamp
        - wsus_sync_success

      Storage:
        - wsus_database_size_bytes
        - wsus_content_size_bytes

      Group Compliance:
        - wsus_computer_group_compliance{group}

      Exporter:
        - wsus_exporter_duration_seconds
        - wsus_exporter_last_run_timestamp
#>

param(
    [ValidateSet('stdout', 'textfile', 'http')]
    [string]$Mode = 'stdout',

    [int]$Port = 9516,

    [string]$TextfileDir = 'C:\ProgramData\node_exporter',

    [switch]$InstallScheduledTask,

    [int]$TaskIntervalMinutes = 2
)

# Create a scheduled task to run this script every $TaskIntervalMinutes minutes
if ($InstallScheduledTask) {
    $taskName = "WsusComplianceExporter"
    $existingTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue

    if (-not $existingTask) {
        $taskAction = New-ScheduledTaskAction -Execute "powershell.exe" -Argument "-NoProfile -ExecutionPolicy Bypass -File `"$($MyInvocation.MyCommand.Path)`" -Mode textfile"

        if (-not $TaskIntervalMinutes -or $TaskIntervalMinutes -le 0) {
            throw "TaskIntervalMinutes must be a positive integer"
        }

        $taskTrigger = New-ScheduledTaskTrigger -Once -At (Get-Date).AddMinutes(1) -RepetitionInterval (New-TimeSpan -Minutes $TaskIntervalMinutes) -RepetitionDuration (New-TimeSpan -Days 365)
        $taskPrincipal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -LogonType ServiceAccount -RunLevel Highest

        try {
            Write-Host "Creating scheduled task: $taskName"
            Register-ScheduledTask -TaskName $taskName -Action $taskAction -Trigger $taskTrigger -Principal $taskPrincipal -Description "Exports WSUS compliance metrics for Prometheus every $TaskIntervalMinutes minutes"

            $createdTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
            if (-not $createdTask) {
                throw "Failed to verify scheduled task creation"
            }
            Write-Host "Successfully created scheduled task: $taskName" -ForegroundColor Green
        } catch {
            Write-Error "Failed to create scheduled task: $($_.Exception.Message)"
            throw
        }
    } else {
        Write-Host "Scheduled task '$taskName' already exists, skipping creation"
    }
}

$ErrorActionPreference = 'SilentlyContinue'

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

function Get-UnixTimestamp {
    [int][double]::Parse((Get-Date -UFormat '%s'))
}

function Format-MetricValue {
    param([double]$Value, [int]$Decimals = 2)
    [math]::Round($Value, $Decimals)
}

# ============================================================================
# WSUS METRICS
# ============================================================================

function Get-WsusMetrics {
    $sb = [System.Text.StringBuilder]::new()

    # Connect to WSUS
    $wsus = $null
    try {
        $wsus = Get-WsusServer -ErrorAction Stop
    } catch {
        Write-Warning "Failed to connect to WSUS server: $_"
    }

    # --- wsus_up ---
    [void]$sb.AppendLine('# HELP wsus_up WSUS server connection status (1=connected, 0=down)')
    [void]$sb.AppendLine('# TYPE wsus_up gauge')
    $upVal = if ($wsus) { 1 } else { 0 }
    [void]$sb.AppendLine("wsus_up $upVal")
    [void]$sb.AppendLine('')

    # --- wsus_exporter_info ---
    [void]$sb.AppendLine('# HELP wsus_exporter_info Exporter version information')
    [void]$sb.AppendLine('# TYPE wsus_exporter_info gauge')
    [void]$sb.AppendLine('wsus_exporter_info{version="1.0"} 1')
    [void]$sb.AppendLine('')

    if (-not $wsus) {
        return $sb.ToString()
    }

    # --- Computer metrics ---
    $computers = @()
    try {
        $computers = $wsus.GetComputerTargets()
    } catch {}

    [void]$sb.AppendLine('# HELP wsus_computers_total Total number of computer targets')
    [void]$sb.AppendLine('# TYPE wsus_computers_total gauge')
    [void]$sb.AppendLine("wsus_computers_total $($computers.Count)")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_computers_needing_updates Computers with pending updates')
    [void]$sb.AppendLine('# TYPE wsus_computers_needing_updates gauge')
    try {
        $needingUpdates = ($computers | Where-Object { $_.GetUpdateInstallationInfoPerUpdate() | Where-Object { $_.UpdateInstallationState -eq 'NotInstalled' -or $_.UpdateInstallationState -eq 'Downloaded' } } | Select-Object -Unique | Measure-Object).Count
    } catch { $needingUpdates = 0 }
    [void]$sb.AppendLine("wsus_computers_needing_updates $needingUpdates")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_computers_with_errors Computers reporting update errors')
    [void]$sb.AppendLine('# TYPE wsus_computers_with_errors gauge')
    try {
        $withErrors = ($computers | Where-Object { $_.GetUpdateInstallationInfoPerUpdate() | Where-Object { $_.UpdateInstallationState -eq 'Failed' } } | Select-Object -Unique | Measure-Object).Count
    } catch { $withErrors = 0 }
    [void]$sb.AppendLine("wsus_computers_with_errors $withErrors")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_computers_not_contacted Computers not contacted in 30 days')
    [void]$sb.AppendLine('# TYPE wsus_computers_not_contacted gauge')
    try {
        $cutoffDate = (Get-Date).AddDays(-30)
        $notContacted = ($computers | Where-Object { $_.LastReportedStatusTime -lt $cutoffDate } | Measure-Object).Count
    } catch { $notContacted = 0 }
    [void]$sb.AppendLine("wsus_computers_not_contacted $notContacted")
    [void]$sb.AppendLine('')

    # --- Update metrics ---
    $allUpdates = @()
    try {
        $allUpdates = $wsus.GetUpdates()
    } catch {}

    [void]$sb.AppendLine('# HELP wsus_updates_approved_total Number of approved updates')
    [void]$sb.AppendLine('# TYPE wsus_updates_approved_total gauge')
    try {
        $approvedCount = ($allUpdates | Where-Object { $_.IsApproved -eq $true } | Measure-Object).Count
    } catch { $approvedCount = 0 }
    [void]$sb.AppendLine("wsus_updates_approved_total $approvedCount")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_updates_declined_total Number of declined updates')
    [void]$sb.AppendLine('# TYPE wsus_updates_declined_total gauge')
    try {
        $declinedCount = ($allUpdates | Where-Object { $_.IsDeclined -eq $true } | Measure-Object).Count
    } catch { $declinedCount = 0 }
    [void]$sb.AppendLine("wsus_updates_declined_total $declinedCount")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_updates_not_approved_total Number of updates not approved')
    [void]$sb.AppendLine('# TYPE wsus_updates_not_approved_total gauge')
    try {
        $notApprovedCount = ($allUpdates | Where-Object { $_.IsApproved -eq $false -and $_.IsDeclined -eq $false } | Measure-Object).Count
    } catch { $notApprovedCount = 0 }
    [void]$sb.AppendLine("wsus_updates_not_approved_total $notApprovedCount")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_updates_needed_total Updates needed by at least one computer')
    [void]$sb.AppendLine('# TYPE wsus_updates_needed_total gauge')
    try {
        $neededCount = ($allUpdates | Where-Object { $_.IsApproved -eq $true -and ($_.GetUpdateInstallationInfoPerComputerTarget() | Where-Object { $_.UpdateInstallationState -eq 'NotInstalled' -or $_.UpdateInstallationState -eq 'Downloaded' }) } | Measure-Object).Count
    } catch { $neededCount = 0 }
    [void]$sb.AppendLine("wsus_updates_needed_total $neededCount")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_updates_installed_total Successfully installed update instances')
    [void]$sb.AppendLine('# TYPE wsus_updates_installed_total gauge')
    try {
        $installedCount = ($allUpdates | Where-Object { $_.IsApproved -eq $true } | ForEach-Object { $_.GetUpdateInstallationInfoPerComputerTarget() } | Where-Object { $_.UpdateInstallationState -eq 'Installed' } | Measure-Object).Count
    } catch { $installedCount = 0 }
    [void]$sb.AppendLine("wsus_updates_installed_total $installedCount")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_updates_failed_total Failed update installation instances')
    [void]$sb.AppendLine('# TYPE wsus_updates_failed_total gauge')
    try {
        $failedCount = ($allUpdates | Where-Object { $_.IsApproved -eq $true } | ForEach-Object { $_.GetUpdateInstallationInfoPerComputerTarget() } | Where-Object { $_.UpdateInstallationState -eq 'Failed' } | Measure-Object).Count
    } catch { $failedCount = 0 }
    [void]$sb.AppendLine("wsus_updates_failed_total $failedCount")
    [void]$sb.AppendLine('')

    # --- Sync metrics ---
    [void]$sb.AppendLine('# HELP wsus_last_sync_timestamp Unix timestamp of last WSUS synchronization')
    [void]$sb.AppendLine('# TYPE wsus_last_sync_timestamp gauge')
    try {
        $lastSync = $wsus.GetSubscription().LastSynchronizationTime
        $syncTimestamp = [int][double]::Parse((Get-Date $lastSync -UFormat '%s'))
    } catch { $syncTimestamp = 0 }
    [void]$sb.AppendLine("wsus_last_sync_timestamp $syncTimestamp")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_sync_success Last synchronization result (1=success, 0=failure)')
    [void]$sb.AppendLine('# TYPE wsus_sync_success gauge')
    try {
        $lastSyncResult = $wsus.GetSubscription().LastSynchronizationResult
        $syncSuccess = if ($lastSyncResult -eq 'Succeeded') { 1 } else { 0 }
    } catch { $syncSuccess = 0 }
    [void]$sb.AppendLine("wsus_sync_success $syncSuccess")
    [void]$sb.AppendLine('')

    # --- Storage metrics ---
    [void]$sb.AppendLine('# HELP wsus_database_size_bytes WSUS database file size in bytes')
    [void]$sb.AppendLine('# TYPE wsus_database_size_bytes gauge')
    try {
        $dbPath = 'C:\Windows\WID\Data\SUSDB.mdf'
        if (Test-Path $dbPath) {
            $dbSize = (Get-Item $dbPath -ErrorAction Stop).Length
        } else {
            $dbSize = 0
        }
    } catch { $dbSize = 0 }
    [void]$sb.AppendLine("wsus_database_size_bytes $dbSize")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP wsus_content_size_bytes WSUS content store disk usage in bytes')
    [void]$sb.AppendLine('# TYPE wsus_content_size_bytes gauge')
    try {
        $wsusConfig = $wsus.GetConfiguration()
        $contentDir = $wsusConfig.LocalContentCachePath
        if (Test-Path $contentDir) {
            $contentSize = (Get-ChildItem -Path $contentDir -Recurse -File -ErrorAction Stop | Measure-Object -Property Length -Sum).Sum
            if (-not $contentSize) { $contentSize = 0 }
        } else {
            $contentSize = 0
        }
    } catch { $contentSize = 0 }
    [void]$sb.AppendLine("wsus_content_size_bytes $contentSize")
    [void]$sb.AppendLine('')

    # --- Computer group compliance ---
    [void]$sb.AppendLine('# HELP wsus_computer_group_compliance Compliance percentage per computer group (0-100)')
    [void]$sb.AppendLine('# TYPE wsus_computer_group_compliance gauge')
    try {
        $groups = $wsus.GetComputerTargetGroups()
        foreach ($group in $groups) {
            $groupName = $group.Name -replace '[\\"]', ''
            $groupComputers = $wsus.GetComputerTargets([Microsoft.UpdateServices.Administration.ComputerTargetScope]@{ ComputerTargetGroups = @($group) })
            $totalInGroup = $groupComputers.Count

            if ($totalInGroup -gt 0) {
                $compliantCount = 0
                foreach ($computer in $groupComputers) {
                    try {
                        $pendingUpdates = $computer.GetUpdateInstallationInfoPerUpdate() | Where-Object {
                            $_.UpdateApprovalAction -eq 'Install' -and $_.UpdateInstallationState -ne 'Installed'
                        }
                        if (-not $pendingUpdates) {
                            $compliantCount++
                        }
                    } catch {}
                }
                $compliancePct = Format-MetricValue (($compliantCount / $totalInGroup) * 100)
            } else {
                $compliancePct = 100
            }
            [void]$sb.AppendLine("wsus_computer_group_compliance{group=`"$groupName`"} $compliancePct")
        }
    } catch {}
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# COLLECT ALL METRICS
# ============================================================================

function Get-AllMetrics {
    $scriptStart = Get-Date
    $sb = [System.Text.StringBuilder]::new()

    # Collect WSUS metrics
    [void]$sb.Append((Get-WsusMetrics))

    # Exporter runtime
    $scriptEnd = Get-Date
    $duration = Format-MetricValue ($scriptEnd - $scriptStart).TotalSeconds
    $timestamp = Get-UnixTimestamp

    [void]$sb.AppendLine('# HELP wsus_exporter_duration_seconds Time to generate all metrics')
    [void]$sb.AppendLine('# TYPE wsus_exporter_duration_seconds gauge')
    [void]$sb.AppendLine("wsus_exporter_duration_seconds $duration")
    [void]$sb.AppendLine('')
    [void]$sb.AppendLine('# HELP wsus_exporter_last_run_timestamp Unix timestamp of last successful run')
    [void]$sb.AppendLine('# TYPE wsus_exporter_last_run_timestamp gauge')
    [void]$sb.AppendLine("wsus_exporter_last_run_timestamp $timestamp")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# HTTP SERVER MODE
# ============================================================================

function Start-HttpServer {
    param([int]$ListenPort)

    $prefix = "http://+:$ListenPort/"
    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add($prefix)

    try {
        $listener.Start()
        Write-Host "Starting WSUS Compliance exporter on port $ListenPort..." -ForegroundColor Green
        Write-Host "Metrics available at http://localhost:$ListenPort/metrics"

        while ($listener.IsListening) {
            $context  = $listener.GetContext()
            $request  = $context.Request
            $response = $context.Response

            if ($request.Url.AbsolutePath -eq '/metrics') {
                $metrics = Get-AllMetrics
                $buffer  = [System.Text.Encoding]::UTF8.GetBytes($metrics)
                $response.ContentType = 'text/plain; version=0.0.4; charset=utf-8'
            }
            else {
                $html = @"
<!DOCTYPE html>
<html>
<head><title>WSUS Compliance Exporter v1.0</title></head>
<body>
<h1>WSUS Compliance Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<h2>Metrics</h2>
<ul>
<li>WSUS server connection status</li>
<li>Computer targets and compliance state</li>
<li>Update approval, installation, and failure counts</li>
<li>Synchronization status and timing</li>
<li>Database and content store sizes</li>
<li>Per-group compliance percentages</li>
</ul>
</body>
</html>
"@
                $buffer = [System.Text.Encoding]::UTF8.GetBytes($html)
                $response.ContentType = 'text/html; charset=utf-8'
            }

            $response.ContentLength64 = $buffer.Length
            $response.OutputStream.Write($buffer, 0, $buffer.Length)
            $response.OutputStream.Close()
        }
    }
    catch {
        Write-Error "HTTP server error: $_"
        Write-Error "If access denied, run: netsh http add urlacl url=http://+:$ListenPort/ user=Everyone"
    }
    finally {
        if ($listener.IsListening) {
            $listener.Stop()
        }
    }
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

switch ($Mode) {
    'http' {
        Start-HttpServer -ListenPort $Port
    }
    'textfile' {
        $OutputFile = Join-Path $TextfileDir 'wsus_compliance.prom'

        $outputDir = Split-Path $OutputFile -Parent
        if (-not (Test-Path $outputDir)) {
            New-Item -Path $outputDir -ItemType Directory -Force | Out-Null
        }

        $tempFile = Join-Path $outputDir ".wsus_compliance.$PID.tmp"

        try {
            $metrics = Get-AllMetrics
            $metrics | Out-File -FilePath $tempFile -Encoding utf8 -NoNewline

            $lineCount = ($metrics -split "`n").Count
            if ($lineCount -lt 10) {
                Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
                Write-Error "Metrics file too small ($lineCount lines), keeping previous"
                exit 1
            }

            Move-Item -Path $tempFile -Destination $OutputFile -Force
            Write-Host "Metrics written to $OutputFile ($lineCount lines)" -ForegroundColor Green
        }
        catch {
            Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
            Write-Error "Failed to generate metrics: $_"
            exit 1
        }
    }
    default {
        Get-AllMetrics | Write-Output
    }
}