<#
.SYNOPSIS
    Windows Deployment Services Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for WDS - active deployments, completed/failed jobs,
    image counts, multicast sessions, PXE requests, storage usage. 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:
        - wds_up
        - wds_exporter_info{version}

      Images:
        - wds_boot_images_total
        - wds_install_images_total
        - wds_image_groups_total

      Deployments:
        - wds_deployments_active
        - wds_deployments_completed_total
        - wds_deployments_failed_total

      Multicast:
        - wds_multicast_transmissions_total
        - wds_multicast_sessions_active

      PXE:
        - wds_pxe_requests_total

      Storage:
        - wds_image_store_size_bytes

      Devices:
        - wds_pending_devices

      Exporter:
        - wds_exporter_duration_seconds
        - wds_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 = "WdsMetricsExporter"
    $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 WDS 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)
}

# ============================================================================
# WDS METRICS
# ============================================================================

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

    # --- wds_boot_images_total ---
    [void]$sb.AppendLine('# HELP wds_boot_images_total Number of boot images available in WDS')
    [void]$sb.AppendLine('# TYPE wds_boot_images_total gauge')
    try {
        $bootOutput = & WDSUTIL /Get-AllImages /ImageType:Boot 2>&1
        $bootCount = ($bootOutput | Select-String -Pattern 'Image name:' | Measure-Object).Count
        [void]$sb.AppendLine("wds_boot_images_total $bootCount")
    } catch {
        [void]$sb.AppendLine("wds_boot_images_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_install_images_total ---
    [void]$sb.AppendLine('# HELP wds_install_images_total Number of install images available in WDS')
    [void]$sb.AppendLine('# TYPE wds_install_images_total gauge')
    try {
        $installOutput = & WDSUTIL /Get-AllImages /ImageType:Install 2>&1
        $installCount = ($installOutput | Select-String -Pattern 'Image name:' | Measure-Object).Count
        [void]$sb.AppendLine("wds_install_images_total $installCount")
    } catch {
        [void]$sb.AppendLine("wds_install_images_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_image_groups_total ---
    [void]$sb.AppendLine('# HELP wds_image_groups_total Number of image groups configured in WDS')
    [void]$sb.AppendLine('# TYPE wds_image_groups_total gauge')
    try {
        $groupOutput = & WDSUTIL /Get-AllImageGroups 2>&1
        $groupCount = ($groupOutput | Select-String -Pattern 'Group name:' | Measure-Object).Count
        [void]$sb.AppendLine("wds_image_groups_total $groupCount")
    } catch {
        [void]$sb.AppendLine("wds_image_groups_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_deployments_active ---
    [void]$sb.AppendLine('# HELP wds_deployments_active Number of currently active WDS deployments')
    [void]$sb.AppendLine('# TYPE wds_deployments_active gauge')
    try {
        $activeClients = Get-CimInstance -Namespace 'root\cimv2' -ClassName 'MSFT_WdsClient' -Filter "Status='Active'" -ErrorAction Stop
        $activeCount = ($activeClients | Measure-Object).Count
        [void]$sb.AppendLine("wds_deployments_active $activeCount")
    } catch {
        [void]$sb.AppendLine("wds_deployments_active 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_deployments_completed_total ---
    [void]$sb.AppendLine('# HELP wds_deployments_completed_total Total number of completed WDS deployments')
    [void]$sb.AppendLine('# TYPE wds_deployments_completed_total counter')
    try {
        $completedEvents = Get-WinEvent -FilterHashtable @{
            LogName   = 'Microsoft-Windows-Deployment-Services-Diagnostics/Operational'
            Id        = 257
        } -ErrorAction Stop
        $completedCount = ($completedEvents | Measure-Object).Count
        [void]$sb.AppendLine("wds_deployments_completed_total $completedCount")
    } catch {
        [void]$sb.AppendLine("wds_deployments_completed_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_deployments_failed_total ---
    [void]$sb.AppendLine('# HELP wds_deployments_failed_total Total number of failed WDS deployments')
    [void]$sb.AppendLine('# TYPE wds_deployments_failed_total counter')
    try {
        $failedEvents = Get-WinEvent -FilterHashtable @{
            LogName   = 'Microsoft-Windows-Deployment-Services-Diagnostics/Operational'
            Id        = 258
        } -ErrorAction Stop
        $failedCount = ($failedEvents | Measure-Object).Count
        [void]$sb.AppendLine("wds_deployments_failed_total $failedCount")
    } catch {
        [void]$sb.AppendLine("wds_deployments_failed_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_multicast_transmissions_total ---
    [void]$sb.AppendLine('# HELP wds_multicast_transmissions_total Total number of multicast transmissions configured')
    [void]$sb.AppendLine('# TYPE wds_multicast_transmissions_total gauge')
    try {
        $mcastOutput = & WDSUTIL /Get-AllMulticastTransmissions 2>&1
        $mcastCount = ($mcastOutput | Select-String -Pattern 'Transmission name:' | Measure-Object).Count
        [void]$sb.AppendLine("wds_multicast_transmissions_total $mcastCount")
    } catch {
        [void]$sb.AppendLine("wds_multicast_transmissions_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_multicast_sessions_active ---
    [void]$sb.AppendLine('# HELP wds_multicast_sessions_active Number of active multicast sessions with connected clients')
    [void]$sb.AppendLine('# TYPE wds_multicast_sessions_active gauge')
    try {
        $mcastDetailOutput = & WDSUTIL /Get-AllMulticastTransmissions /Show:Clients 2>&1
        $activeSessionCount = ($mcastDetailOutput | Select-String -Pattern 'Client count:' | ForEach-Object {
            if ($_ -match 'Client count:\s+(\d+)') { [int]$Matches[1] }
        } | Where-Object { $_ -gt 0 } | Measure-Object).Count
        [void]$sb.AppendLine("wds_multicast_sessions_active $activeSessionCount")
    } catch {
        [void]$sb.AppendLine("wds_multicast_sessions_active 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_pxe_requests_total ---
    [void]$sb.AppendLine('# HELP wds_pxe_requests_total Total number of PXE boot requests received')
    [void]$sb.AppendLine('# TYPE wds_pxe_requests_total counter')
    try {
        $pxeEvents = Get-WinEvent -FilterHashtable @{
            LogName   = 'Microsoft-Windows-Deployment-Services-Diagnostics/Operational'
            Id        = 513
        } -ErrorAction Stop
        $pxeCount = ($pxeEvents | Measure-Object).Count
        [void]$sb.AppendLine("wds_pxe_requests_total $pxeCount")
    } catch {
        [void]$sb.AppendLine("wds_pxe_requests_total 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_image_store_size_bytes ---
    [void]$sb.AppendLine('# HELP wds_image_store_size_bytes Total size of the WDS image store in bytes')
    [void]$sb.AppendLine('# TYPE wds_image_store_size_bytes gauge')
    try {
        $wdsConfig = & WDSUTIL /Get-Server /Show:Config 2>&1
        $remInstallPath = ($wdsConfig | Select-String -Pattern 'RemoteInstall location:\s+(.+)' | ForEach-Object { $_.Matches[0].Groups[1].Value.Trim() })
        if (-not $remInstallPath) {
            $remInstallPath = 'C:\RemoteInstall'
        }
        $imagesPath = Join-Path $remInstallPath 'Images'
        $storeSize = (Get-ChildItem -Path $imagesPath -Recurse -File -ErrorAction Stop | Measure-Object -Property Length -Sum).Sum
        if (-not $storeSize) { $storeSize = 0 }
        [void]$sb.AppendLine("wds_image_store_size_bytes $storeSize")
    } catch {
        [void]$sb.AppendLine("wds_image_store_size_bytes 0")
    }
    [void]$sb.AppendLine('')

    # --- wds_pending_devices ---
    [void]$sb.AppendLine('# HELP wds_pending_devices Number of devices pending approval in WDS')
    [void]$sb.AppendLine('# TYPE wds_pending_devices gauge')
    try {
        $pendingOutput = & WDSUTIL /Get-AllDevices /DeviceType:PendingDevices 2>&1
        $pendingCount = ($pendingOutput | Select-String -Pattern 'Device name:' | Measure-Object).Count
        [void]$sb.AppendLine("wds_pending_devices $pendingCount")
    } catch {
        [void]$sb.AppendLine("wds_pending_devices 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 wds_up WDS service status (1=running, 0=stopped)')
    [void]$sb.AppendLine('# TYPE wds_up gauge')
    try {
        $wdsSvc = Get-Service -Name WDSServer -ErrorAction Stop
        $upVal = if ($wdsSvc.Status -eq 'Running') { 1 } else { 0 }
        [void]$sb.AppendLine("wds_up $upVal")
    } catch {
        [void]$sb.AppendLine("wds_up 0")
    }
    [void]$sb.AppendLine('')

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

    # Collect WDS metrics
    [void]$sb.Append((Get-WdsMetrics))

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

    [void]$sb.AppendLine('# HELP wds_exporter_duration_seconds Time to generate all metrics')
    [void]$sb.AppendLine('# TYPE wds_exporter_duration_seconds gauge')
    [void]$sb.AppendLine("wds_exporter_duration_seconds $duration")
    [void]$sb.AppendLine('')
    [void]$sb.AppendLine('# HELP wds_exporter_last_run_timestamp Unix timestamp of last successful run')
    [void]$sb.AppendLine('# TYPE wds_exporter_last_run_timestamp gauge')
    [void]$sb.AppendLine("wds_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 WDS metrics 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>WDS Metrics Exporter v1.0</title></head>
<body>
<h1>WDS Metrics Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<h2>Metrics</h2>
<ul>
<li>Boot and install image counts</li>
<li>Image groups</li>
<li>Active, completed, and failed deployments</li>
<li>Multicast transmissions and active sessions</li>
<li>PXE boot requests</li>
<li>Image store size</li>
<li>Pending devices</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 'wds_metrics.prom'

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

        $tempFile = Join-Path $outputDir ".wds_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
    }
}