<#
.SYNOPSIS
    Windows Update Compliance Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for Windows Update compliance - pending updates,
    installed update counts, reboot required status, last update timestamp,
    update categories, WSUS configuration, and Windows Update Agent version.
    Exports metrics as Prometheus-compatible text format.
.PARAMETER Mode
    Output mode: 'stdout' (default), 'textfile', or 'http'
.PARAMETER Port
    HTTP port for http mode (default: 9536)
.PARAMETER TextfileDir
    Directory for textfile collector output (default: C:\ProgramData\node_exporter)
.PARAMETER InstallScheduledTask
    Switch to create a scheduled task for automatic execution
.PARAMETER TaskIntervalMinutes
    Interval in minutes for the scheduled task (default: 60)
.NOTES
    Author:  Phil Connor
    Contact: contact@mylinux.work
    Website: https://mylinux.work
    License: MIT
    Version: 1.0

    Metrics Exported:
      Core Status:
        - windows_update_up
        - windows_update_exporter_info{version}

      Update Counts:
        - windows_update_pending_total
        - windows_update_pending_critical_total
        - windows_update_pending_security_total
        - windows_update_pending_optional_total
        - windows_update_installed_total
        - windows_update_hidden_total

      Update State:
        - windows_update_reboot_required
        - windows_update_last_install_timestamp
        - windows_update_last_search_timestamp
        - windows_update_days_since_last_install

      WSUS:
        - windows_update_wsus_configured
        - windows_update_auto_update_enabled

      Exporter:
        - windows_update_exporter_duration_seconds
        - windows_update_exporter_last_run_timestamp
#>

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

    [int]$Port = 9536,

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

    [switch]$InstallScheduledTask,

    [int]$TaskIntervalMinutes = 60
)

if ($InstallScheduledTask) {
    $taskName = "WindowsUpdateComplianceExporter"
    $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 Windows Update 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)
}

# ============================================================================
# UPDATE METRICS
# ============================================================================

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

    # --- Pending updates ---
    try {
        $session = New-Object -ComObject Microsoft.Update.Session
        $searcher = $session.CreateUpdateSearcher()

        $searchResult = $searcher.Search("IsInstalled=0 AND IsHidden=0")
        $pendingUpdates = $searchResult.Updates

        $pendingTotal = $pendingUpdates.Count
        $pendingCritical = 0
        $pendingSecurity = 0
        $pendingOptional = 0

        foreach ($update in $pendingUpdates) {
            foreach ($category in $update.Categories) {
                switch ($category.Name) {
                    'Critical Updates'  { $pendingCritical++ }
                    'Security Updates'  { $pendingSecurity++ }
                }
            }
            if ($update.BrowseOnly) { $pendingOptional++ }
        }

        [void]$sb.AppendLine('# HELP windows_update_pending_total Total pending updates not yet installed')
        [void]$sb.AppendLine('# TYPE windows_update_pending_total gauge')
        [void]$sb.AppendLine("windows_update_pending_total $pendingTotal")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_update_pending_critical_total Pending critical updates')
        [void]$sb.AppendLine('# TYPE windows_update_pending_critical_total gauge')
        [void]$sb.AppendLine("windows_update_pending_critical_total $pendingCritical")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_update_pending_security_total Pending security updates')
        [void]$sb.AppendLine('# TYPE windows_update_pending_security_total gauge')
        [void]$sb.AppendLine("windows_update_pending_security_total $pendingSecurity")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_update_pending_optional_total Pending optional updates')
        [void]$sb.AppendLine('# TYPE windows_update_pending_optional_total gauge')
        [void]$sb.AppendLine("windows_update_pending_optional_total $pendingOptional")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_pending_total Total pending updates not yet installed')
        [void]$sb.AppendLine('# TYPE windows_update_pending_total gauge')
        [void]$sb.AppendLine("windows_update_pending_total 0")
        [void]$sb.AppendLine('')
    }

    # --- Installed updates count ---
    try {
        $installedResult = $searcher.Search("IsInstalled=1")
        $installedTotal = $installedResult.Updates.Count

        [void]$sb.AppendLine('# HELP windows_update_installed_total Total installed updates')
        [void]$sb.AppendLine('# TYPE windows_update_installed_total gauge')
        [void]$sb.AppendLine("windows_update_installed_total $installedTotal")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_installed_total Total installed updates')
        [void]$sb.AppendLine('# TYPE windows_update_installed_total gauge')
        [void]$sb.AppendLine("windows_update_installed_total 0")
        [void]$sb.AppendLine('')
    }

    # --- Hidden updates ---
    try {
        $hiddenResult = $searcher.Search("IsHidden=1")
        $hiddenTotal = $hiddenResult.Updates.Count

        [void]$sb.AppendLine('# HELP windows_update_hidden_total Total hidden (declined) updates')
        [void]$sb.AppendLine('# TYPE windows_update_hidden_total gauge')
        [void]$sb.AppendLine("windows_update_hidden_total $hiddenTotal")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_hidden_total Total hidden (declined) updates')
        [void]$sb.AppendLine('# TYPE windows_update_hidden_total gauge')
        [void]$sb.AppendLine("windows_update_hidden_total 0")
        [void]$sb.AppendLine('')
    }

    # --- Reboot required ---
    try {
        $rebootRequired = 0
        if (Test-Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\WindowsUpdate\Auto Update\RebootRequired') {
            $rebootRequired = 1
        }
        if (Test-Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Component Based Servicing\RebootPending') {
            $rebootRequired = 1
        }

        [void]$sb.AppendLine('# HELP windows_update_reboot_required Whether a reboot is required to complete updates (1=yes, 0=no)')
        [void]$sb.AppendLine('# TYPE windows_update_reboot_required gauge')
        [void]$sb.AppendLine("windows_update_reboot_required $rebootRequired")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_reboot_required Whether a reboot is required to complete updates (1=yes, 0=no)')
        [void]$sb.AppendLine('# TYPE windows_update_reboot_required gauge')
        [void]$sb.AppendLine("windows_update_reboot_required 0")
        [void]$sb.AppendLine('')
    }

    # --- Last install and search timestamps ---
    try {
        $autoUpdate = New-Object -ComObject Microsoft.Update.AutoUpdate
        $results = $autoUpdate.Results

        $lastInstall = 0
        $lastSearch = 0
        $daysSinceInstall = -1

        if ($results.LastInstallationSuccessDate) {
            $lastInstallDate = [DateTime]$results.LastInstallationSuccessDate
            $lastInstall = [int][double]::Parse(($lastInstallDate).ToUniversalTime().Subtract([DateTime]'1970-01-01').TotalSeconds.ToString())
            $daysSinceInstall = Format-MetricValue ((Get-Date) - $lastInstallDate).TotalDays 0
        }

        if ($results.LastSearchSuccessDate) {
            $lastSearchDate = [DateTime]$results.LastSearchSuccessDate
            $lastSearch = [int][double]::Parse(($lastSearchDate).ToUniversalTime().Subtract([DateTime]'1970-01-01').TotalSeconds.ToString())
        }

        [void]$sb.AppendLine('# HELP windows_update_last_install_timestamp Unix timestamp of last successful update installation')
        [void]$sb.AppendLine('# TYPE windows_update_last_install_timestamp gauge')
        [void]$sb.AppendLine("windows_update_last_install_timestamp $lastInstall")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_update_last_search_timestamp Unix timestamp of last successful update search')
        [void]$sb.AppendLine('# TYPE windows_update_last_search_timestamp gauge')
        [void]$sb.AppendLine("windows_update_last_search_timestamp $lastSearch")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_update_days_since_last_install Days since last successful update installation')
        [void]$sb.AppendLine('# TYPE windows_update_days_since_last_install gauge')
        [void]$sb.AppendLine("windows_update_days_since_last_install $daysSinceInstall")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_last_install_timestamp Unix timestamp of last successful update installation')
        [void]$sb.AppendLine('# TYPE windows_update_last_install_timestamp gauge')
        [void]$sb.AppendLine("windows_update_last_install_timestamp 0")
        [void]$sb.AppendLine('')
    }

    # --- WSUS configuration ---
    try {
        $wsusConfigured = 0
        $wsusReg = Get-ItemProperty -Path 'HKLM:\SOFTWARE\Policies\Microsoft\Windows\WindowsUpdate' -ErrorAction SilentlyContinue
        if ($wsusReg -and $wsusReg.WUServer) {
            $wsusConfigured = 1
        }

        [void]$sb.AppendLine('# HELP windows_update_wsus_configured Whether WSUS is configured (1=yes, 0=no)')
        [void]$sb.AppendLine('# TYPE windows_update_wsus_configured gauge')
        [void]$sb.AppendLine("windows_update_wsus_configured $wsusConfigured")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_wsus_configured Whether WSUS is configured (1=yes, 0=no)')
        [void]$sb.AppendLine('# TYPE windows_update_wsus_configured gauge')
        [void]$sb.AppendLine("windows_update_wsus_configured 0")
        [void]$sb.AppendLine('')
    }

    # --- Auto update enabled ---
    try {
        $autoUpdateEnabled = 0
        $auReg = Get-ItemProperty -Path 'HKLM:\SOFTWARE\Policies\Microsoft\Windows\WindowsUpdate\AU' -ErrorAction SilentlyContinue
        if ($auReg -and $auReg.NoAutoUpdate -eq 0) {
            $autoUpdateEnabled = 1
        } elseif (-not $auReg) {
            $autoUpdateEnabled = 1
        }

        [void]$sb.AppendLine('# HELP windows_update_auto_update_enabled Whether automatic updates are enabled (1=yes, 0=no)')
        [void]$sb.AppendLine('# TYPE windows_update_auto_update_enabled gauge')
        [void]$sb.AppendLine("windows_update_auto_update_enabled $autoUpdateEnabled")
        [void]$sb.AppendLine('')
    } catch {
        [void]$sb.AppendLine('# HELP windows_update_auto_update_enabled Whether automatic updates are enabled (1=yes, 0=no)')
        [void]$sb.AppendLine('# TYPE windows_update_auto_update_enabled gauge')
        [void]$sb.AppendLine("windows_update_auto_update_enabled 0")
        [void]$sb.AppendLine('')
    }

    $sb.ToString()
}

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

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

    # Exporter up
    [void]$sb.AppendLine('# HELP windows_update_up Exporter status (1=up, 0=down)')
    [void]$sb.AppendLine('# TYPE windows_update_up gauge')
    [void]$sb.AppendLine('windows_update_up 1')
    [void]$sb.AppendLine('')

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

    # Collect update metrics
    [void]$sb.Append((Get-UpdateMetrics))

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

    [void]$sb.AppendLine('# HELP windows_update_exporter_duration_seconds Time to generate all metrics')
    [void]$sb.AppendLine('# TYPE windows_update_exporter_duration_seconds gauge')
    [void]$sb.AppendLine("windows_update_exporter_duration_seconds $duration")
    [void]$sb.AppendLine('')
    [void]$sb.AppendLine('# HELP windows_update_exporter_last_run_timestamp Unix timestamp of last successful run')
    [void]$sb.AppendLine('# TYPE windows_update_exporter_last_run_timestamp gauge')
    [void]$sb.AppendLine("windows_update_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 Windows Update 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>Windows Update Compliance Exporter v1.0</title></head>
<body>
<h1>Windows Update Compliance Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<h2>Metrics</h2>
<ul>
<li>Pending updates (total, critical, security, optional)</li>
<li>Installed and hidden update counts</li>
<li>Reboot required status</li>
<li>Last install/search timestamps</li>
<li>WSUS and auto-update configuration</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 'windows_update_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 ".windows_update_compliance_metrics.$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
    }
}