Param(
    [Parameter(Mandatory = $true, Position = 0)]
    [string]$Tenant,

    [Parameter()]
    [int]$Port = 22,

    [Parameter()]
    [string]$Name,

    [Parameter()]
    [string]$Tag,

    [Parameter()]
    [switch]$IgnoreSslCert
)

$ErrorActionPreference = "Stop"

function Write-Info($message) {
    Write-Host "[INFO] $message" -ForegroundColor Green
}

function Write-Warn($message) {
    Write-Host "[WARN] $message" -ForegroundColor Yellow
}

function Write-ErrorAndExit($message) {
    Write-Host "[ERROR] $message" -ForegroundColor Red
    exit 1
}

function Enable-InsecureTls {
    if ($IgnoreSslCert) {
        [System.Net.ServicePointManager]::ServerCertificateValidationCallback = { $true }
    }
}

function Get-PrimaryIPv4 {
    $ip = Get-NetIPAddress -AddressFamily IPv4 -ErrorAction SilentlyContinue |
        Where-Object { $_.IPAddress -ne "127.0.0.1" -and $_.IPAddress -notlike "169.254*" } |
        Select-Object -First 1

    if ($null -ne $ip) {
        return $ip.IPAddress
    }

    return "unknown"
}

function Get-HostKeyFingerprints {
    $scan = & ssh-keyscan -p $Port localhost 2>$null
    if ($LASTEXITCODE -ne 0 -or [string]::IsNullOrWhiteSpace($scan)) {
        Write-ErrorAndExit "Could not probe local SSH keys on port $Port"
    }

    $fps = @()
    foreach ($line in ($scan -split "\r?\n")) {
        if (-not [string]::IsNullOrWhiteSpace($line)) {
            $fpLine = $line | & ssh-keygen -lf - 2>$null
            if ($LASTEXITCODE -eq 0 -and $fpLine) {
                $parts = $fpLine -split "\s+"
                if ($parts.Length -ge 2) {
                    $fps += $parts[1].Trim()
                }
            }
        }
    }

    $uniqueFps = $fps | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | Sort-Object -Unique
    if (-not $uniqueFps -or $uniqueFps.Count -eq 0) {
        Write-ErrorAndExit "No host key fingerprints found from ssh-keyscan output."
    }

    $uniqueFps
}

function Build-FormBody([hashtable]$fields, [string[]]$fingerprints) {
    Add-Type -AssemblyName System.Web
    $parts = @()
    foreach ($key in $fields.Keys) {
        $val = $fields[$key]
        if ($null -ne $val -and $val.ToString().Length -gt 0) {
            $parts += "{0}={1}" -f [System.Web.HttpUtility]::UrlEncode($key), [System.Web.HttpUtility]::UrlEncode($val)
        }
    }
    foreach ($fp in $fingerprints) {
        if ($fp) {
            $parts += "{0}={1}" -f [System.Web.HttpUtility]::UrlEncode("fingerprint"), [System.Web.HttpUtility]::UrlEncode($fp)
        }
    }
    $parts -join "&"
}

function Restart-OpenSsh {
    Write-Info "Restarting OpenSSH service..."
    Restart-Service -Name sshd -Force -ErrorAction Stop
}

Enable-InsecureTls

Write-Info "Collecting server metadata..."
$hostname = [System.Net.Dns]::GetHostName()
$serverName = if ($Name) { $Name } else { $hostname }
$ipAddr = Get-PrimaryIPv4
$osFamily = "windows"
$osVersion = [System.Environment]::OSVersion.VersionString

Write-Info "Requesting device authorization..."
$deviceBody = Build-FormBody @{
    client_id = "sshteam"
    scope = "sshteam-ca-authority"
} @()

$deviceResponse = Invoke-RestMethod -Method Post -Uri "https://$Tenant/oauth2/device" -Body $deviceBody -ContentType "application/x-www-form-urlencoded"

$verifyUri = $deviceResponse.verification_uri_complete
$userCode = $deviceResponse.user_code
$deviceCode = $deviceResponse.device_code
$interval = if ($deviceResponse.interval) { [int]$deviceResponse.interval } else { 5 }

Write-Host ""
Write-Host "ACTION REQUIRED: Please visit: $verifyUri"
Write-Host "Confirm code: $userCode"
Read-Host "Press [Enter] after approving" | Out-Null

Write-Info "Waiting for authorization..."
$accessToken = $null
while ($true) {
    $tokenBody = Build-FormBody @{
        grant_type = "urn:ietf:params:oauth:grant-type:device_code"
        device_code = $deviceCode
        client_id = "sshteam"
    } @()

    $tokenResponse = Invoke-RestMethod -Method Post -Uri "https://$Tenant/oauth2/token" -Body $tokenBody -ContentType "application/x-www-form-urlencoded"

    if ($tokenResponse.access_token) {
        $accessToken = $tokenResponse.access_token
        break
    }

    if ($tokenResponse.error -eq "authorization_pending" -or $tokenResponse.error -eq "slow_down") {
        Start-Sleep -Seconds $interval
        continue
    }

    Write-ErrorAndExit "Failed to retrieve access token: $($tokenResponse | ConvertTo-Json -Compress)"
}

Write-Info "Probing local SSH host keys..."
$fingerprints = Get-HostKeyFingerprints
foreach ($fp in $fingerprints) {
    if ($fp) {
        Write-Info "Found host key fingerprint: $fp"
    }
}

$joinUrl = "https://$Tenant/ssh/join"
Write-Info "Registering server '$serverName' ($ipAddr) with $Tenant..."
$joinFields = @{
    port = $Port
    name = $serverName
    hostname = $hostname
    ipAddress = $ipAddr
    osFamily = $osFamily
    osVersion = $osVersion
}
if ($Tag) {
    $joinFields.tag = $Tag
}

$joinBody = Build-FormBody $joinFields $fingerprints
$headers = @{ Authorization = "Bearer $accessToken" }

try {
    $joinResponse = Invoke-WebRequest -Method Post -Uri $joinUrl -Headers $headers -Body $joinBody -ContentType "application/x-www-form-urlencoded"
} catch {
    $statusCode = $_.Exception.Response.StatusCode.value__
    Write-ErrorAndExit "Join request failed with HTTP $statusCode for $joinUrl"
}

$caData = $joinResponse.Content
if ([string]::IsNullOrWhiteSpace($caData)) {
    Write-ErrorAndExit "No CA data returned from $joinUrl"
}

$caPath = "C:\ProgramData\ssh\sshteam_ca.pub"
$sshdConfig = "C:\ProgramData\ssh\sshd_config"

Write-Info "Writing CA keys to $caPath..."
$caData | Out-File -FilePath $caPath -Encoding ascii -Force

if (-not (Select-String -Path $sshdConfig -Pattern "TrustedUserCAKeys $caPath" -Quiet)) {
    Write-Info "Updating $sshdConfig with TrustedUserCAKeys..."
    Add-Content -Path $sshdConfig -Value "`r`n# SSHTEAM CA TRUST`r`nTrustedUserCAKeys $caPath"
} else {
    Write-Info "SaaS CA trust already configured in $sshdConfig."
}

Restart-OpenSsh
$accessToken = $null
Write-Info "Server joined and OpenSSH configured."
