import boto3, os, time
from pathlib import Path
d={}
for l in Path('/opt/desineuron-ops-control-plane/.env').read_text().splitlines():
  if '=' in l and not l.startswith('#'):
    k,v=l.split('=',1)
    d[k.strip()]=v.strip()
os.environ['AWS_ACCESS_KEY_ID']=d.get('AWS_ACCESS_KEY_ID','')
os.environ['AWS_SECRET_ACCESS_KEY']=d.get('AWS_SECRET_ACCESS_KEY','')
ec2=boto3.client('ec2', region_name='us-east-1')

def get_gpu():
    for r in ec2.describe_instances()['Reservations']:
        for i in r['Instances']:
            if any(t['Key'] == 'Name' and t['Value'] == 'desineuron-comfy-gpu' for t in i.get('Tags', [])):
                return i
    return None

def main():
    while True:
        i = get_gpu()
        if not i:
            print('Not found')
            break
        state = i['State']['Name']
        print(f"Instance {i['InstanceId']} is {state}")
        if state == 'stopped':
            print('Starting instance...')
            ec2.start_instances(InstanceIds=[i['InstanceId']])
            time.sleep(5)
        elif state == 'stopping':
            print('Waiting for extremely aggressive stop sequence gracefully...')
            time.sleep(10)
        elif state == 'running':
            print('Instance successfully running payload on IP:', i.get('PrivateIpAddress'))
            break
        else:
            print('Waiting eagerly...')
            time.sleep(10)

if __name__ == '__main__':
    main()