#!/usr/bin/env bash

# Scale VM Scale Set
# This script scales a VM Scale Set up or down

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/../lib/init.sh"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

# Configuration
RESOURCE_GROUP="${RESOURCE_GROUP:-defi-oracle-mainnet-rg}"
VMSS_NAME="${1:-}"
NEW_CAPACITY="${2:-}"


if [ -z "$VMSS_NAME" ] || [ -z "$NEW_CAPACITY" ]; then
    log_error "Error: VMSS name and capacity required"
    echo "Usage: $0 <vmss-name> <new-capacity>"
    echo "Example:"
    echo "  $0 besu-rpc-vmss 5"
    exit 1
fi

# Validate capacity
if ! [[ "$NEW_CAPACITY" =~ ^[0-9]+$ ]]; then
    log_error "Error: Capacity must be a number"
    exit 1
fi

if [ "$NEW_CAPACITY" -lt 1 ]; then
    log_error "Error: Capacity must be at least 1"
    exit 1
fi

log_success "Scaling VMSS: $VMSS_NAME to $NEW_CAPACITY instances"

# Get current capacity
CURRENT_CAPACITY=$(az vmss show --resource-group "$RESOURCE_GROUP" --name "$VMSS_NAME" --query "sku.capacity" -o tsv 2>/dev/null || echo "0")

if [ "$CURRENT_CAPACITY" == "0" ]; then
    log_error "Error: VMSS not found: $VMSS_NAME"
    exit 1
fi

log_warn "Current capacity: $CURRENT_CAPACITY"
log_warn "New capacity: $NEW_CAPACITY"

if [ "$CURRENT_CAPACITY" -eq "$NEW_CAPACITY" ]; then
    log_warn "Capacity is already $NEW_CAPACITY"
    exit 0
fi

# Confirm scaling
read -p "Do you want to scale from $CURRENT_CAPACITY to $NEW_CAPACITY? (y/N) " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Scaling cancelled"
    exit 0
fi

# Scale VMSS
log_warn "Scaling VMSS..."
az vmss scale \
    --resource-group "$RESOURCE_GROUP" \
    --name "$VMSS_NAME" \
    --new-capacity "$NEW_CAPACITY"

if [ $? -eq 0 ]; then
    log_success "✓ VMSS scaled successfully"
    
    # Wait for instances to be ready
    log_warn "Waiting for instances to be ready..."
    sleep 30
    
    # Check instance status
    INSTANCES=$(az vmss list-instances --resource-group "$RESOURCE_GROUP" --name "$VMSS_NAME" --query "[].{Name:name, ProvisioningState:provisioningState}" -o json 2>/dev/null || echo "[]")
    echo "$INSTANCES" | jq -r '.[] | "\(.Name): \(.ProvisioningState)"'
else
    log_error "✗ Failed to scale VMSS"
    exit 1
fi

log_success "Scaling complete!"