from botocore.exceptions import ClientError
import boto3
import time
import datetime
import logging
import sys
import os

def get_tags():
    return [
        {
            'Key': 'platform',
            'Value': 'Delphi'
        },
        {
            'Key': 'environment',
            'Value': 'Common'
        },
        {
            'Key': 'service',
            'Value': 'EC2'
        },
        {
            'Key': 'project',
            'Value': 'Infrastructure'
        },
        {
            'Key': 'plat_env_project_service',
            'Value': 'DLP_CMN_INFRA_EC2'
        },
    ]

def remove_image_with_snapshots(image_id, ec2, loggerName = 'Main'):
    """
    Deregisters image and invoke removing snapshots based on block device mappings.
    Requires ec2 resource object to run
    """
    logger = logging.getLogger(loggerName)
    image = ec2.Image(image_id)
    mappings = image.block_device_mappings
    logger.info(f"Deregistering image {image_id}")
    image.deregister()
    remove_snapshots_bdm(mappings, ec2)

def remove_snapshots_bdm(block_device_mappings, ec2, loggerName = 'Main'):
    """
    Remove snapshots from block device mappings.
    Requires ec2 resource object to run
    """
    logger = logging.getLogger(loggerName)
    for block_device in block_device_mappings:
        ebs = block_device.get('Ebs', None)
        snapshot_id = ebs.get('SnapshotId', None) if ebs else None
        if snapshot_id:
            logger.info(
                f"Removing snapshot {snapshot_id}")
            snapshot = ec2.Snapshot(snapshot_id)
            try:
                snapshot.delete()
            except ClientError as e:
                if e.response['Error']['Code'] == 'InvalidSnapshot.InUse':
                    logger.info(e.response['Error']['Message'])
                elif e.response['Error']['Code'] == 'InvalidSnapshot.NotFound':
                    logger.info(e.response['Error']['Code'])
                else:
                    raise

def tag_ami_snapshots(block_device_mappings, ec2, loggerName = 'Main'):
    """
    Tag snapshots from block device mappings.
    Requires ec2 resource object to run
    """
    logger = logging.getLogger(loggerName)
    for block_device in block_device_mappings:
        ebs = block_device.get('Ebs', None)
        snapshot_id = ebs.get('SnapshotId', None) if ebs else None
        if snapshot_id:
            logger.info(
                f"Tagging snapshot {snapshot_id}")
            snapshot = ec2.Snapshot(snapshot_id)
            snapshot.create_tags(Tags=get_tags())

def get_instance_name(instance):
    """
    Get instance name from the name tag,
    return empty string if tag doesn't exist.
    """
    try:
        instance_name = next(
            filter(lambda k: k['Key'] == 'Name', instance.tags)
            )['Value']
        return instance_name
    except StopIteration:
        return ""

def create_image(
        ec2_client, instance, ec2,
        date_format = "%Y-%m-%d/%H-%M",
        loggerName = 'Main',
        noreboot = True,
        dryrun = False,
        wait_image_ready = False):
    """
    Creates image for selected instance and set 'Retention' and 'Instance_name'
    tags on it. Function applying retention policy will use these tags
    to filter images. By default image is made with noreboot option.
    Also you could wait until image is available with option wait_image_ready.
    """
    logger = logging.getLogger(loggerName)
    dstamp = datetime.datetime.now().strftime(date_format)
    instance_name = get_instance_name(instance)
    if not instance_name:
        logger.info(
            f"Name tag doesn't exist for {instance.id}, using id as name")
        instance_name = instance.id
    try:
        image = instance.create_image(
            Description=f"Created by backup lambda from "
                        f"Source: {instance.id} Name: {instance_name}",
            Name=f"{instance_name}_{dstamp}",
            NoReboot=True,
            DryRun=dryrun)
        if not dryrun:
            if wait_image_ready:
                waiter = ec2_client.get_waiter('image_available')
                waiter.wait(ImageIds = [image.image_id])
            image.create_tags(Tags=[{'Key': 'Retention', 'Value':'Enabled'}])
            image.create_tags(Tags=[{'Key': 'Instance_Name', 'Value':instance_name}])
            image.create_tags(Tags=get_tags())
            image.reload()
            tag_ami_snapshots(image.block_device_mappings, ec2)
    except ClientError as e:
        print(e.response['Error']['Message'])

def get_expired_amis(
        ec2, instance,
        account_id, retention_limit = 3,
        loggerName = 'Main'):
    """
    Get the list of amis that are not fitting retention_limit.
    Sorting is made via creation date
    """
    logger = logging.getLogger(loggerName)
    instance_name = get_instance_name(instance)
    filters = [
        {
            'Name': 'tag:Retention',
            'Values': ['Enabled']
        },
        {
            'Name': 'tag:Instance_Name',
            'Values': [instance_name]
        }
    ]
    image_list = sorted(
        list(ec2.images.filter(Filters=filters, Owners=[account_id])),
        key=lambda k: k.creation_date)
    logger.info(f"All images:\n {[x.name for x in image_list]}" )
    remove_index = len(image_list) - retention_limit
    if remove_index > 0:
        return tuple(image_list[:remove_index])

def lambda_handler(event, context):
    # Get environment variables, default values if not
    # TODO remove default values; add logger name; set default logger level;

    # Set up logging. Use INFO by default, DEBUG if DEBUG has been set to true.
    debug = os.getenv('DEBUG', 'FALSE').upper() == 'TRUE'
    logger = logging.getLogger('Main')
    if debug:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)

    # Dryrun is disabled by default.
    dryrun = os.getenv('DRYRUN', 'FALSE').upper() == 'TRUE'

    # Retention limit is mandatory env variable
    try:
        RETENTION_LIMIT = int(os.getenv('RETENTION_LIMIT'))
    except:
        logger.error("No retention limit specified. Set 'RETENTION_LIMIT' environment variable to proceed")
        exit(1)

    # INSTANCES is mandatory too
    INSTANCES = os.getenv('INSTANCES', '').lower().replace(' ','')
    if not INSTANCES:
        logger.error("No instances specified. Set 'INSTANCES' environment variable")
        exit(1)

    # Initialize client for autoscale and resource for ec2
    logger.debug(f"Event:\n{event}\nContext:\n{context}")
    logger.info("Creating clients and resources...")
    ec2 = boto3.resource('ec2')
    ec2_client = boto3.client('ec2')
    account_id = boto3.client('sts').get_caller_identity().get('Account')
    logger.info(f"Account id is {account_id}")
    instances = INSTANCES.split(',')

    for instance_id in instances:
        if not instance_id:
            logger.error("""Empty string was given as an instance id.
Check your INSTANCES variable - it should be comma separated list of instance IDs.
""")
            break

        instance = ec2.Instance(instance_id)
        create_image(ec2_client, instance, ec2, dryrun=dryrun)
        images_to_delete = get_expired_amis(
            ec2, instance, account_id,
            retention_limit=RETENTION_LIMIT)
        if images_to_delete:
            logger.info(f"""
List of images which will be deregistred:
{[image.name for image in images_to_delete]}
            """)
            if not dryrun:
                for image in images_to_delete:
                    remove_image_with_snapshots(image.id, ec2)