import boto3
import argparse
import datetime
import collections

def automated_ami_backup_and_retention(
    instance_tag_key,
    instance_tag_value,
    retention_days,
    region_name='us-east-1'
):
    """
    Automates the creation of AMIs for EC2 instances identified by a tag
    and manages their retention by deleting old AMIs and their associated snapshots.

    Args:
        instance_tag_key (str): The tag key to identify target EC2 instances for backup.
        instance_tag_value (str): The tag value to identify target EC2 instances for backup.
        retention_days (int): The number of days to retain AMIs. AMIs older than this will be deleted.
        region_name (str): The AWS region.
    """
    ec2_client = boto3.client('ec2', region_name=region_name)

    print(f"Starting automated AMI backup and retention in region {region_name}...")

    # ==========================================================================
    # STEP 1: Identify EC2 instances for backup based on tags.
    # ==========================================================================
    print("\n>>> Step 1: Identifying EC2 instances for backup...")
    instances_to_backup = []
    try:
        response = ec2_client.describe_instances(
            Filters=[
                {'Name': f'tag:{instance_tag_key}', 'Values': [instance_tag_value]},
                {'Name': 'instance-state-name', 'Values': ['running', 'stopped']}
            ]
        )
        for reservation in response['Reservations']:
            for instance in reservation['Instances']:
                instance_id = instance['InstanceId']
                instance_name = 'N/A'
                for tag in instance.get('Tags', []):
                    if tag['Key'] == 'Name':
                        instance_name = tag['Value']
                        break
                instances_to_backup.append({'id': instance_id, 'name': instance_name})
        
        if not instances_to_backup:
            print(f"   No EC2 instances found with tag {instance_tag_key}={instance_tag_value}. Exiting.")
            return
        print(f"   Found {len(instances_to_backup)} instances to backup.")

    except Exception as e:
        print(f"Error identifying instances: {e}")
        return

    # ==========================================================================
    # STEP 2: Create AMIs for identified instances.
    # ==========================================================================
    print("\n>>> Step 2: Creating AMIs for identified instances...")
    created_amis = []
    for instance in instances_to_backup:
        instance_id = instance['id']
        instance_name = instance['name']
        ami_name = f"{instance_name}-{instance_id}-AMI-{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}"
        ami_description = f"Automated AMI backup for {instance_name} ({instance_id}) on {datetime.date.today().isoformat()}"
        
        try:
            print(f"   Creating AMI for instance {instance_id} ({instance_name})...")
            create_ami_response = ec2_client.create_image(
                InstanceId=instance_id,
                Name=ami_name,
                Description=ami_description,
                NoReboot=True, # Set to False if you want to reboot the instance during AMI creation
                TagSpecifications=[
                    {
                        'ResourceType': 'image',
                        'Tags': [
                            {'Key': 'Name', 'Value': ami_name},
                            {'Key': 'CreatedBy', 'Value': 'AutomatedAMIBackupScript'},
                            {'Key': 'InstanceId', 'Value': instance_id},
                            {'Key': 'RetentionDays', 'Value': str(retention_days)}
                        ]
                    },
                ]
            )
            ami_id = create_ami_response['ImageId']
            created_amis.append(ami_id)
            print(f"   AMI {ami_id} created for {instance_id}.")
        except Exception as e:
            print(f"Error creating AMI for {instance_id}: {e}")

    print(f"   Successfully initiated creation of {len(created_amis)} AMIs.")

    # ==========================================================================
    # STEP 3: Manage AMI retention - deregister old AMIs and delete their associated snapshots.
    # ==========================================================================
    print(f"\n>>> Step 3: Managing AMI retention (deleting AMIs older than {retention_days} days)...")
    try:
        # Get all AMIs created by this script
        all_amis = ec2_client.describe_images(
            Filters=[
                {'Name': 'tag:CreatedBy', 'Values': ['AutomatedAMIBackupScript']},
                {'Name': 'state', 'Values': ['available']}
            ],
            Owners=['self']
        )['Images']

        if not all_amis:
            print("   No AMIs created by this script found for retention management.")
            return

        now = datetime.datetime.now(datetime.timezone.utc)
        deleted_ami_count = 0
        deleted_snapshot_count = 0

        for ami in all_amis:
            ami_id = ami['ImageId']
            creation_date_str = ami['CreationDate'] # Format: YYYY-MM-DDTHH:MM:SS.000Z
            creation_date = datetime.datetime.strptime(creation_date_str, '%Y-%m-%dT%H:%M:%S.%fZ').replace(tzinfo=datetime.timezone.utc)
            
            age = now - creation_date

            if age.days > retention_days:
                print(f"   AMI {ami_id} ({ami['Name']}) is {age.days} days old (older than {retention_days} days). Initiating deletion...")
                
                # Deregister the AMI
                ec2_client.deregister_image(ImageId=ami_id)
                print(f"      Deregistered AMI: {ami_id}.")
                deleted_ami_count += 1

                # Delete associated snapshots
                for bdm in ami['BlockDeviceMappings']:
                    if 'Ebs' in bdm and 'SnapshotId' in bdm['Ebs']:
                        snapshot_id = bdm['Ebs']['SnapshotId']
                        print(f"      Deleting associated snapshot: {snapshot_id}.")
                        ec2_client.delete_snapshot(SnapshotId=snapshot_id)
                        deleted_snapshot_count += 1
            else:
                print(f"   AMI {ami_id} ({ami['Name']}) is {age.days} days old (within retention). Skipping deletion.")

        print(f"   Successfully deregistered {deleted_ami_count} AMIs and deleted {deleted_snapshot_count} associated snapshots.")

    except Exception as e:
        print(f"Error managing AMI retention: {e}")

    print("\nAutomated AMI backup and retention completed.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Automate EC2 AMI backup and retention policy.")
    parser.add_argument("--instance-tag-key", required=True, help="Tag key to identify target EC2 instances for backup (e.g., 'Backup').")
    parser.add_argument("--instance-tag-value", required=True, help="Tag value to identify target EC2 instances for backup (e.g., 'true').")
    parser.add_argument("--retention-days", type=int, default=30, help="Number of days to retain AMIs (default: 30).")
    parser.add_argument("--region", default="us-east-1", help="AWS region (default: us-east-1).")

    args = parser.parse_args()

    automated_ami_backup_and_retention(
        instance_tag_key=args.instance_tag_key,
        instance_tag_value=args.instance_tag_value,
        retention_days=args.retention_days,
        region_name=args.region
    )