update-pm-files.ps1

# Simple script to update files on all instances using EC2 Instance Connect
# 
# Usage Examples:
#   .\update-exposure-spread.ps1                                    # Updates exposure_spread.py (default)
#   .\update-exposure-spread.ps1 -FileName "predict_dupes.py"       # Updates predict_dupes.py
#   .\update-exposure-spread.ps1 -FileName "app.py" -TargetPath "/home/ec2-user/AWS_Portfolio_Manager/"  # Updates app.py in root directory
#
param(
    [Parameter(Mandatory=$false)]
    [string]$FileName = "exposure_spread.py",
    
    [Parameter(Mandatory=$false)]
    [string]$TargetPath = "/home/ec2-user/AWS_Portfolio_Manager/global_func/"
)

# Handle special case: app.py maps to application.py on the remote instances
$remoteFileName = $FileName
if ($FileName -eq "app.py") {
    $remoteFileName = "application.py"
    Write-Host "Note: app.py will be uploaded as application.py on remote instances" -ForegroundColor Yellow
}

# Determine local file path - check global_func first, then root directory
$localFilePath = ""
if (Test-Path ".\global_func\$FileName") {
    $localFilePath = ".\global_func\$FileName"
} elseif (Test-Path ".\$FileName") {
    $localFilePath = ".\$FileName"
    # If uploading from root and no target path specified, default to root on remote
    if ($TargetPath -eq "/home/ec2-user/AWS_Portfolio_Manager/global_func/") {
        $TargetPath = "/home/ec2-user/AWS_Portfolio_Manager/"
    }
} else {
    Write-Host "Error: File '$FileName' not found!" -ForegroundColor Red
    Write-Host "Checked locations:" -ForegroundColor Yellow
    Write-Host "  .\global_func\$FileName" -ForegroundColor Cyan
    Write-Host "  .\$FileName" -ForegroundColor Cyan
    Write-Host "`nAvailable files in global_func:" -ForegroundColor Yellow
    Get-ChildItem ".\global_func\*.py" -ErrorAction SilentlyContinue | ForEach-Object { Write-Host "  $($_.Name)" -ForegroundColor Cyan }
    Write-Host "`nAvailable files in root:" -ForegroundColor Yellow
    Get-ChildItem ".\*.py" -ErrorAction SilentlyContinue | ForEach-Object { Write-Host "  $($_.Name)" -ForegroundColor Cyan }
    exit 1
}

Write-Host "Updating file: $FileName -> $remoteFileName" -ForegroundColor Green
Write-Host "Local path: $localFilePath" -ForegroundColor Cyan
Write-Host "Target path: $TargetPath$remoteFileName" -ForegroundColor Cyan

# Get instance IPs and IDs
. .\get-ips.ps1
$instanceIds = aws autoscaling describe-auto-scaling-groups --auto-scaling-group-names portfolio-manager-asg --query 'AutoScalingGroups[0].Instances[?HealthStatus==`Healthy`].InstanceId' --output text

if ($global:instances.Count -eq 0) {
    Write-Host "No instances found" -ForegroundColor Red
    exit 1
}

Write-Host "Found $($global:instances.Count) instances: $($global:instances -join ', ')" -ForegroundColor Green

# Generate a temporary SSH key for this session
Write-Host "Generating temporary SSH key..." -ForegroundColor Cyan
ssh-keygen -t rsa -f temp_update_key -N '""' -q

$instanceIdArray = $instanceIds -split "`t"

# Update each instance one at a time to maintain uptime
for ($i = 0; $i -lt $global:instances.Count; $i++) {
    $ip = $global:instances[$i]
    $instanceId = $instanceIdArray[$i]
    
    Write-Host "Updating instance $instanceId ($ip)..." -ForegroundColor Yellow
    
    # Send SSH public key via EC2 Instance Connect
    Write-Host "Sending SSH key to $instanceId..."
    aws ec2-instance-connect send-ssh-public-key --instance-id $instanceId --instance-os-user ec2-user --ssh-public-key file://temp_update_key.pub
    
    if ($LASTEXITCODE -eq 0) {
        # Copy the file using SCP
        Write-Host "Copying $FileName -> $remoteFileName to $ip..."
        scp -i temp_update_key -o StrictHostKeyChecking=no "$localFilePath" ec2-user@${ip}:$TargetPath$remoteFileName
        
        if ($LASTEXITCODE -eq 0) {
            Write-Host "$remoteFileName copied successfully to $ip" -ForegroundColor Green
            
            # Only restart Streamlit if we're updating a Python file
            if ($remoteFileName -like "*.py") {
                Write-Host "Restarting Streamlit service on $ip..."
                ssh -i temp_update_key -o StrictHostKeyChecking=no ec2-user@$ip "pkill -f 'streamlit run' && cd /home/ec2-user/AWS_Portfolio_Manager && nohup ./venv/bin/streamlit run application.py --server.port 5000 --server.address 0.0.0.0 > /dev/null 2>&1 &"
                
                if ($LASTEXITCODE -eq 0) {
                    Write-Host "Service restarted successfully on $ip" -ForegroundColor Green
                } else {
                    Write-Host "Failed to restart service on $ip" -ForegroundColor Red
                }
            } else {
                Write-Host "Non-Python file - skipping service restart" -ForegroundColor Yellow
            }
        } else {
            Write-Host "Failed to copy $remoteFileName to $ip" -ForegroundColor Red
        }
    } else {
        Write-Host "Failed to send SSH key to $instanceId" -ForegroundColor Red
    }
    
    Write-Host "Waiting 15 seconds before next instance..." -ForegroundColor Cyan
    Start-Sleep -Seconds 15
}

# Clean up temporary key files
Remove-Item temp_update_key, temp_update_key.pub -ErrorAction SilentlyContinue

Write-Host "All instances updated with rolling restart!" -ForegroundColor Green