[CmdletBinding()]
param(
    [string]$ServerHost = "192.168.31.51",
    [string]$ServerUser = "jsxq",
    [int]$SshPort = 22
)

$ErrorActionPreference = "Stop"
$knownHostsFile = Join-Path ([IO.Path]::GetTempPath()) "codex-company-linux-known-hosts"

function Require-Command([string]$Name) {
    if (-not (Get-Command $Name -ErrorAction SilentlyContinue)) {
        throw "Required command not found: $Name"
    }
}

Require-Command "ssh.exe"
Require-Command "ssh-keygen.exe"

if ([string]::IsNullOrWhiteSpace($env:COMPANY_LINUX_PASSWORD)) {
    throw "Set COMPANY_LINUX_PASSWORD in the current process before running first-time key setup."
}

$sshDirectory = Join-Path $env:USERPROFILE ".ssh"
$privateKey = Join-Path $sshDirectory "id_ed25519"
$publicKey = "$privateKey.pub"

if (-not (Test-Path -LiteralPath $publicKey)) {
    New-Item -ItemType Directory -Path $sshDirectory -Force | Out-Null
    & ssh-keygen.exe -q -t ed25519 -f $privateKey -N "" -C "$env:USERNAME@codex-company-linux"
    if ($LASTEXITCODE -ne 0) {
        throw "ssh-keygen failed with exit code $LASTEXITCODE"
    }
}

$keyText = (Get-Content -LiteralPath $publicKey -Raw).Trim()
if ([string]::IsNullOrWhiteSpace($keyText)) {
    throw "Public key is empty: $publicKey"
}

$keyBase64 = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($keyText))
$askPass = Join-Path $PSScriptRoot "ssh-askpass.cmd"
$oldAskPass = $env:SSH_ASKPASS
$oldAskPassRequire = $env:SSH_ASKPASS_REQUIRE
$oldDisplay = $env:DISPLAY

try {
    $env:SSH_ASKPASS = $askPass
    $env:SSH_ASKPASS_REQUIRE = "force"
    $env:DISPLAY = "codex"

    $remote = "set -eu; umask 077; mkdir -p ~/.ssh; touch ~/.ssh/authorized_keys; key=`$(printf '%s' '$keyBase64' | base64 -d); grep -qxF -- `"`$key`" ~/.ssh/authorized_keys || printf '%s\n' `"`$key`" >> ~/.ssh/authorized_keys; chmod 700 ~/.ssh; chmod 600 ~/.ssh/authorized_keys"
    & ssh.exe -T -p $SshPort -o StrictHostKeyChecking=accept-new -o "UserKnownHostsFile=$knownHostsFile" -o ConnectTimeout=10 -o NumberOfPasswordPrompts=1 "$ServerUser@$ServerHost" $remote
    if ($LASTEXITCODE -ne 0) {
        throw "Unable to install the SSH public key."
    }
} finally {
    $env:SSH_ASKPASS = $oldAskPass
    $env:SSH_ASKPASS_REQUIRE = $oldAskPassRequire
    $env:DISPLAY = $oldDisplay
}

& ssh.exe -T -p $SshPort -o BatchMode=yes -o StrictHostKeyChecking=accept-new -o "UserKnownHostsFile=$knownHostsFile" -o ConnectTimeout=10 "$ServerUser@$ServerHost" "echo SSH_KEY_READY"
if ($LASTEXITCODE -ne 0) {
    throw "The key was installed, but key authentication verification failed."
}

Write-Output "SSH key authentication is ready for $ServerUser@$ServerHost."