param(
  [Parameter(Mandatory = $false)]
  [string]$PairingCode,

  [string]$RelayUrl = "http://127.0.0.1:4173",
  [int]$IntervalSeconds = 60,
  [string]$InstallMode = "manual",
  [switch]$Once
)

$ErrorActionPreference = "Stop"

$UserStateDir = Join-Path $env:LOCALAPPDATA "AllReach"
$MachineStateDir = Join-Path $env:ProgramData "AllReach"
$UserStatePath = Join-Path $UserStateDir "agent-state.json"
$MachineStatePath = Join-Path $MachineStateDir "agent-state.json"
$StatePath = if (Test-Path $MachineStatePath) { $MachineStatePath } else { $UserStatePath }
$StateDir = Split-Path -Parent $StatePath

function Get-IPv4Addresses {
  Get-NetIPAddress -AddressFamily IPv4 |
    Where-Object {
      $_.IPAddress -notlike "169.254.*" -and
      $_.IPAddress -ne "127.0.0.1"
    } |
    Select-Object -ExpandProperty IPAddress
}

function Test-RdpEnabled {
  $key = "HKLM:\System\CurrentControlSet\Control\Terminal Server"
  $value = Get-ItemProperty -Path $key -Name "fDenyTSConnections"
  return $value.fDenyTSConnections -eq 0
}

function Test-PortListening {
  param([int]$ListenPort)

  $connection = Get-NetTCPConnection -LocalPort $ListenPort -State Listen -ErrorAction SilentlyContinue
  return $null -ne $connection
}

function Protect-AgentSecret {
  param([string]$AgentSecret)

  if (-not $AgentSecret) {
    return ""
  }

  $bytes = [System.Text.Encoding]::UTF8.GetBytes($AgentSecret)
  $protected = [System.Security.Cryptography.ProtectedData]::Protect(
    $bytes,
    $null,
    [System.Security.Cryptography.DataProtectionScope]::CurrentUser
  )
  return [Convert]::ToBase64String($protected)
}

function Unprotect-AgentSecret {
  param([string]$ProtectedSecret)

  if (-not $ProtectedSecret) {
    return ""
  }

  try {
    $protected = [Convert]::FromBase64String($ProtectedSecret)
    $bytes = [System.Security.Cryptography.ProtectedData]::Unprotect(
      $protected,
      $null,
      [System.Security.Cryptography.DataProtectionScope]::CurrentUser
    )
    return [System.Text.Encoding]::UTF8.GetString($bytes)
  } catch {
    return ""
  }
}

function Get-AgentSecretFromState {
  param([object]$State)

  if (-not $State) {
    return ""
  }

  if ($State.agentSecret) {
    return $State.agentSecret
  }

  return Unprotect-AgentSecret -ProtectedSecret $State.agentSecretProtected
}

function Get-AgentStatus {
  $rdpPort = 3389
  $edition = (Get-ComputerInfo).WindowsProductName
  $rdpEnabled = Test-RdpEnabled
  $rdpListening = Test-PortListening -ListenPort $rdpPort

  [ordered]@{
    deviceName = $env:COMPUTERNAME
    platform = "windows"
    edition = $edition
    installMode = $InstallMode
    online = $true
    addresses = @(Get-IPv4Addresses)
    protocols = @(
      [ordered]@{
        name = "rdp"
        enabled = $rdpEnabled
        port = $rdpPort
        listening = $rdpListening
      }
    )
    routes = [ordered]@{
      lan = $true
      mesh = $false
      relay = $true
    }
    generatedAt = (Get-Date).ToUniversalTime().ToString("o")
  }
}

function Read-AgentState {
  if (-not (Test-Path $StatePath)) {
    return $null
  }

  Get-Content $StatePath -Raw | ConvertFrom-Json
}

function Save-AgentState {
  param([object]$State)

  New-Item -ItemType Directory -Path $StateDir -Force | Out-Null
  $State | ConvertTo-Json -Depth 8 | Set-Content -Path $StatePath -Encoding UTF8
}

function Update-AgentStateFromResponse {
  param(
    [object]$Response,
    [object]$ExistingState
  )

  Save-AgentState -State ([ordered]@{
    agentId = if ($Response.id) { $Response.id } else { $ExistingState.agentId }
    agentSecretProtected = Protect-AgentSecret -AgentSecret $(if ($Response.agentSecret) { $Response.agentSecret } else { Get-AgentSecretFromState -State $ExistingState })
    deviceId = if ($Response.deviceId) { $Response.deviceId } else { $ExistingState.deviceId }
    deviceName = if ($Response.deviceName) { $Response.deviceName } else { $ExistingState.deviceName }
    relayUrl = if ($ExistingState.relayUrl) { $ExistingState.relayUrl } else { $RelayUrl }
    pairedAt = $ExistingState.pairedAt
    lastHeartbeatAt = if ($Response.lastSeen) { $Response.lastSeen } else { (Get-Date).ToUniversalTime().ToString("o") }
  })
}

function Invoke-AgentRegister {
  param([string]$Code)

  $status = Get-AgentStatus
  $body = [ordered]@{
    pairingCode = $Code
    deviceName = $status.deviceName
    platform = $status.platform
    edition = $status.edition
    addresses = $status.addresses
    protocols = $status.protocols
    routes = $status.routes
  }

  $response = Invoke-RestMethod `
    -Method Post `
    -Uri "$RelayUrl/api/agents/register" `
    -ContentType "application/json" `
    -Body ($body | ConvertTo-Json -Depth 8)

  Save-AgentState -State ([ordered]@{
    agentId = $response.id
    agentSecretProtected = Protect-AgentSecret -AgentSecret $response.agentSecret
    deviceId = $response.deviceId
    deviceName = $response.deviceName
    relayUrl = $RelayUrl
    pairedAt = (Get-Date).ToUniversalTime().ToString("o")
  })

  return $response
}

function Invoke-AgentHeartbeat {
  param(
    [string]$AgentId,
    [string]$AgentSecret
  )

  $status = Get-AgentStatus
  $headers = @{}
  if ($AgentSecret) {
    $headers["X-AllReach-Agent-Token"] = $AgentSecret
  }

  Invoke-RestMethod `
    -Method Post `
    -Uri "$RelayUrl/api/agents/$AgentId/heartbeat" `
    -Headers $headers `
    -ContentType "application/json" `
    -Body ($status | ConvertTo-Json -Depth 8)
}

function Invoke-AgentSessions {
  param(
    [string]$AgentId,
    [string]$AgentSecret
  )

  $headers = @{}
  if ($AgentSecret) {
    $headers["X-AllReach-Agent-Token"] = $AgentSecret
  }

  Invoke-RestMethod `
    -Method Get `
    -Headers $headers `
    -Uri "$RelayUrl/api/agents/$AgentId/sessions"
}

function Invoke-ClaimSession {
  param(
    [string]$AgentId,
    [object]$Session
  )

  $body = [ordered]@{
    agentId = $AgentId
    agentToken = $Session.agentToken
  }

  Invoke-RestMethod `
    -Method Post `
    -Uri "$RelayUrl/api/sessions/$($Session.sessionId)/claim" `
    -ContentType "application/json" `
    -Body ($body | ConvertTo-Json -Depth 4)
}

function Sync-AgentSessions {
  param(
    [string]$AgentId,
    [string]$AgentSecret
  )

  $payload = Invoke-AgentSessions -AgentId $AgentId -AgentSecret $AgentSecret
  $sessions = @($payload.sessions)

  foreach ($session in $sessions) {
    if ($session.status -ne "pending") {
      continue
    }

    $claimed = Invoke-ClaimSession -AgentId $AgentId -Session $session
    Write-Host "Claimed $($claimed.protocol.ToUpper()) session $($claimed.sessionId) on local port $($claimed.connectPort)."
  }
}

if ($PairingCode) {
  $registered = Invoke-AgentRegister -Code $PairingCode
  Write-Host "AllReach agent paired as $($registered.id) for $($registered.deviceName)."
}

$state = Read-AgentState
if (-not $state -or -not $state.agentId) {
  throw "No paired AllReach agent was found. Run with -PairingCode NN-NN-NN first."
}

if (-not $PairingCode -and $state.relayUrl -and $RelayUrl -eq "http://127.0.0.1:4173") {
  $RelayUrl = $state.relayUrl
}

do {
  $agentSecret = Get-AgentSecretFromState -State $state
  $heartbeat = Invoke-AgentHeartbeat -AgentId $state.agentId -AgentSecret $agentSecret
  Update-AgentStateFromResponse -Response $heartbeat -ExistingState $state
  $state = Read-AgentState
  Write-Host "Heartbeat sent for $($heartbeat.deviceName) at $($heartbeat.lastSeen)."
  Sync-AgentSessions -AgentId $state.agentId -AgentSecret (Get-AgentSecretFromState -State $state)

  if ($Once) {
    break
  }

  Start-Sleep -Seconds $IntervalSeconds
} while ($true)