import unittest

from boto3 import Session
from moto import mock_ec2

from smartaws.ec2 import base
from smartaws.ec2 import internet_gateway as ec2_igw
from smartaws.ec2 import vpc as ec2_vpc


class VpcTestCase(unittest.TestCase):

    credentials = {
        'aws_access_key_id': 'access_key',
        'aws_secret_access_key': 'secret_access_key',
        'region_name': 'us-west-2'
    }

    def setUp(self):
        self.mock = mock_ec2()
        self.mock.start()

        self.session = Session(**self.credentials)
        self.ec2 = self.session.resource('ec2')

        # Create a default VPC for each test
        self.ec2.create_vpc(CidrBlock='172.31.0.0/16')

    def tearDown(self):
        self.mock.stop()

        # Clean the cache for EC2 objects after each test
        base.BaseHandler._cache.clear()

    def test_create_vpc(self):
        ctx = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'tags': {
                'description': 'VPC vpc01 (subnet01a & subnet01b)'
            }
        }

        tags = [
            {
                'Key': 'Name',
                'Value': 'vpc01'
            },
            {
                'Key': 'Description',
                'Value': 'VPC vpc01 (subnet01a & subnet01b)'
            }
        ]

        filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the VPC
        h = ec2_vpc.create_handler(ctx, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]

        self.assertEqual(len(vpcs), 1)
        self.assertEqual(vpc.name, 'vpc01')
        self.assertEqual(vpc.cidr_block, '10.0.10.0/24')
        self.assertCountEqual(vpc.tags, tags)

        # Test default values for VPC attributes
        response = vpc.describe_attribute(Attribute='enableDnsSupport')
        attr = response.get('EnableDnsSupport')
        self.assertEqual(attr.get('Value'), True)

        response = vpc.describe_attribute(Attribute='enableDnsHostnames')
        attr = response.get('EnableDnsHostnames')
        self.assertEqual(attr.get('Value'), False)

    def test_update_vpc(self):
        ctx1 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'tags': {
                'description': 'VPC vpc01 (subnet01a & subnet01b)',
                'stack': 'Test',
                'owner': 'Team A'
            }
        }

        ctx2 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'tags': {
                'description': 'VPC vpc01 (subnet01a & subnet01b)',
                'stack': 'Production'
            }
        }

        tags1 = [
            {
                'Key': 'Name',
                'Value': 'vpc01'
            },
            {
                'Key': 'Description',
                'Value': 'VPC vpc01 (subnet01a & subnet01b)'
            },
            {
                'Key': 'Stack',
                'Value': 'Test'
            },
            {
                'Key': 'Owner',
                'Value': 'Team A'
            }
        ]

        tags2 = [
            {
                'Key': 'Name',
                'Value': 'vpc01'
            },
            {
                'Key': 'Description',
                'Value': 'VPC vpc01 (subnet01a & subnet01b)'
            },
            {
                'Key': 'Stack',
                'Value': 'Production'
            }
        ]

        filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the VPC
        h = ec2_vpc.create_handler(ctx1, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]
        vpc_id = vpc.id

        self.assertEqual(len(vpcs), 1)
        self.assertCountEqual(vpc.tags, tags1)

        # We clear the resource cache to simulate a new
        # program execution with the 'update' option
        base.BaseHandler._cache.clear()

        # Update the VPC
        h = ec2_vpc.create_handler(ctx2, self.credentials)
        h.update_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]

        self.assertEqual(len(vpcs), 1)
        self.assertEqual(vpc.id, vpc_id)
        self.assertCountEqual(vpc.tags, tags2)

    def test_update_enable_dns_support(self):
        ctx1 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24'
        }

        ctx2 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'enable_dns_support': False
        }

        filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the VPC
        h = ec2_vpc.create_handler(ctx1, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]

        # Test default values for 'enableDnsSupport' attribute
        response = vpc.describe_attribute(Attribute='enableDnsSupport')
        attr = response.get('EnableDnsSupport')
        self.assertEqual(attr.get('Value'), True)

        # We clear the resource cache to simulate a new
        # program execution with the 'update' option
        base.BaseHandler._cache.clear()

        # Update the VPC
        h = ec2_vpc.create_handler(ctx2, self.credentials)
        h.update_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]

        response = vpc.describe_attribute(Attribute='enableDnsSupport')
        attr = response.get('EnableDnsSupport')
        self.assertEqual(attr.get('Value'), False)

    def test_update_enable_dns_hostnames(self):
        ctx1 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24'
        }

        ctx2 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'enable_dns_hostnames': True
        }

        filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the VPC
        h = ec2_vpc.create_handler(ctx1, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]

        # Test default values for 'enableDnsHostnames' attribute
        response = vpc.describe_attribute(Attribute='enableDnsHostnames')
        attr = response.get('EnableDnsHostnames')
        self.assertEqual(attr.get('Value'), False)

        # We clear the resource cache to simulate a new
        # program execution with the 'update' option
        base.BaseHandler._cache.clear()

        # Update the VPC
        h = ec2_vpc.create_handler(ctx2, self.credentials)
        h.update_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))
        vpc = vpcs[0]

        response = vpc.describe_attribute(Attribute='enableDnsHostnames')
        attr = response.get('EnableDnsHostnames')
        self.assertEqual(attr.get('Value'), True)

    def test_delete_vpc(self):
        ctx = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24'
        }

        filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the VPC
        h = ec2_vpc.create_handler(ctx, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))

        self.assertEqual(len(vpcs), 1)

        # We clear the resource cache to simulate a new
        # program execution with the 'delete' option
        base.BaseHandler._cache.clear()

        # Delete the VPC
        h.delete_resource()

        vpcs = list(ec2.vpcs.filter(Filters=filters))

        self.assertEqual(len(vpcs), 0)

    def test_delete_vpc_with_attached_internet_gateway(self):
        igw_ctx = {
            'name': 'igw01'
        }

        vpc_ctx = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'internet_gateway': 'igw01'
        }

        igw_filters = [{'Name': 'tag:Name', 'Values': ['igw01']}]
        vpc_filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the internet gateway
        h = ec2_igw.create_handler(igw_ctx, self.credentials)
        h.create_resource()

        gateways = list(ec2.internet_gateways.filter(Filters=igw_filters))
        igw = gateways[0]

        self.assertCountEqual(igw.attachments, [])

        # Create the VPC
        h = ec2_vpc.create_handler(vpc_ctx, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc_filters))
        vpc = vpcs[0]

        # Test that the internet gateway has been attached
        igw.reload()
        attachments = [{'VpcId': vpc.id, 'State': 'available'}]
        self.assertCountEqual(igw.attachments, attachments)

        # We clear the resource cache to simulate a new
        # program execution with the 'delete' option
        base.BaseHandler._cache.clear()

        # Delete the VPC
        h = ec2_vpc.create_handler(vpc_ctx, self.credentials)
        h.delete_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc_filters))

        self.assertEqual(len(vpcs), 0)

    def test_attach_internet_gateway(self):
        igw_ctx = {
            'name': 'igw01'
        }

        vpc_ctx = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'internet_gateway': 'igw01'
        }

        igw_filters = [{'Name': 'tag:Name', 'Values': ['igw01']}]
        vpc_filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the internet gateway
        h = ec2_igw.create_handler(igw_ctx, self.credentials)
        h.create_resource()

        gateways = list(ec2.internet_gateways.filter(Filters=igw_filters))
        igw = gateways[0]

        self.assertCountEqual(igw.attachments, [])

        # Create the VPC
        h = ec2_vpc.create_handler(vpc_ctx, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc_filters))
        vpc = vpcs[0]

        # Test that the internet gateway has been attached
        igw.reload()
        attachments = [{'VpcId': vpc.id, 'State': 'available'}]
        self.assertCountEqual(igw.attachments, attachments)

    def test_internet_gateway_already_attached(self):
        igw01_ctx = {
            'name': 'igw01'
        }

        vpc01_ctx = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'internet_gateway': 'igw01'
        }

        vpc02_ctx = {
            'name': 'vpc02',
            'cidr_block': '10.0.20.0/24',
            'internet_gateway': 'igw01'
        }

        igw01_filters = [{'Name': 'tag:Name', 'Values': ['igw01']}]
        vpc01_filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]
        vpc02_filters = [{'Name': 'tag:Name', 'Values': ['vpc02']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the internet gateway
        h = ec2_igw.create_handler(igw01_ctx, self.credentials)
        h.create_resource()

        gateways = list(ec2.internet_gateways.filter(Filters=igw01_filters))
        igw01 = gateways[0]

        self.assertCountEqual(igw01.attachments, [])

        # Create the first VPC (vpc01)
        h = ec2_vpc.create_handler(vpc01_ctx, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc01_filters))
        vpc01 = vpcs[0]

        attachments = [{'VpcId': vpc01.id, 'State': 'available'}]

        # Test that the internet gateway has been attached
        igw01.reload()
        self.assertCountEqual(igw01.attachments, attachments)

        # Create the second VPC (vpc02)
        h = ec2_vpc.create_handler(vpc02_ctx, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc02_filters))
        vpc02 = vpcs[0]

        self.assertNotEqual(vpc01.id, vpc02.id)

        # Test that the internet gateway is still attached to the first VPC
        igw01.reload()
        self.assertCountEqual(igw01.attachments, attachments)

    def test_another_internet_gateway_already_attached_1(self):
        # Case 1: an internet gateway is attached to the VPC but not
        # defined in the configuration. We leave everything unchanged.
        igw01_ctx = {
            'name': 'igw01'
        }

        vpc01_ctx1 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'internet_gateway': 'igw01'
        }

        vpc01_ctx2 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24'
        }

        igw01_filters = [{'Name': 'tag:Name', 'Values': ['igw01']}]
        vpc01_filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the internet gateway
        h = ec2_igw.create_handler(igw01_ctx, self.credentials)
        h.create_resource()

        gateways = list(ec2.internet_gateways.filter(Filters=igw01_filters))
        igw01 = gateways[0]

        self.assertCountEqual(igw01.attachments, [])

        # Create the VPC
        h = ec2_vpc.create_handler(vpc01_ctx1, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc01_filters))
        vpc01 = vpcs[0]

        attachments = [{'VpcId': vpc01.id, 'State': 'available'}]

        # Test that the internet gateway has been attached
        igw01.reload()
        self.assertCountEqual(igw01.attachments, attachments)

        # We clear the resource cache to simulate a new
        # program execution with the 'update' option
        base.BaseHandler._cache.clear()

        # Update the VPC with the new context
        h = ec2_vpc.create_handler(vpc01_ctx2, self.credentials)
        h.update_resource()

        # Test that the internet gateway is still attached to the VPC
        igw01.reload()
        self.assertCountEqual(igw01.attachments, attachments)

    def test_another_internet_gateway_already_attached_2(self):
        # Case 2: an internet gateway is attached to the VPC and there
        # is another internet gateway defined in the configuration. We
        # also leave everything unchanged.
        igw01_ctx = {
            'name': 'igw01'
        }

        igw02_ctx = {
            'name': 'igw02'
        }

        vpc01_ctx1 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'internet_gateway': 'igw01'
        }

        vpc01_ctx2 = {
            'name': 'vpc01',
            'cidr_block': '10.0.10.0/24',
            'internet_gateway': 'igw02'
        }

        igw01_filters = [{'Name': 'tag:Name', 'Values': ['igw01']}]
        igw02_filters = [{'Name': 'tag:Name', 'Values': ['igw02']}]
        vpc01_filters = [{'Name': 'tag:Name', 'Values': ['vpc01']}]

        def _add_wrapper(base_classes, **kwargs):
            base_classes.insert(0, ec2_vpc.VpcWrapper)

        event = 'creating-resource-class.ec2.Vpc'
        session = Session(**self.credentials)
        session.events.register(event, _add_wrapper)
        ec2 = session.resource('ec2')

        # Create the first internet gateway
        h = ec2_igw.create_handler(igw01_ctx, self.credentials)
        h.create_resource()

        gateways = list(ec2.internet_gateways.filter(Filters=igw01_filters))
        igw01 = gateways[0]

        self.assertCountEqual(igw01.attachments, [])

        # Create the VPC
        h = ec2_vpc.create_handler(vpc01_ctx1, self.credentials)
        h.create_resource()

        vpcs = list(ec2.vpcs.filter(Filters=vpc01_filters))
        vpc01 = vpcs[0]

        attachments = [{'VpcId': vpc01.id, 'State': 'available'}]

        # Test that the internet gateway has been attached
        igw01.reload()
        self.assertCountEqual(igw01.attachments, attachments)

        # Create the second internet gateway
        h = ec2_igw.create_handler(igw02_ctx, self.credentials)
        h.create_resource()

        gateways = list(ec2.internet_gateways.filter(Filters=igw02_filters))
        igw02 = gateways[0]

        self.assertCountEqual(igw02.attachments, [])

        # We clear the resource cache to simulate a new
        # program execution with the 'update' option
        base.BaseHandler._cache.clear()

        # Update the VPC with the new context
        h = ec2_vpc.create_handler(vpc01_ctx2, self.credentials)
        h.update_resource()

        # Test that the first internet gateway is still attached to the VPC
        igw01.reload()
        self.assertCountEqual(igw01.attachments, attachments)

        # Test that the second internet gateway is not attached
        igw02.reload()
        self.assertCountEqual(igw02.attachments, [])
