Function Start-AdfsServerTrace {
    param (
        [string]$ActivityId,
        [switch]$IncludeDebug,
        [string[]]$ComputerName
    )
    
    $jobs = @()
    foreach ($computer in $ComputerName) {
        $jobs += Start-Job -ScriptBlock {
            param ($ActivityId, $IncludeDebug, $computer)
            
            $logNames = @("Security", "Application", "System")
            if ($IncludeDebug) { $logNames += "AD FS Tracing/Debug" }
            
            $events = @()
            foreach ($log in $logNames) {
                try {
                    $events += Get-WinEvent -ComputerName $computer -LogName $log -ErrorAction SilentlyContinue | 
                                Where-Object { $_.ActivityId -eq $ActivityId }
                } catch {
                    Write-Error "Failed to retrieve events from $computer for log $log"
                }
            }
            
            return [PSCustomObject]@{
                ComputerName = $computer
                Events = $events
            }
        } -ArgumentList $ActivityId, $IncludeDebug, $computer
    }
    
    return $jobs
}

Function Receive-AdfsServerTrace {
    param (
        [Parameter(Mandatory = $true)]
        [System.Management.Automation.Job[]]$Jobs
    )
    
    $results = @()
    foreach ($job in $Jobs) {
        $jobResult = Receive-Job -Job $job -Wait -AutoRemoveJob
        $results += $jobResult
    }
    
    return $results
}

Function Get-AdfsServerTrace {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$ActivityId,

        [switch]$IncludeDebug,

        [string]$OutHtmlFilePath,

        [string[]]$ComputerName = @("localhost")
    )

    # Start jobs to search all computers in parallel and retrieve results
    $jobs = Start-AdfsServerTrace -ActivityId $ActivityId -IncludeDebug:$IncludeDebug -ComputerName $ComputerName
    $results = Receive-AdfsServerTrace -Jobs $jobs

    if ($OutHtmlFilePath) {
        $results | ConvertTo-Html -Property ComputerName, Events | Out-File $OutHtmlFilePath -Force
        Write-Output "Report Generated at $OutHtmlFilePath"
        Start-Process $OutHtmlFilePath
    } else {
        Write-Output $results
    }
}