#!/bin/bash

# Launch cheap AWS spot instance for testing OPTION A ensemble
# Searches for cheapest available spot instances

set -e

echo "========================================="
echo "AWS Spot Instance Launcher - Test OPTION A"
echo "========================================="

# Configuration
INSTANCE_TYPES=(
    "t3.medium"      # 2 vCPU, 4GB RAM - ~$0.01/hr
    "t3a.medium"     # 2 vCPU, 4GB RAM - ~$0.009/hr (cheaper AMD)
    "t3.large"       # 2 vCPU, 8GB RAM - ~$0.02/hr
    "t3a.large"      # 2 vCPU, 8GB RAM - ~$0.018/hr
    "c6a.large"      # 2 vCPU, 4GB RAM, compute optimized - ~$0.015/hr
)

AMI_ID="ami-0c55b159cbfafe1f0"  # Ubuntu 22.04 LTS (us-east-1)
REGION="us-east-1"
KEY_NAME="ensemble-test-key"
SECURITY_GROUP="ensemble-test-sg"

# Colors
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'

# Function to get spot price
get_spot_price() {
    local instance_type=$1
    aws ec2 describe-spot-price-history \
        --instance-types "$instance_type" \
        --product-descriptions "Linux/UNIX" \
        --region "$REGION" \
        --max-results 1 \
        --query 'SpotPriceHistory[0].SpotPrice' \
        --output text
}

# Find cheapest instance
echo -e "\n${YELLOW}Finding cheapest spot instance...${NC}"
echo ""

cheapest_type=""
cheapest_price=999999
prices=()

for instance_type in "${INSTANCE_TYPES[@]}"; do
    price=$(get_spot_price "$instance_type")

    if [ -n "$price" ]; then
        echo "  $instance_type: \$$price/hr"
        prices+=("$instance_type:$price")

        # Check if cheaper
        if (( $(echo "$price < $cheapest_price" | bc -l) )); then
            cheapest_price=$price
            cheapest_type=$instance_type
        fi
    fi
done

echo ""
echo -e "${GREEN}Cheapest: $cheapest_type at \$$cheapest_price/hr${NC}"
echo ""

# Confirm
read -p "Launch $cheapest_type spot instance? (y/n) " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Cancelled."
    exit 0
fi

# Create key pair if doesn't exist
echo ""
echo "Checking SSH key..."
if ! aws ec2 describe-key-pairs --key-names "$KEY_NAME" --region "$REGION" &> /dev/null; then
    echo "Creating key pair..."
    aws ec2 create-key-pair \
        --key-name "$KEY_NAME" \
        --region "$REGION" \
        --query 'KeyMaterial' \
        --output text > ~/.ssh/${KEY_NAME}.pem
    chmod 400 ~/.ssh/${KEY_NAME}.pem
    echo -e "${GREEN}✓ Key created: ~/.ssh/${KEY_NAME}.pem${NC}"
else
    echo -e "${GREEN}✓ Key exists${NC}"
fi

# Create security group if doesn't exist
echo ""
echo "Checking security group..."
if ! aws ec2 describe-security-groups --group-names "$SECURITY_GROUP" --region "$REGION" &> /dev/null; then
    echo "Creating security group..."

    vpc_id=$(aws ec2 describe-vpcs \
        --region "$REGION" \
        --filters "Name=isDefault,Values=true" \
        --query 'Vpcs[0].VpcId' \
        --output text)

    sg_id=$(aws ec2 create-security-group \
        --group-name "$SECURITY_GROUP" \
        --description "Security group for ensemble testing" \
        --vpc-id "$vpc_id" \
        --region "$REGION" \
        --query 'GroupId' \
        --output text)

    # Allow SSH
    aws ec2 authorize-security-group-ingress \
        --group-id "$sg_id" \
        --protocol tcp \
        --port 22 \
        --cidr 0.0.0.0/0 \
        --region "$REGION"

    echo -e "${GREEN}✓ Security group created${NC}"
else
    echo -e "${GREEN}✓ Security group exists${NC}"
fi

# Create user data script
cat > /tmp/user-data.sh << 'USERDATA'
#!/bin/bash

# Update system
apt-get update
apt-get install -y python3.10 python3-pip git

# Install dependencies
pip3 install --upgrade pip

# Clone repository
cd /home/ubuntu
git clone https://huggingface.co/marcosremar2/ensemble-tts-annotation
cd ensemble-tts-annotation

# Install requirements
pip3 install -r requirements.txt

# Create results directory
mkdir -p /home/ubuntu/test-results

echo "✅ Setup complete"
echo "Run: python3 scripts/test/test_quick.py > /home/ubuntu/test-results/test.log 2>&1"
USERDATA

# Request spot instance
echo ""
echo "Requesting spot instance..."

MAX_PRICE=$(echo "$cheapest_price * 1.5" | bc)

spot_request=$(aws ec2 request-spot-instances \
    --instance-count 1 \
    --type "one-time" \
    --launch-specification "{
        \"ImageId\": \"$AMI_ID\",
        \"InstanceType\": \"$cheapest_type\",
        \"KeyName\": \"$KEY_NAME\",
        \"SecurityGroups\": [\"$SECURITY_GROUP\"],
        \"UserData\": \"$(base64 /tmp/user-data.sh | tr -d '\n')\"
    }" \
    --spot-price "$MAX_PRICE" \
    --region "$REGION" \
    --query 'SpotInstanceRequests[0].SpotInstanceRequestId' \
    --output text)

echo -e "${GREEN}✓ Spot request created: $spot_request${NC}"

# Wait for fulfillment
echo ""
echo "Waiting for spot instance to launch..."

while true; do
    status=$(aws ec2 describe-spot-instance-requests \
        --spot-instance-request-ids "$spot_request" \
        --region "$REGION" \
        --query 'SpotInstanceRequests[0].Status.Code' \
        --output text)

    if [ "$status" == "fulfilled" ]; then
        echo -e "${GREEN}✓ Spot instance fulfilled!${NC}"
        break
    elif [ "$status" == "price-too-low" ] || [ "$status" == "capacity-not-available" ]; then
        echo -e "${RED}✗ Spot request failed: $status${NC}"
        exit 1
    fi

    echo "  Status: $status"
    sleep 5
done

# Get instance ID
instance_id=$(aws ec2 describe-spot-instance-requests \
    --spot-instance-request-ids "$spot_request" \
    --region "$REGION" \
    --query 'SpotInstanceRequests[0].InstanceId' \
    --output text)

echo "Instance ID: $instance_id"

# Wait for instance to be running
echo ""
echo "Waiting for instance to be running..."
aws ec2 wait instance-running --instance-ids "$instance_id" --region "$REGION"

# Get public IP
public_ip=$(aws ec2 describe-instances \
    --instance-ids "$instance_id" \
    --region "$REGION" \
    --query 'Reservations[0].Instances[0].PublicIpAddress' \
    --output text)

echo ""
echo "========================================="
echo -e "${GREEN}✓ Instance launched successfully!${NC}"
echo "========================================="
echo ""
echo "Instance Type: $cheapest_type"
echo "Cost: ~\$$cheapest_price/hr"
echo "Instance ID: $instance_id"
echo "Public IP: $public_ip"
echo ""
echo "SSH Command:"
echo "  ssh -i ~/.ssh/${KEY_NAME}.pem ubuntu@$public_ip"
echo ""
echo "Run test:"
echo "  ssh -i ~/.ssh/${KEY_NAME}.pem ubuntu@$public_ip 'cd ensemble-tts-annotation && python3 scripts/test/test_quick.py'"
echo ""
echo "Stop instance:"
echo "  aws ec2 terminate-instances --instance-ids $instance_id --region $REGION"
echo ""
echo "========================================="

# Save instance info
cat > /tmp/spot-instance-info.txt << EOF
Instance Type: $cheapest_type
Cost: \$$cheapest_price/hr
Instance ID: $instance_id
Public IP: $public_ip
SSH Key: ~/.ssh/${KEY_NAME}.pem
Region: $REGION

SSH: ssh -i ~/.ssh/${KEY_NAME}.pem ubuntu@$public_ip
Terminate: aws ec2 terminate-instances --instance-ids $instance_id --region $REGION
EOF

echo "Instance info saved to: /tmp/spot-instance-info.txt"
echo ""