<#
.SYNOPSIS
    Windows Remote Desktop Services Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for Windows RDS. Monitors active/disconnected sessions,
    per-user session details, license availability, session host resources,
    Connection Broker status, and RD Gateway metrics.
.PARAMETER TextFile
    Write to windows_exporter textfile directory
.PARAMETER OutFile
    Write metrics to a specific file path
.PARAMETER Install
    Create a scheduled task for automatic collection
.PARAMETER Listen
    Start HTTP listener on specified address:port
.PARAMETER Interval
    Collection interval in seconds for scheduled task (default: 120)
.NOTES
    Author:  Phil Connor
    Contact: contact@mylinux.work
    Website: https://mylinux.work
    License: MIT
    Version: 1.0
#>

param(
    [switch]$TextFile,
    [string]$OutFile = "",
    [switch]$Install,
    [string]$Listen = "",
    [int]$Interval = 120
)

$ErrorActionPreference = "SilentlyContinue"
$Version = "1.0"
$TextfileDir = "C:\ProgramData\node_exporter"
$MetricPrefix = "rds"

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

function Get-PrometheusEscape {
    param([string]$Value)
    $Value -replace '\\', '\\' -replace '"', '\"' -replace "`n", ''
}

function Write-MetricHeader {
    param([string]$Name, [string]$Help, [string]$Type)
    "# HELP $Name $Help"
    "# TYPE $Name $Type"
}

# ============================================================================
# METRIC GENERATION
# ============================================================================

function Get-RdsMetrics {
    $startTime = Get-Date
    $metrics = [System.Collections.ArrayList]::new()

    # Check if RDS module is available
    $rdsAvailable = $false
    if (Get-Module -ListAvailable -Name "RemoteDesktop" -ErrorAction SilentlyContinue) {
        Import-Module RemoteDesktop -ErrorAction SilentlyContinue
        $rdsAvailable = $true
    }

    $up = if ($rdsAvailable) { 1 } else { 0 }

    [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_up" "RDS exporter status (1=up, 0=down)" "gauge"))
    [void]$metrics.Add("${MetricPrefix}_up $up")
    [void]$metrics.Add("")

    [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_exporter_info" "Exporter version information" "gauge"))
    [void]$metrics.Add("${MetricPrefix}_exporter_info{version=`"$Version`"} 1")
    [void]$metrics.Add("")

    if ($up -eq 0) {
        $endTime = Get-Date
        $duration = [math]::Round(($endTime - $startTime).TotalSeconds, 2)
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_exporter_duration_seconds" "Script execution time" "gauge"))
        [void]$metrics.Add("${MetricPrefix}_exporter_duration_seconds $duration")
        return ($metrics -join "`n")
    }

    # ========================================================================
    # SESSION METRICS
    # ========================================================================

    try {
        $sessions = Get-RDUserSession -ErrorAction Stop
        $sessionHosts = $sessions | Group-Object -Property HostServer

        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_sessions_active" "Active RDS sessions per host" "gauge"))
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_sessions_disconnected" "Disconnected RDS sessions per host" "gauge"))
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_sessions_total" "Total RDS sessions per host" "gauge"))

        $totalActive = 0
        $totalDisconnected = 0

        foreach ($hostGroup in $sessionHosts) {
            $host = Get-PrometheusEscape $hostGroup.Name
            $active = @($hostGroup.Group | Where-Object { $_.SessionState -eq "STATE_ACTIVE" }).Count
            $disconnected = @($hostGroup.Group | Where-Object { $_.SessionState -eq "STATE_DISCONNECTED" }).Count
            $total = $hostGroup.Count

            [void]$metrics.Add("${MetricPrefix}_sessions_active{host=`"$host`"} $active")
            [void]$metrics.Add("${MetricPrefix}_sessions_disconnected{host=`"$host`"} $disconnected")
            [void]$metrics.Add("${MetricPrefix}_sessions_total{host=`"$host`"} $total")

            $totalActive += $active
            $totalDisconnected += $disconnected
        }
        [void]$metrics.Add("")

        # Per-user session details
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_session_state" "Per-user session state (1=current)" "gauge"))
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_session_idle_seconds" "Session idle time in seconds" "gauge"))

        foreach ($session in $sessions) {
            $host = Get-PrometheusEscape $session.HostServer
            $user = Get-PrometheusEscape $session.UserName
            $state = $session.SessionState -replace "STATE_", ""

            [void]$metrics.Add("${MetricPrefix}_session_state{host=`"$host`",user=`"$user`",state=`"$state`"} 1")

            if ($session.IdleTime) {
                $idleSeconds = [math]::Round($session.IdleTime.TotalSeconds)
                [void]$metrics.Add("${MetricPrefix}_session_idle_seconds{host=`"$host`",user=`"$user`"} $idleSeconds")
            }
        }
        [void]$metrics.Add("")

    } catch {
        [void]$metrics.Add("# INFO: Could not query RDS sessions: $_")
        [void]$metrics.Add("")
    }

    # ========================================================================
    # LICENSE METRICS
    # ========================================================================

    try {
        $licenseServers = Get-RDLicenseConfiguration -ErrorAction Stop

        if ($licenseServers) {
            $licenses = Get-WmiObject -Class Win32_TSLicenseKeyPack -ErrorAction Stop

            [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_license_total" "Total RDS licenses by type" "gauge"))
            [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_license_available" "Available RDS licenses by type" "gauge"))
            [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_license_used" "Used RDS licenses by type" "gauge"))

            foreach ($lic in $licenses) {
                $type = switch ($lic.ProductVersionID) {
                    0 { "PerDevice" }
                    1 { "PerUser" }
                    2 { "PerDevice" }
                    default { "Unknown" }
                }

                $total = $lic.TotalLicenses
                $available = $lic.AvailableLicenses
                $used = $total - $available

                [void]$metrics.Add("${MetricPrefix}_license_total{type=`"$type`"} $total")
                [void]$metrics.Add("${MetricPrefix}_license_available{type=`"$type`"} $available")
                [void]$metrics.Add("${MetricPrefix}_license_used{type=`"$type`"} $used")
            }
            [void]$metrics.Add("")
        }
    } catch {
        # License server may not be on this host
    }

    # ========================================================================
    # SESSION HOST RESOURCE METRICS
    # ========================================================================

    try {
        $rdServers = Get-RDServer -Role "RDS-RD-SERVER" -ErrorAction Stop

        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_host_cpu_percent" "Session host CPU usage" "gauge"))
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_host_memory_percent" "Session host memory usage" "gauge"))
        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_host_disk_free_bytes" "Session host free disk space" "gauge"))

        foreach ($server in $rdServers) {
            $hostName = Get-PrometheusEscape $server.Server

            $cpu = Get-Counter -ComputerName $server.Server -Counter "\Processor(_Total)\% Processor Time" -ErrorAction SilentlyContinue
            if ($cpu) {
                $cpuVal = [math]::Round($cpu.CounterSamples[0].CookedValue, 1)
                [void]$metrics.Add("${MetricPrefix}_host_cpu_percent{host=`"$hostName`"} $cpuVal")
            }

            $mem = Get-WmiObject -ComputerName $server.Server -Class Win32_OperatingSystem -ErrorAction SilentlyContinue
            if ($mem) {
                $memPercent = [math]::Round((($mem.TotalVisibleMemorySize - $mem.FreePhysicalMemory) / $mem.TotalVisibleMemorySize) * 100, 1)
                [void]$metrics.Add("${MetricPrefix}_host_memory_percent{host=`"$hostName`"} $memPercent")
            }

            $disks = Get-WmiObject -ComputerName $server.Server -Class Win32_LogicalDisk -Filter "DriveType=3" -ErrorAction SilentlyContinue
            foreach ($disk in $disks) {
                $drive = $disk.DeviceID -replace ":", ""
                [void]$metrics.Add("${MetricPrefix}_host_disk_free_bytes{host=`"$hostName`",drive=`"$drive`"} $($disk.FreeSpace)")
            }
        }
        [void]$metrics.Add("")
    } catch {
        # May not have access to remote servers
    }

    # ========================================================================
    # CONNECTION BROKER
    # ========================================================================

    $brokerSvc = Get-Service -Name "Tssdis" -ErrorAction SilentlyContinue
    $brokerUp = if ($brokerSvc -and $brokerSvc.Status -eq "Running") { 1 } else { 0 }

    [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_connection_broker_up" "Connection Broker service status" "gauge"))
    [void]$metrics.Add("${MetricPrefix}_connection_broker_up $brokerUp")
    [void]$metrics.Add("")

    # Collection health
    try {
        $collections = Get-RDSessionCollection -ErrorAction Stop

        [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_collection_health" "Collection health (1=healthy)" "gauge"))
        foreach ($col in $collections) {
            $colName = Get-PrometheusEscape $col.CollectionName
            $health = 1
            [void]$metrics.Add("${MetricPrefix}_collection_health{collection=`"$colName`"} $health")
        }
        [void]$metrics.Add("")
    } catch {
        # Collections may not be configured
    }

    # ========================================================================
    # RD GATEWAY METRICS
    # ========================================================================

    try {
        $gwCounter = Get-Counter -Counter "\TS Gateway\Current Connections" -ErrorAction Stop
        if ($gwCounter) {
            $gwActive = [math]::Round($gwCounter.CounterSamples[0].CookedValue)
            [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_gateway_connections_active" "Active RD Gateway connections" "gauge"))
            [void]$metrics.Add("${MetricPrefix}_gateway_connections_active $gwActive")
            [void]$metrics.Add("")
        }

        $gwRate = Get-Counter -Counter "\TS Gateway\Connections/Sec" -ErrorAction SilentlyContinue
        if ($gwRate) {
            $gwPerSec = [math]::Round($gwRate.CounterSamples[0].CookedValue, 2)
            [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_gateway_connections_per_second" "RD Gateway connection rate" "gauge"))
            [void]$metrics.Add("${MetricPrefix}_gateway_connections_per_second $gwPerSec")
            [void]$metrics.Add("")
        }

        $gwAuth = Get-Counter -Counter "\TS Gateway\Failed Authentications" -ErrorAction SilentlyContinue
        if ($gwAuth) {
            $gwFailed = [math]::Round($gwAuth.CounterSamples[0].CookedValue)
            [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_gateway_auth_failures_total" "RD Gateway authentication failures" "counter"))
            [void]$metrics.Add("${MetricPrefix}_gateway_auth_failures_total $gwFailed")
            [void]$metrics.Add("")
        }
    } catch {
        # Gateway may not be installed
    }

    # ========================================================================
    # EXPORTER RUNTIME
    # ========================================================================

    $endTime = Get-Date
    $duration = [math]::Round(($endTime - $startTime).TotalSeconds, 2)
    $timestamp = [math]::Round((Get-Date -UFormat %s), 0)

    [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_exporter_duration_seconds" "Script execution time" "gauge"))
    [void]$metrics.Add("${MetricPrefix}_exporter_duration_seconds $duration")
    [void]$metrics.Add("")

    [void]$metrics.Add((Write-MetricHeader "${MetricPrefix}_exporter_last_run_timestamp" "Unix timestamp of last run" "gauge"))
    [void]$metrics.Add("${MetricPrefix}_exporter_last_run_timestamp $timestamp")

    return ($metrics -join "`n")
}

# ============================================================================
# OUTPUT MODES
# ============================================================================

if ($Install) {
    $scriptPath = $MyInvocation.MyCommand.Path
    $action = New-ScheduledTaskAction -Execute "powershell.exe" `
        -Argument "-NoProfile -ExecutionPolicy Bypass -File `"$scriptPath`" -TextFile"
    $trigger = New-ScheduledTaskTrigger -RepetitionInterval (New-TimeSpan -Seconds $Interval) -Once -At (Get-Date)
    $principal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -LogonType ServiceAccount -RunLevel Highest
    $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -StartWhenAvailable
    Register-ScheduledTask -TaskName "RDS Metrics Exporter" -Action $action -Trigger $trigger -Principal $principal -Settings $settings -Force
    Write-Host "Scheduled task 'RDS Metrics Exporter' created (interval: ${Interval}s)"
    exit 0
}

if ($Listen -ne "") {
    $port = $Listen -replace '.*:', ''
    if (-not $port) { $port = "9185" }
    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add("http://+:$port/")
    $listener.Start()
    Write-Host "RDS exporter listening on port $port..."

    while ($listener.IsListening) {
        $context = $listener.GetContext()
        $response = $context.Response
        $output = Get-RdsMetrics
        $buffer = [System.Text.Encoding]::UTF8.GetBytes($output)
        $response.ContentType = "text/plain; version=0.0.4"
        $response.ContentLength64 = $buffer.Length
        $response.OutputStream.Write($buffer, 0, $buffer.Length)
        $response.Close()
    }
} elseif ($TextFile -or $OutFile -ne "") {
    $outputPath = if ($OutFile -ne "") { $OutFile } else { Join-Path $TextfileDir "rds.prom" }
    $outputDir = Split-Path $outputPath -Parent
    if (-not (Test-Path $outputDir)) { New-Item -ItemType Directory -Path $outputDir -Force | Out-Null }
    $tempFile = Join-Path $outputDir ".rds-metrics.tmp"
    $metricsOutput = Get-RdsMetrics
    $metricsOutput | Out-File -FilePath $tempFile -Encoding utf8 -NoNewline
    Move-Item -Path $tempFile -Destination $outputPath -Force
    Write-Host "Metrics written to $outputPath"
} else {
    Get-RdsMetrics
}