"""
$ python update.py

Updates an AWS Lambda function with all the latest goodies.

To use, create a `lambda` directory containing all the files and requirements required 
to execute your Lambda function. Then, place the contents of this file into an `update.py`
file in the root of your project. Setup your AWS CLI profile name and AWS Lambda ARN as
environment variables, then run `python update.py`. It will:

1. Create a ZIP file of your `lambda` directory and place it in a `package` directory
2. Upload the ZIP file (`package/function.zip`) as your function's code to AWS Lambda
3. Poll Lambda to let you know when the code is ready to use
"""

import os
import boto3
import shutil

AWS_PROFILE = os.environ['AWS_PROFILE']
AWS_LAMBDA_ARN = os.environ['AWS_LAMBDA_ARN']

boto3.setup_default_session(profile_name=AWS_PROFILE)
lambda_client = boto3.client('lambda')

def create_zip_archive():
    """
    Creates a ZIP archive of the "lambda" directory
    """
    print('📦 Hang tight! Creating ZIP archive...')
    filename = 'function'
    shutil.make_archive(f'./package/{filename}', 'zip', 'lambda')
    print('📦 ZIP archive created!')
    return f'package/{filename}.zip'
    
def update_lambda(filename):
    """
    Updates the Lambda function with the latest code
    """
    print('🚀 Uploading Lambda function...')
    with open(filename, 'rb') as f:
        lambda_client.update_function_code(
            FunctionName=AWS_LAMBDA_ARN,
            ZipFile=f.read()
        )
    print('✅ Lambda function updated')

def poll_lambda_status():
    """
    Uses a waiter to poll Lambda for up to 100 seconds (5 * 20) to check
    your function's availability
    """
    print('⏰ Waiting for function to become available...')
    waiter = lambda_client.get_waiter('function_updated')
    try:
        waiter.wait(
            FunctionName=AWS_LAMBDA_ARN,
            WaiterConfig={
                'Delay': 5,
                'MaxAttempts': 20
            }
        )
    except:
        print('🤔 Hmm... the Lambda function failed to update')
    else:
        print('💥 Lambda function is now active!')
    
def kickoff():
    f = create_zip_archive()
    update_lambda(f)
    poll_lambda_status()

if __name__ == '__main__':
    kickoff()