#Requires -Version 5.1
#Requires -RunAsAdministrator
[CmdletBinding()]
param(
    [string]$ConnectionName = 'SGU Azure Device',
    [ValidateRange(30,600)][int]$WaitSeconds = 180
)

$ErrorActionPreference = 'Stop'
$computer = Get-CimInstance Win32_ComputerSystem
if (-not $computer.PartOfDomain) { throw 'The device must already be joined to its domain.' }
$logPath = Join-Path $env:ProgramData 'SGU\Enrollment\azure-domain-connectivity.json'
$deadline = (Get-Date).AddSeconds($WaitSeconds)
$restarted = $false
$controller = $null
try {
    do {
        $vpn = Get-VpnConnection -Name $ConnectionName -AllUserConnection -ErrorAction SilentlyContinue
        $reachable = $false
        if ($vpn -and $vpn.ConnectionStatus -eq 'Connected') {
            $record = Resolve-DnsName "_ldap._tcp.dc._msdcs.$($computer.Domain)" -Type SRV -ErrorAction SilentlyContinue |
                Where-Object Type -eq 'SRV' | Select-Object -First 1
            if ($record) {
                $controller = $record.NameTarget.TrimEnd('.')
                $socket = [Net.Sockets.TcpClient]::new()
                try {
                    $connect = $socket.BeginConnect($controller, 389, $null, $null)
                    if ($connect.AsyncWaitHandle.WaitOne(2000)) {
                        $socket.EndConnect($connect)
                        $reachable = $socket.Connected
                    }
                } catch { $reachable = $false }
                finally { $socket.Dispose() }
            }
        }
        if ($reachable) { break }
        Start-Sleep -Seconds 5
    } while ((Get-Date) -lt $deadline)
    if (-not $reachable) { throw "The VPN and a domain controller were not reachable within $WaitSeconds seconds." }

    # An early Netlogon attempt can remain failed after the device VPN connects.
    # Refresh only that service, after confirming the domain is reachable.
    if (-not (Test-ComputerSecureChannel -Server $controller)) {
        Restart-Service -Name Netlogon
        $restarted = $true
    }
    $secure = $false
    for ($attempt = 0; $attempt -lt 6; $attempt++) {
        $secure = Test-ComputerSecureChannel -Server $controller
        if ($secure) { break }
        Start-Sleep -Seconds 5
    }
    if (-not $secure) { throw 'The domain is reachable but the secure channel is still invalid. Administrative repair is required.' }
    $guard = Get-ScheduledTask -TaskName 'SGU-CredentialProvider-EnrollmentGuard' -ErrorAction SilentlyContinue
    $guardResult = $null
    if ($guard) {
        # Domain principal lookup can recover after the secure channel itself.
        # Await the guard and retry a transient failure instead of reporting
        # success while its asynchronous repair is still running or failed.
        $guardDeadline = (Get-Date).AddMinutes(3)
        do {
            $guard = Get-ScheduledTask -TaskName $guard.TaskName
            if ($guard.State -notin @('Running','Queued')) {
                $previousRun = (Get-ScheduledTaskInfo -TaskName $guard.TaskName).LastRunTime
                Start-ScheduledTask -InputObject $guard
                do {
                    Start-Sleep -Seconds 2
                    $guard = Get-ScheduledTask -TaskName $guard.TaskName
                    $info = Get-ScheduledTaskInfo -TaskName $guard.TaskName
                } while (($info.LastRunTime -le $previousRun -or $guard.State -in @('Running','Queued')) -and (Get-Date) -lt $guardDeadline)
                if ($info.LastRunTime -gt $previousRun -and $guard.State -notin @('Running','Queued')) {
                    $guardResult = $info.LastTaskResult
                    if ($guardResult -eq 0) { break }
                }
            }
            Start-Sleep -Seconds 10
        } while ((Get-Date) -lt $guardDeadline)
        if ($guardResult -ne 0) { throw "The secure channel recovered, but the enrollment guard did not succeed (result $guardResult)." }
    }
    [pscustomobject]@{
        CheckedAt = (Get-Date).ToString('o')
        ComputerName = $computer.Name
        Domain = $computer.Domain
        DomainController = $controller
        ConnectionName = $ConnectionName
        NetlogonRestarted = $restarted
        SecureChannel = $secure
        EnrollmentGuardResult = $guardResult
    } | ConvertTo-Json | Set-Content -LiteralPath $logPath
} catch {
    [pscustomobject]@{
        CheckedAt = (Get-Date).ToString('o')
        ConnectionName = $ConnectionName
        NetlogonRestarted = $restarted
        SecureChannel = $false
        Error = $_.Exception.Message
    } | ConvertTo-Json | Set-Content -LiteralPath $logPath
    throw
}