param(
    [string]$Region = "nyc2",
    [string]$Size = "gpu-mi325x8-2048gb",
    [string]$Image = "gpu-amd-base",
    [string]$Name = "riemann-node",
    [Parameter(Mandatory = $true)][string]$SshKeyName,
    [Parameter(Mandatory = $true)][string]$StagingRepo,
    [switch]$StopAfterStage2
)
$ErrorActionPreference = "Stop"
$proj = Split-Path -Parent $PSScriptRoot
$sessionFiles = "$env:RIEMANN_ROOT\cloud_state"
New-Item -ItemType Directory -Force -Path $sessionFiles | Out-Null

$existing = doctl compute droplet list --format ID,Name,PublicIPv4,Status --no-header | Where-Object { $_ -match "\b$Name\b" }
if ($existing) {
    Write-Host "Droplet already exists: $existing"
    $ip = ($existing -split '\s+')[2]
} else {
    $keyId = (doctl compute ssh-key list --format ID,Name --no-header | Where-Object { $_ -match $SshKeyName } | ForEach-Object { ($_ -split '\s+')[0] })
    if (-not $keyId) { throw "SSH key '$SshKeyName' not found on the account" }
    Write-Host "Creating $Size in $Region with image $Image ..."
    doctl compute droplet create $Name --region $Region --size $Size --image $Image --ssh-keys $keyId --wait --format ID,Name,PublicIPv4,Status | Out-Host
    $ip = (doctl compute droplet list --format Name,PublicIPv4 --no-header | Where-Object { $_ -match "\b$Name\b" } | ForEach-Object { ($_ -split '\s+')[1] })
}
Write-Host "Droplet IP: $ip"
Set-Content -Path "$sessionFiles\node_ip.txt" -Value $ip

Write-Host "Waiting for SSH ..."
$deadline = (Get-Date).AddMinutes(10)
do {
    $ok = $false
    try { ssh -o StrictHostKeyChecking=accept-new -o ConnectTimeout=8 -o BatchMode=yes root@$ip "echo ready" 2>$null | Select-String ready | ForEach-Object { $ok = $true } } catch {}
    if (-not $ok) { Start-Sleep 10 }
} while (-not $ok -and (Get-Date) -lt $deadline)
if (-not $ok) { throw "SSH not reachable after 10 minutes" }

Write-Host "Packaging project ..."
$bundle = "$env:TEMP\riemann_bundle.tar.gz"
Push-Location $proj
tar -czf $bundle riemann scripts configs cloud requirements-cuda.txt LICENSE
Pop-Location

Write-Host "Shipping code and HF token ..."
ssh root@$ip "mkdir -p /workspace/riemann ~/.cache/huggingface"
scp -q $bundle root@${ip}:/workspace/riemann_bundle.tar.gz
scp -q "$env:USERPROFILE\.cache\huggingface\token" root@${ip}:~/.cache/huggingface/token
ssh root@$ip "cd /workspace/riemann && tar -xzf /workspace/riemann_bundle.tar.gz && sed -i 's/\r$//' cloud/*.sh cloud/*.py && chmod +x cloud/bootstrap.sh && chmod 600 ~/.cache/huggingface/token && ls"

$stopFlag = if ($StopAfterStage2) { "STOP_AFTER_STAGE2=1" } else { "STOP_AFTER_STAGE2=0" }
Write-Host "Starting bootstrap under nohup ($stopFlag) ..."
ssh root@$ip "cd /workspace/riemann && mkdir -p /workspace/RiemannModel/logs && STAGING_REPO=$StagingRepo $stopFlag nohup bash cloud/bootstrap.sh > /workspace/RiemannModel/logs/bootstrap.log 2>&1 &"
Start-Sleep 5
ssh root@$ip "tail -5 /workspace/RiemannModel/logs/bootstrap.log"
Write-Host ""
Write-Host "Launched. Follow with:  ssh root@$ip tail -f /workspace/RiemannModel/logs/bootstrap.log"
Write-Host "Destroy with:           doctl compute droplet delete $Name --force"