[CmdletBinding(SupportsShouldProcess)]
param(
    [Parameter(Mandatory)]
    [ValidatePattern('^[A-Za-z0-9][A-Za-z0-9.-]*$')]
    [string]$ServerAddress,

    [string]$InstallRoot = "$env:ProgramFiles\SGU\RustDeskServer",
    [string]$DataRoot = "$env:ProgramData\SGU\RustDesk\Server",
    [string]$FirewallRemoteAddress = '192.168.50.0/24',
    [uri]$DownloadUri = 'https://github.com/rustdesk/rustdesk-server/releases/download/1.1.16/rustdesk-server-windows-x86_64-unsigned.zip',
    [ValidatePattern('^[A-Fa-f0-9]{64}$')]
    [string]$ExpectedSha256 = 'B865A3A62FC8755B45480C508F1C4871C3338590408DDA8C58C7E9C373B7ADB0'
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'
$hbbsTaskName = 'SGU-RustDesk-hbbs'
$hbbrTaskName = 'SGU-RustDesk-hbbr'
$downloadRoot = Join-Path $env:ProgramData 'SGU\RustDesk\Downloads'
$archivePath = Join-Path $downloadRoot 'rustdesk-server-windows-x86_64-1.1.16.zip'

function Assert-Administrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    if (-not $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
        throw 'Run this script from an elevated Windows PowerShell session.'
    }
}

function Set-PrivateDirectoryAcl {
    param([Parameter(Mandatory)][string]$Path)

    New-Item -ItemType Directory -Path $Path -Force | Out-Null
    $acl = New-Object Security.AccessControl.DirectorySecurity
    $acl.SetAccessRuleProtection($true, $false)
    $inheritance = [Security.AccessControl.InheritanceFlags]'ContainerInherit,ObjectInherit'
    $allow = [Security.AccessControl.AccessControlType]::Allow
    foreach ($sid in @('S-1-5-18', 'S-1-5-32-544')) {
        $acl.AddAccessRule([Security.AccessControl.FileSystemAccessRule]::new(
            [Security.Principal.SecurityIdentifier]::new($sid),
            [Security.AccessControl.FileSystemRights]::FullControl,
            $inheritance,
            [Security.AccessControl.PropagationFlags]::None,
            $allow))
    }
    Set-Acl -LiteralPath $Path -AclObject $acl
}

function Assert-FileHash {
    param(
        [Parameter(Mandatory)][string]$Path,
        [Parameter(Mandatory)][string]$ExpectedHash
    )

    $actualHash = (Get-FileHash -LiteralPath $Path -Algorithm SHA256).Hash
    if (-not $actualHash.Equals($ExpectedHash, [StringComparison]::OrdinalIgnoreCase)) {
        throw "SHA-256 verification failed for $Path."
    }
}

function Copy-IfDifferent {
    param(
        [Parameter(Mandatory)][string]$Source,
        [Parameter(Mandatory)][string]$Destination
    )

    if (-not (Test-Path -LiteralPath $Destination -PathType Leaf) -or
        (Get-FileHash -LiteralPath $Source -Algorithm SHA256).Hash -ne
        (Get-FileHash -LiteralPath $Destination -Algorithm SHA256).Hash) {
        Copy-Item -LiteralPath $Source -Destination $Destination -Force
        return $true
    }
    return $false
}

function Set-RustDeskFirewallRule {
    param(
        [Parameter(Mandatory)][string]$Name,
        [Parameter(Mandatory)][ValidateSet('TCP', 'UDP')][string]$Protocol,
        [Parameter(Mandatory)][string]$LocalPort
    )

    $rule = Get-NetFirewallRule -DisplayName $Name -ErrorAction SilentlyContinue
    if (-not $rule) {
        $rule = New-NetFirewallRule -DisplayName $Name -Group 'SGU RustDesk' `
            -Direction Inbound -Action Allow -Protocol $Protocol -LocalPort $LocalPort `
            -RemoteAddress $FirewallRemoteAddress -Profile Domain -Enabled True
    }
    else {
        $rule | Set-NetFirewallRule -Enabled True -Profile Domain -Action Allow | Out-Null
        $rule | Get-NetFirewallPortFilter | Set-NetFirewallPortFilter `
            -Protocol $Protocol -LocalPort $LocalPort | Out-Null
        $rule | Get-NetFirewallAddressFilter | Set-NetFirewallAddressFilter `
            -RemoteAddress $FirewallRemoteAddress | Out-Null
    }
}

function Stop-RustDeskTasks {
    foreach ($taskName in @($hbbsTaskName, $hbbrTaskName)) {
        $task = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
        if ($task -and $task.State -eq 'Running') {
            Stop-ScheduledTask -TaskName $taskName
        }
    }
    Start-Sleep -Seconds 1
}

function Register-RustDeskServerTask {
    param(
        [Parameter(Mandatory)][string]$TaskName,
        [Parameter(Mandatory)][string]$Executable,
        [string]$Arguments
    )

    # New-ScheduledTaskAction rejects an empty -Argument value.  hbbr has no
    # command-line arguments, whereas hbbs needs the relay endpoint, so add
    # the parameter only when it is meaningful.
    $actionParameters = @{
        Execute = $Executable
        WorkingDirectory = $DataRoot
    }
    if (-not [string]::IsNullOrWhiteSpace($Arguments)) {
        $actionParameters.Argument = $Arguments
    }
    $action = New-ScheduledTaskAction @actionParameters
    $trigger = New-ScheduledTaskTrigger -AtStartup
    $trigger.Delay = 'PT30S'
    $settings = New-ScheduledTaskSettingsSet -StartWhenAvailable `
        -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries `
        -ExecutionTimeLimit ([TimeSpan]::Zero) `
        -RestartCount 3 -RestartInterval (New-TimeSpan -Minutes 1)
    Register-ScheduledTask -TaskName $TaskName -Action $action -Trigger $trigger `
        -Settings $settings -User 'SYSTEM' -RunLevel Highest -Force | Out-Null
    $registeredTask = Get-ScheduledTask -TaskName $TaskName -ErrorAction Stop
    if ($registeredTask.State -ne 'Running') {
        Start-ScheduledTask -TaskName $TaskName
    }
}

function Wait-ForRustDeskServer {
    for ($attempt = 1; $attempt -le 30; $attempt++) {
        $hbbsListening = [bool](Get-NetTCPConnection -LocalPort 21116 -State Listen `
            -ErrorAction SilentlyContinue)
        $hbbrListening = [bool](Get-NetTCPConnection -LocalPort 21117 -State Listen `
            -ErrorAction SilentlyContinue)
        $publicKeyReady = Test-Path -LiteralPath (Join-Path $DataRoot 'id_ed25519.pub') -PathType Leaf
        if ($hbbsListening -and $hbbrListening -and $publicKeyReady) {
            return
        }
        Start-Sleep -Seconds 2
    }
    throw 'RustDesk hbbs/hbbr did not become ready within 60 seconds.'
}

Assert-Administrator
if (-not $PSCmdlet.ShouldProcess($env:COMPUTERNAME, 'Install and configure the RustDesk OSS rendezvous and relay server')) {
    return
}

Set-PrivateDirectoryAcl -Path $DataRoot
$managementRoot = Split-Path $DataRoot -Parent
Set-PrivateDirectoryAcl -Path $managementRoot
New-Item -ItemType Directory -Path $InstallRoot,$downloadRoot -Force | Out-Null

if (-not (Test-Path -LiteralPath $archivePath -PathType Leaf) -or
    (Get-FileHash -LiteralPath $archivePath -Algorithm SHA256).Hash -ne $ExpectedSha256) {
    Invoke-WebRequest -Uri $DownloadUri -OutFile $archivePath -UseBasicParsing
}
Assert-FileHash -Path $archivePath -ExpectedHash $ExpectedSha256

$stagingRoot = Join-Path $env:TEMP ('sgu-rustdesk-server-' + [Guid]::NewGuid().ToString('N'))
try {
    Expand-Archive -LiteralPath $archivePath -DestinationPath $stagingRoot -Force
    $payloadRoot = Join-Path $stagingRoot 'x86_64'
    $sourceHbbs = Join-Path $payloadRoot 'hbbs.exe'
    $sourceHbbr = Join-Path $payloadRoot 'hbbr.exe'
    foreach ($required in @($sourceHbbs, $sourceHbbr)) {
        if (-not (Test-Path -LiteralPath $required -PathType Leaf)) {
            throw "The verified RustDesk archive is missing $required."
        }
    }

    $targetHbbs = Join-Path $InstallRoot 'hbbs.exe'
    $targetHbbr = Join-Path $InstallRoot 'hbbr.exe'
    $requiresBinaryUpdate =
        -not (Test-Path -LiteralPath $targetHbbs) -or
        -not (Test-Path -LiteralPath $targetHbbr) -or
        (Get-FileHash -LiteralPath $sourceHbbs -Algorithm SHA256).Hash -ne
            (Get-FileHash -LiteralPath $targetHbbs -Algorithm SHA256).Hash -or
        (Get-FileHash -LiteralPath $sourceHbbr -Algorithm SHA256).Hash -ne
            (Get-FileHash -LiteralPath $targetHbbr -Algorithm SHA256).Hash
    if ($requiresBinaryUpdate) {
        Stop-RustDeskTasks
        Copy-IfDifferent -Source $sourceHbbs -Destination $targetHbbs | Out-Null
        Copy-IfDifferent -Source $sourceHbbr -Destination $targetHbbr | Out-Null
    }
}
finally {
    if (Test-Path -LiteralPath $stagingRoot) {
        Remove-Item -LiteralPath $stagingRoot -Recurse -Force
    }
}

Set-RustDeskFirewallRule -Name 'SGU RustDesk hbbs (TCP)' -Protocol TCP -LocalPort '21115-21116'
Set-RustDeskFirewallRule -Name 'SGU RustDesk hbbr (TCP)' -Protocol TCP -LocalPort '21117'
Set-RustDeskFirewallRule -Name 'SGU RustDesk hbbs (UDP)' -Protocol UDP -LocalPort '21116'

Register-RustDeskServerTask -TaskName $hbbrTaskName -Executable (Join-Path $InstallRoot 'hbbr.exe')
Register-RustDeskServerTask -TaskName $hbbsTaskName -Executable (Join-Path $InstallRoot 'hbbs.exe') `
    -Arguments "-r $ServerAddress`:21117"
Wait-ForRustDeskServer

$publicKey = (Get-Content -LiteralPath (Join-Path $DataRoot 'id_ed25519.pub') -Raw).Trim()
if ([string]::IsNullOrWhiteSpace($publicKey)) {
    throw 'RustDesk generated an empty public key.'
}
$statusPath = Join-Path (Split-Path $DataRoot -Parent) 'server.json'
$status = [ordered]@{
    ServerAddress = $ServerAddress
    PublicKey = $publicKey
    PublicKeySha256 = ([Security.Cryptography.SHA256]::Create().ComputeHash(
        [Text.Encoding]::UTF8.GetBytes($publicKey)) | ForEach-Object ToString x2) -join ''
    HbbsTaskName = $hbbsTaskName
    HbbrTaskName = $hbbrTaskName
    HbbsTcpPort = 21116
    HbbrTcpPort = 21117
    InstalledAt = (Get-Date).ToString('o')
}
[IO.File]::WriteAllText($statusPath, ($status | ConvertTo-Json), [Text.UTF8Encoding]::new($false))

[pscustomobject]@{
    ServerAddress = $ServerAddress
    PublicKey = $publicKey
    PublicKeySha256 = $status.PublicKeySha256
    HbbsTask = (Get-ScheduledTask -TaskName $hbbsTaskName).State.ToString()
    HbbrTask = (Get-ScheduledTask -TaskName $hbbrTaskName).State.ToString()
    HbbsListening = [bool](Get-NetTCPConnection -LocalPort 21116 -State Listen -ErrorAction SilentlyContinue)
    HbbrListening = [bool](Get-NetTCPConnection -LocalPort 21117 -State Listen -ErrorAction SilentlyContinue)
    StatusPath = $statusPath
}