import boto3
import requests
import time

METADATA_URL = "http://169.254.169.254/latest/meta-data"

def check_spot_interruption():
    """
    Poll for spot interruption notice.
    AWS gives us 2 minutes warning—use it wisely.
    """
    try:
        resp = requests.get(
            f"{METADATA_URL}/spot/instance-action",
            timeout=1
        )
        if resp.status_code == 200:
            return resp.json()
    except requests.exceptions.RequestException:
        pass
    return None

def handle_interruption(action_data):
    """
    Graceful shutdown sequence:
    1. Stop accepting new requests
    2. Drain existing connections
    3. Checkpoint any in-flight work
    4. Signal ready for termination
    """
    print(f"⚠️ Spot interruption in {action_data['time']}")
    
    # 1. Deregister from load balancer
    deregister_from_alb()
    
    # 2. Wait for in-flight requests (max 30s)
    drain_connections(timeout=30)
    
    # 3. Flush metrics and logs
    flush_telemetry()
    
    # 4. Notify orchestrator
    notify_replacement_needed()
    
    print("✓ Graceful shutdown complete")

# Main loop - check every 5 seconds
while True:
    interruption = check_spot_interruption()
    if interruption:
        handle_interruption(interruption)
        break
    time.sleep(5)