"""Update ECR reference in containerized lambda funcion to force update."""

import os

from botocore.exceptions import ClientError
import sys

import boto3


AWS_REGION = os.environ.get('AWS_REGION', 'us-east-1')
LAMBDA_DIR = os.environ.get('LAMBDA_DIR')
ENV = os.environ.get('ENV')
PREFIX = os.environ.get('PREFIX', 'lambda')
ECR_ACCOUNT_ID = os.environ.get(
    'ECR_REGISTRY_ACCOUNT_ID', '437795906767')
ECR_REGION = os.environ.get('ECR_REGISTRY_REGION', 'us-east-1')
ECR_REGISTRY_URL = f'{ECR_ACCOUNT_ID}.dkr.ecr.{ECR_REGION}.amazonaws.com'
IMAGE_NAME = os.environ.get(
    'IMAGE_NAME_OVERRIDE', f"{PREFIX}-{LAMBDA_DIR.replace('_', '-')}")


def main():
    """Entrypoint."""
    try:
        new_tag = sys.argv[1]
    except IndexError:
        raise SystemExit('\nargv[1] containing new tag is required')
    function_name = get_function_name()
    image_uri = f'{ECR_REGISTRY_URL}/{IMAGE_NAME}:{new_tag}'
    lambda_client = boto3.client('lambda', region_name=AWS_REGION)

    print(f'\nFetching function {function_name}')
    response = lambda_client.get_function(FunctionName=function_name)
    if response['Code']['RepositoryType'] != 'ECR':
        raise SystemExit('\nLambda function is not configured to use ECR')

    print(f'\nValidating image {image_uri} exists')
    ecr_client = boto3.client('ecr', region_name=AWS_REGION)
    try:
        ecr_client.describe_images(
            registryId=ECR_ACCOUNT_ID,
            repositoryName=IMAGE_NAME,
            imageIds=[
                {'imageTag': new_tag}
            ]
        )
    except ClientError:
        raise SystemExit(
            f'\nFailed to validate existence of image {image_uri}. '
            f'ECR Repository does not exist or does not permit access.'
        )

    print(f'\nUpdating {function_name} image URI to {image_uri}')
    response = lambda_client.update_function_code(
        FunctionName=function_name, ImageUri=image_uri)
    if response['ResponseMetadata']['HTTPStatusCode'] != 200:
        print(f"\n{response['ResponseMetadata']}")
        raise SystemExit('\nError updating lambda, non-200 response')


def get_function_name():
    """Parse out the function name."""
    if len(sys.argv) > 2:
        return sys.argv[2]
    assert ENV, 'ENV must be set.'
    assert PREFIX is not None, 'PREFIX must be set.'
    assert LAMBDA_DIR, 'LAMBDA_DIR must be set.'
    parts = [
        ENV,
        PREFIX,
        LAMBDA_DIR.replace('_', '-')
    ]
    return '-'.join([x for x in parts if x])


if __name__ == '__main__':
    main()