"""Check that Airflow environment successfully updated"""
import datetime
import os
import time

import boto3

environment_name = os.environ.get('ENVIRONMENT_NAME')
aws_region = os.environ.get('AWS_REGION', 'us-east-1')
update_timeout = int(os.environ.get('UPDATE_TIMEOUT', 1800))


def poll_airflow_status():
    """Poll service to see if task has launched."""
    client = boto3.client('mwaa', region_name=aws_region)
    timeout = time.time() + update_timeout
    print('\nPolling for new environment status. Timeout: {} seconds'.format(
        update_timeout))
    while time.time() < timeout:
        environment = client.get_environment(
            Name=environment_name)['Environment']
        if environment['LastUpdate']['Status'] == 'SUCCESS':
            print('Environment update successful.')
            return True
        elif environment['LastUpdate']['Status'] == 'PENDING':
            print('Environment still in PENDING status.')
            time.sleep(30)
            continue
        else:
            print('Environment update failed. Please troubleshoot and recreate environment.')
            print(environment['LastUpdate']['Error'])
            raise SystemExit()
    raise SystemExit('Environment update failed by timeout(30 minutes). Please troubleshoot and recreate environment.')


def main():
    """Main entrypoint function."""
    poll_airflow_status()


if __name__ == "__main__":
    main()