import boto3
import datetime
from time import sleep

from asgebsmanager.mds_utils import get_imds_token, get_instance_metadata


class AsgEbsManager:

    def __init__(self,region,asg_name):
        self.asg_name = asg_name
        self.session = boto3.Session(region_name=region)
        self.ec2_client = self.session.client('ec2')

    def __get_name():
        return 'AsgEbsManager'
    
    def latest_volume(self,az,state=None):
        status = ["available","in-use"] if state is None else [state]
        volumes = self.ec2_client.client.describe_volumes(
            Filters=[
                {
                    'Name': 'tag:BelongsTo',
                    'Values': [
                        self.asg_name,
                    ]
                },
                {
                    'Name': 'tag:ManagedBy',
                    'Values': [
                        self.__get_name(),
                    ]
                },
                {
                    'Name': 'availability-zone',
                    'Values': [
                        az,
                    ]
                },
                {
                    'Name': 'status',
                    'Values': [
                        status,
                    ]
                }
            ],
            MaxResults=100
        )

        if len(volumes['Volumes']) > 0:
            sorted_volume = [volume['VolumeId'] for volume in sorted(volumes['Volumes'], key=lambda volume: datetime.strptime(volume['CreateTime'], '%Y-%m-%dT%H:%M:%S.%f%z'))]
            return sorted_volume[0]
        return None

    def create_volume(self,az,size,iops,throughput,snapshot_id=None):
        response = self.ec2_client.create_volume(
            SnapshotId=snapshot_id,
            AvailabilityZone=az,
            Encrypted=True,
            Iops=iops,
            Size=size,
            VolumeType='gp3',
            Throughput=throughput,
        )
        return response['VolumeId']
    
    def describe_volume(self,volume_id):
        args = dict(
            Filters=[
                {
                    'Name': 'tag:BelongsTo',
                    'Values': [
                        self.asg_name,
                    ]
                },
                {
                    'Name': 'tag:ManagedBy',
                    'Values': [
                        self.__get_name(),
                    ]
                }
            ],
            MaxResults=1,
            VolumeIds=[volume_id]
        )
        volumes = self.ec2_client.client.describe_volumes(**args)['Volumes']
        volume = None
        if len(volumes) > 0:
            volume = volumes[0]
        if volume is None: 
            return None 
        return volume
    


    def attach_volume(self,volume_id,instance_id,device):
        response = self.ec2_client.attach_volume(
            Device=device,
            InstanceId=instance_id,
            VolumeId=volume_id
        )
        return response['VolumeId']
    
    def delete_volume(self,volume_id):
        self.ec2_client.delete_volume(
            VolumeId=volume_id
        )

    def create_snapshot(self,volume_id):
        response = self.ec2_client.create_snapshot(
            Description=f'SnapshotOf{volume_id}',
            VolumeId=volume_id
        )
        return response['SnapshotId']

    
    def delete_snapshot(self,snapshot_id):
        self.ec2_client.delete_snapshot(
            SnapshotId=snapshot_id
        )


    def describe_snapshot(self,snapshot_id):
        args = dict(
            Filters=[
                {
                    'Name': 'tag:BelongsTo',
                    'Values': [
                        self.asg_name,
                    ]
                },
                {
                    'Name': 'tag:ManagedBy',
                    'Values': [
                        self.__get_name(),
                    ]
                }
            ],
            MaxResults=1,
            SnapshotIds=[snapshot_id]
        )
        snapshots = self.ec2_client.client.describe_snapshots(**args)['Snapshots']
        snapshot = None
        if len(snapshots) > 0:
            snapshot = snapshots[0]
        if snapshot is None: 
            return None 
        return snapshot
    
    def tag_resource(self,resource_id):
        self.ec2_client.create_tags(
            Resources=[
                resource_id,
            ],
            Tags=[
                {
                    'Name': 'tag:BelongsTo',
                    'Values': [
                        self.asg_name,
                    ]
                },
                {
                    'Name': 'tag:ManagedBy',
                    'Values': [
                        self.__get_name(),
                    ]
                }
            ]
        )

def manage_ebs(args):
    asg_ebs_manager = AsgEbsManager(region=args.region,asg_name=args.asg_name)

    # get the current instance az and id
    mds_token = get_imds_token()

    az =  get_instance_metadata(mds_token,"placement/availability-zone")
    instance_id =  get_instance_metadata(mds_token,"instance-id")

    volume_id = asg_ebs_manager.latest_volume(az)
    # if volume_id is not None, there is chance we can get the available volume
    if volume_id is not None:
        # find the available volume
        volume_id = asg_ebs_manager.latest_volume(az,'available')
        # wait for some other volume to be available
        max_retries = args.check_timeout
        while volume_id is None and max_retries > 0:
                sleep(60)
                volume_id = asg_ebs_manager.latest_volume(az,'available')
                max_retries -= 1
        if max_retries == 0:
            volume_id = None
    
    # check the difference in the volume configuration and if there is a diff, 
    # create a new volume from the snapshot of the existing volume and
    # delete the older volume
    if volume_id:
        volume = asg_ebs_manager.describe_volume(volume_id)
        if volume['Size'] != args.size or volume['Iops'] != args.iops or volume['Throughput'] != args.throughput:
            snapshot_id = asg_ebs_manager.create_snapshot(volume_id)
            snapshot = asg_ebs_manager.describe_snapshot(snapshot_id)
            max_retries = args.check_timeout
            while snapshot['State'] != 'completed' and max_retries > 0:
                sleep(60)
                max_retries -= 1
                snapshot = asg_ebs_manager.describe_snapshot(snapshot_id)
            # if snapshot is not getting created, attach the same volume
            if max_retries == 0:
                snapshot_id = None

        # if the snapshot is created, create a new volume
        if snapshot_id:
            old_volume_id = volume_id
            asg_ebs_manager.tag_resource(snapshot_id)
            volume_id = asg_ebs_manager.create_volume(az,args.size,args.iops,args.throughput,snapshot_id=snapshot_id)
            max_retries = args.check_timeout
            volume = asg_ebs_manager.describe_volume(volume_id)
            while volume['State'] != 'available' and max_retries > 0:
                sleep(60)
                max_retries -= 1
                volume = asg_ebs_manager.describe_volume(volume_id)
            if max_retries == 0:
                volume_id = None
            if args.delete:
                asg_ebs_manager.delete_volume(old_volume_id)
                asg_ebs_manager.delete_snapshot(snapshot_id)
            


    # In this case , we cannot do anything other than creating a new volume
    if volume_id is None and args.force: 
        volume_id = asg_ebs_manager.create_volume(az,args.size,args.iops,args.throughput)
        max_retries = args.check_timeout
        volume = asg_ebs_manager.describe_volume(volume_id)
        while volume['State'] != 'available' and max_retries > 0:
            sleep(60)
            max_retries -= 1
            volume = asg_ebs_manager.describe_volume(volume_id)
        if max_retries == 0:
            volume_id = None


    # if we are able to create/ reuse old volume, attach the volume to the instance and tag it
    if volume_id is not None:
        volume_id = asg_ebs_manager.attach_volume(volume_id,instance_id,args.device)
        max_retries = args.check_timeout
        volume = asg_ebs_manager.describe_volume(volume_id)
        while volume['State'] != 'in-use' and max_retries > 0:
            sleep(60)
            max_retries -= 1
            volume = asg_ebs_manager.describe_volume(volume_id)
        if max_retries == 0:
            volume_id = None
        
        # if we are able to attach volume successfully to the instance, tag the volume 
        if volume_id is not None:
            asg_ebs_manager.tag_resource(volume_id)
            return

    # Exiting because we were not able to create or attach volume successfully
    exit(1)


