# Setup SSH Key-based Authentication
# This script configures passwordless SSH login to Ubuntu Server

$ErrorActionPreference = "Stop"

Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  SSH Key Setup for Ubuntu Server" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan
Write-Host ""

$SSH_DIR = "$env:USERPROFILE\.ssh"
$SSH_KEY = "$SSH_DIR\id_rsa"
$SSH_PUB = "$SSH_KEY.pub"
$REMOTE_HOST = "ubuntu@10.1.0.254"

# Step 1: Check if SSH directory exists
Write-Host "[1/4] Checking SSH directory..." -ForegroundColor Yellow

if (-not (Test-Path $SSH_DIR)) {
    Write-Host "  Creating .ssh directory..." -ForegroundColor Gray
    New-Item -ItemType Directory -Path $SSH_DIR -Force | Out-Null
    Write-Host "  [OK] Directory created" -ForegroundColor Green
} else {
    Write-Host "  [OK] Directory exists" -ForegroundColor Green
}

Write-Host ""

# Step 2: Generate SSH key if not exists
Write-Host "[2/4] Checking for SSH key..." -ForegroundColor Yellow

if (Test-Path $SSH_PUB) {
    Write-Host "  [OK] SSH key already exists" -ForegroundColor Green
    Write-Host "  Key location: $SSH_KEY" -ForegroundColor Gray
    Write-Host ""
    Write-Host "  Public key content:" -ForegroundColor Gray
    Get-Content $SSH_PUB | Write-Host -ForegroundColor White
} else {
    Write-Host "  Generating new SSH key pair..." -ForegroundColor Gray
    Write-Host "  (Press Enter when prompted for passphrase to skip)" -ForegroundColor Yellow

    ssh-keygen -t rsa -b 4096 -f $SSH_KEY -C "windows-to-ubuntu"

    if ($LASTEXITCODE -eq 0) {
        Write-Host "  [OK] SSH key generated successfully" -ForegroundColor Green
    } else {
        Write-Host "  [FAIL] Failed to generate SSH key" -ForegroundColor Red
        exit 1
    }
}

Write-Host ""

# Step 3: Copy public key to remote server
Write-Host "[3/4] Copying public key to Ubuntu Server..." -ForegroundColor Yellow
Write-Host "  You will need to enter the Ubuntu password ONE TIME" -ForegroundColor Yellow
Write-Host ""

try {
    # Read public key content
    $pubKeyContent = Get-Content $SSH_PUB -Raw

    # Create command to add key to authorized_keys
    $remoteCommand = @"
mkdir -p ~/.ssh && \
chmod 700 ~/.ssh && \
echo '$pubKeyContent' >> ~/.ssh/authorized_keys && \
chmod 600 ~/.ssh/authorized_keys && \
echo 'SSH key added successfully'
"@

    # Execute remote command
    Write-Host "  Connecting to $REMOTE_HOST..." -ForegroundColor Gray
    ssh $REMOTE_HOST $remoteCommand

    if ($LASTEXITCODE -eq 0) {
        Write-Host "  [OK] Public key copied successfully" -ForegroundColor Green
    } else {
        Write-Host "  [FAIL] Failed to copy public key" -ForegroundColor Red
        exit 1
    }
} catch {
    Write-Host "  [ERROR] $($_.Exception.Message)" -ForegroundColor Red
    exit 1
}

Write-Host ""

# Step 4: Test passwordless login
Write-Host "[4/4] Testing passwordless SSH login..." -ForegroundColor Yellow

try {
    $testResult = ssh -o BatchMode=yes -o ConnectTimeout=5 $REMOTE_HOST "echo 'Success'"

    if ($LASTEXITCODE -eq 0 -and $testResult -eq "Success") {
        Write-Host "  [OK] Passwordless login works!" -ForegroundColor Green
    } else {
        Write-Host "  [FAIL] Passwordless login test failed" -ForegroundColor Red
        Write-Host "  You may need to restart SSH service on Ubuntu Server" -ForegroundColor Yellow
        Write-Host "  Command: sudo systemctl restart ssh" -ForegroundColor Gray
        exit 1
    }
} catch {
    Write-Host "  [ERROR] Test failed: $($_.Exception.Message)" -ForegroundColor Red
    exit 1
}

Write-Host ""
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  Setup Completed!" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan
Write-Host ""
Write-Host "You can now run commands without password:" -ForegroundColor Green
Write-Host "  ssh $REMOTE_HOST 'docker ps'" -ForegroundColor White
Write-Host "  ssh $REMOTE_HOST 'docker exec postgres psql -U postgres -l'" -ForegroundColor White
Write-Host ""
Write-Host "Next step:" -ForegroundColor Yellow
Write-Host "  Run check-postgres.ps1 to verify PostgreSQL connection" -ForegroundColor White
Write-Host ""