import netaddr
from django.test import TestCase

from ipam.choices import IPAddressStatusChoices, PrefixStatusChoices
from ipam.models import VRF, IPAddress, Prefix


class OrderingTestBase(TestCase):
    vrfs = None

    @classmethod
    def setUpTestData(cls):
        """
        Setup the VRFs for the class as a whole
        """
        vrfs = (
            VRF(name='VRF 1'),
            VRF(name='VRF 2'),
            VRF(name='VRF 3'),
        )
        VRF.objects.bulk_create(vrfs)

    def _compare(self, queryset, objectset):
        """
        Perform the comparison of the queryset object and the object used to instantiate the queryset.
        """
        for i, obj in enumerate(queryset):
            self.assertEqual(obj, objectset[i])

    def _compare_ne(self, queryset, objectset):
        """
        Perform the comparison of the queryset object and the object used to instantiate the queryset.
        """
        for i, obj in enumerate(queryset):
            self.assertNotEqual(obj, objectset[i])


class PrefixOrderingTestCase(OrderingTestBase):

    def test_prefix_vrf_ordering(self):
        """
        This is a very basic test, which tests both prefixes without VRFs and prefixes with VRFs
        """
        vrf1, vrf2 = VRF.objects.all()[:2]
        prefixes = (
            Prefix(status=PrefixStatusChoices.STATUS_CONTAINER, vrf=None, prefix=netaddr.IPNetwork('192.168.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.0.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.2.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.3.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.4.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.5.0/24')),

            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.0.0/8')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.0.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.2.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.3.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.4.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.2.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.3.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.4.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.2.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.2.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.2.2.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.2.3.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.2.4.0/24')),

            Prefix(status=PrefixStatusChoices.STATUS_CONTAINER, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.0.0/12')),
            Prefix(status=PrefixStatusChoices.STATUS_CONTAINER, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.0.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.2.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.3.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.16.4.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_CONTAINER, vrf=vrf2, prefix=netaddr.IPNetwork('172.17.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.17.0.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.17.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.17.2.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.17.3.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf2, prefix=netaddr.IPNetwork('172.17.4.0/24')),
        )

        Prefix.objects.bulk_create(prefixes)

        # Test
        self._compare(Prefix.objects.all(), prefixes)

    def test_prefix_complex_ordering(self):
        """
        This function tests a complex ordering of interwoven prefixes and vrfs.  This is the current expected ordering
         of VRFs. This includes the testing of the Container status.

        The proper ordering, to get proper containerization should be:
            None:10.0.0.0/8
            None:10.0.0.0/16
            VRF A:10.0.0.0/24
            VRF A:10.0.1.0/24
            VRF A:10.0.1.0/25
            None:10.1.0.0/16
            VRF A:10.1.0.0/24
            VRF A:10.1.1.0/24
            None: 192.168.0.0/16
        """
        vrf1 = VRF.objects.first()
        prefixes = [
            Prefix(status=PrefixStatusChoices.STATUS_CONTAINER, vrf=None, prefix=netaddr.IPNetwork('10.0.0.0/8')),
            Prefix(status=PrefixStatusChoices.STATUS_CONTAINER, vrf=None, prefix=netaddr.IPNetwork('10.0.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('10.1.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=None, prefix=netaddr.IPNetwork('192.168.0.0/16')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.0.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.1.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.0.1.0/25')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.0.0/24')),
            Prefix(status=PrefixStatusChoices.STATUS_ACTIVE, vrf=vrf1, prefix=netaddr.IPNetwork('10.1.1.0/24')),
        ]
        Prefix.objects.bulk_create(prefixes)

        # Test
        self._compare(Prefix.objects.all(), prefixes)


class IPAddressOrderingTestCase(OrderingTestBase):
    def test_address_vrf_ordering(self):
        """
        This function tests ordering with the inclusion of vrfs
        """
        vrf1, vrf2 = VRF.objects.all()[:2]
        addresses = (
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.0.0.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.0.1.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.0.2.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.0.3.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.0.4.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.1.0.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.1.1.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.1.2.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.1.3.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.1.4.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.2.0.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.2.1.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.2.2.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.2.3.1/24')),
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf1, address=netaddr.IPNetwork('10.2.4.1/24')),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.16.0.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.16.1.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.16.2.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.16.3.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.16.4.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.17.0.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.17.1.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.17.2.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.17.3.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=vrf2, address=netaddr.IPNetwork('172.17.4.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=None, address=netaddr.IPNetwork('192.168.0.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=None, address=netaddr.IPNetwork('192.168.1.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=None, address=netaddr.IPNetwork('192.168.2.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=None, address=netaddr.IPNetwork('192.168.3.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=None, address=netaddr.IPNetwork('192.168.4.1/24')
            ),
            IPAddress(
                status=IPAddressStatusChoices.STATUS_ACTIVE, vrf=None, address=netaddr.IPNetwork('192.168.5.1/24')
            ),
        )
        IPAddress.objects.bulk_create(addresses)

        # Test
        self._compare(IPAddress.objects.all(), addresses)

    def test_duplicate_address_ordering(self):
        """
        Host addresses are not unique, so tied addresses must be ordered by primary key to yield a stable, total
        ordering. Without a tiebreaker the database may return tied rows in a different order from one query to the
        next, duplicating or omitting objects across paginated requests.
        """
        # Create several duplicates of each address, interleaved so that primary key order does not follow
        # address order.
        addresses = [
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, address=netaddr.IPNetwork(f'10.0.{i}.1/24'))
            for _ in range(4)
            for i in range(100)
        ]
        IPAddress.objects.bulk_create(addresses)

        pks = list(IPAddress.objects.values_list('pk', flat=True))
        expected = [
            ip.pk for ip in sorted(IPAddress.objects.all(), key=lambda ip: (ip.address.ip, ip.pk))
        ]

        self.assertEqual(pks, expected)

    def test_duplicate_address_pagination(self):
        """
        Paginating over duplicate addresses must not return the same object on two pages, nor omit any object.
        """
        addresses = [
            IPAddress(status=IPAddressStatusChoices.STATUS_ACTIVE, address=netaddr.IPNetwork(f'10.0.{i}.1/24'))
            for _ in range(4)
            for i in range(100)
        ]
        IPAddress.objects.bulk_create(addresses)

        queryset = IPAddress.objects.values_list('pk', flat=True)
        page_size = 37
        paginated = []
        for offset in range(0, len(addresses), page_size):
            paginated.extend(queryset[offset:offset + page_size])

        self.assertEqual(len(paginated), len(addresses))
        self.assertEqual(len(set(paginated)), len(addresses))
        self.assertEqual(paginated, list(queryset))
