import boto3


class EC2Wrapper:
    def get_vpc_subnets(self, region_name):
        client = boto3.client(
            'ec2',
            region_name=region_name
        )

        vpc_subnets = {}
        vpcs = client.describe_vpcs()

        for vpc in vpcs['Vpcs']:
            vpc_id = vpc['VpcId']
            subnets = client.describe_subnets(
                Filters=[
                    {
                        'Name': 'vpc-id',
                        'Values': [
                            vpc_id,
                        ]
                    },
                ]
            )
            subnet_ids = [subnet['SubnetId'] for subnet in subnets['Subnets']]
            # Only include VPCs with subnets
            if subnet_ids:
                vpc_subnets[vpc_id] = subnet_ids

        return vpc_subnets

    def list_all_regions(self):
        # Creates an EC2 client for users default region
        ec2_client = boto3.client('ec2')

        # Get the list of available regions
        response = ec2_client.describe_regions()

        # Return just the regions names
        return [region['RegionName'] for region in response['Regions']]

    def get_current_region(self):
        # Creates an EC2 client for users default region
        ec2_client = boto3.client('ec2')

        # Return the default region
        return ec2_client.meta.region_name

    def get_all_key_pairs_in_region(self, region_name):
        # Creates an EC2 client for users default region where key pairs should be located
        ec2_client = boto3.client('ec2', region_name=region_name)

        # Get the list of key pairs
        response = ec2_client.describe_key_pairs()
        key_pairs = response['KeyPairs']

        # Return just the key pair names
        return [kp['KeyName'] for kp in key_pairs]

    def create_key_pair(self, key_name, region_name):
        # Creates an EC2 client for the region where new key pair should be created
        ec2_client = boto3.client('ec2', region_name=region_name)

        # Generates new key pair
        response = ec2_client.create_key_pair(KeyName=key_name)

        # Returns the private key from the new key pair
        return response['KeyMaterial']

    def get_availability_zones(self, region_name):
        # Creates an EC2 client for the region of which the availability zones should be returned
        ec2_client = boto3.client('ec2', region_name=region_name)

        # Get availability zones
        response = ec2_client.describe_availability_zones()

        # Return just the availability zone names for zones that are available
        return [zone['ZoneName'] for zone in response['AvailabilityZones'] if zone['State'] == 'available']
