import asyncio
import time
import unittest

import preseem.grpc_nm_model
from preseem import PreseemNetworkMetadataModel, NetworkMetadataEntity, NetworkMetadataIndex, NetworkMetadataLowercaseIndex, NetworkMetadataReference, Reference
from preseem_grpc_model import network_metadata_pb2

def _make_entity(net):
    """Helper to make a NetworkMetadataEntity in a standard way."""
    return NetworkMetadataEntity(net, time.time(), {
        'session': Reference(net),
        'snmp_session': Reference(net),
    })

class FakeMetadataClient(object):
    """Fake nm client"""
    def __init__(self):
        self.nms = []  # entities to set in the fake server
        self.refs = []  # references to set in the fake server

    def close_subscriptions(self):
        pass

    async def populate_ents(self):
        """Fake sending entities via a subscription."""
        syn = True
        async def send_msg(ent):
            """Send a message.  If ent is "None", send the LOADED action."""
            nonlocal syn
            await self.send_ent_entity_action(ent, syn=syn)
            syn = False

        for ent in self.nms:
            await send_msg(ent)
        await send_msg(None)
        self.done_fut.set_result(None)

    async def send_ent_entity_action(self, ent, delete=False, syn=False, fut=None):
        """Send a message to the subscription.  If ent is None, send the LOADED action."""
        ts = network_metadata_pb2.google_dot_protobuf_dot_timestamp__pb2.Timestamp()             
        ts.FromNanoseconds(int(time.time() * 1000000000))
        action = network_metadata_pb2.NetworkMetadataEntityAction.SET
        set_ent = None
        if ent:
            set_ent = network_metadata_pb2.NetworkMetadataEntity(address_space='default', network_prefix=ent.network_prefix, start=ts)
            preseem.grpc_nm_model._nm_set_attrs(set_ent.attributes, ent.attributes)
            if delete:
                action = network_metadata_pb2.NetworkMetadataEntityAction.DELETE
        else:
            action = network_metadata_pb2.NetworkMetadataEntityAction.LOADED
        msg = network_metadata_pb2.NetworkMetadataEntityAction(action=action, entity=set_ent, timestamp=ts)
        await self.ent_cbk(msg, syn)
        if fut:
            fut.set_result(None)

    async def send_ref_entity_action(self, ref, delete=False, syn=False, fut=None):
        """Send a message to the subscription.  If ref is None, send the LOADED action."""
        ts = network_metadata_pb2.google_dot_protobuf_dot_timestamp__pb2.Timestamp()             
        ts.FromNanoseconds(int(time.time() * 1000000000))
        action = network_metadata_pb2.NetworkMetadataReferenceEntityAction.SET
        set_ref = None
        if ref:
            set_ref = network_metadata_pb2.NetworkMetadataReferenceEntity(type=ref.type, value=ref.value)
            preseem.grpc_nm_model._nm_set_attrs(set_ref.attributes, ref.attributes)
            if delete:
                action = network_metadata_pb2.NetworkMetadataReferenceEntityAction.DELETE
        else:
            action = network_metadata_pb2.NetworkMetadataReferenceEntityAction.LOADED
        msg = network_metadata_pb2.NetworkMetadataReferenceEntityAction(action=action, entity=set_ref, timestamp=ts)
        await self.ref_cbk(msg, syn)
        if fut:
            fut.set_result(None)

    async def del_nm_async(self, ent):
        """HAve an entity asynchronously deleted from the subscription."""
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        loop.create_task(self.send_ent_entity_action(ent, delete=True, fut=fut))
        await fut

    async def set_nm_async(self, ent):
        """HAve an entity asynchronously deleted from the subscription."""
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        loop.create_task(self.send_ent_entity_action(ent, fut=fut))
        await fut

    async def del_ref_async(self, ref):
        """Have a reference asynchronously deleted from the subscription."""
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        loop.create_task(self.send_ref_entity_action(ref, delete=True, fut=fut))
        await fut

    async def set_ref_async(self, ref):
        """Have a reference asynchronously set into the subscription."""
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        loop.create_task(self.send_ref_entity_action(ref, fut=fut))
        await fut

    async def populate_refs(self):
        """Fake sending refs via a subscription."""
        syn = True
        async def send_msg(ref):
            """Send a message.  If ref is "None", send the LOADED action."""
            nonlocal syn
            await self.send_ref_entity_action(ref, syn=syn)
            syn = False

        for ref in self.refs:
            await send_msg(ref)
        await send_msg(None)
        self.done_fut.set_result(None)

    async def nm_ent_subscribe(self, cbk, company_id=None):
        self.ent_cbk = cbk
        loop = asyncio.get_running_loop()
        self.done_fut = loop.create_future()
        loop.create_task(self.populate_ents())
        return self.done_fut

    async def nm_ref_subscribe(self, cbk, company_id=None):
        self.ref_cbk = cbk
        loop = asyncio.get_running_loop()
        self.done_fut = loop.create_future()
        loop.create_task(self.populate_refs())
        return self.done_fut

class TestNetworkMetadataModel(unittest.TestCase):
    def setUp(self):
        self.loop = asyncio.new_event_loop() # needed to initialize asyncio
        asyncio.set_event_loop(self.loop)
        self.client = FakeMetadataClient()
        self.model = PreseemNetworkMetadataModel(self.client, refs_by_type=True)

    def tearDown(self):
        self.model.close()
        self.loop.run_until_complete(self.loop.shutdown_asyncgens())
        self.loop.close()

    def update(self):
        return self.loop.run_until_complete(self.model.update())

    def _wait_for_queue(self):
        """Wait for the flex model queue to be emptied."""
        self.loop.run_until_complete(self.model._flex_q.join())

    def _wait(self, sec):
        self.loop.run_until_complete(asyncio.sleep(sec))

    def check_req(self, pb, flexes):
        """Verify that a list of flexes is the same as a flex request."""
        req = metrics_pb2.FlexPushRequest()
        req.flexes.extend([self.model._flex_pb(x) for x in flexes])
        self.assertEqual(pb, req)

    def del_nm(self, ent):
        self.loop.run_until_complete(self.client.del_nm_async(ent))

    def set_nm(self, ent):
        self.loop.run_until_complete(self.client.set_nm_async(ent))

    def del_ref(self, ref):
        self.loop.run_until_complete(self.client.del_ref_async(ref))

    def set_ref(self, ref):
        self.loop.run_until_complete(self.client.set_ref_async(ref))

    def test_lowercase_index(self):
        """Test lowercase reference index"""
        # Load fake refs
        self.model.close()
        indexes = [NetworkMetadataLowercaseIndex('service', 'username'),]
        self.model = PreseemNetworkMetadataModel(self.client, refs_by_type=True, indexes=indexes)
        self.client.refs = [
            NetworkMetadataReference('service', 'service1', {}),
            NetworkMetadataReference('service', 'service2', {'username': 'User1'})
        ]
        self.update()

        # Lookup from initial load
        refs = self.model.get_ref_index('service', 'username', 'user1') or {}
        self.assertEqual(len(refs), 1)
        self.assertEqual(refs.get('service2', None), self.client.refs[1])

        # Load in a new ref to update the index
        new_ref = NetworkMetadataReference('service', 'service8', {'username': 'User8'})
        self.set_ref(new_ref)
        refs = self.model.get_ref_index('service', 'username', 'user8') or {}
        self.assertEqual(len(refs), 1)
        self.assertEqual(refs.get('service8', None), new_ref)

        # Delete a ref and make sure it's removed from the index
        self.del_ref(new_ref)
        refs = self.model.get_ref_index('service', 'username', 'user8') or {}
        self.assertEqual(len(refs), 0)

    def test_nm_trie(self):
        # Load fake entities
        self.model.close()
        """Make sure network metadata entity changes update the trie index."""
        self.model = PreseemNetworkMetadataModel(self.client, refs_by_type=True)
        self.client.nms = [
            _make_entity('192.0.2.1'),
            _make_entity('192.0.2.54'),
            _make_entity('192.0.2.52/30'),
        ]
        self.update()

        # Check results of initial load
        self.assertEqual(len(self.model.nms), 3)
        self.assertEqual(len(self.model.nms), len(self.model.pt4))
        key = self.model.pt4.get_key('192.0.2.54')
        self.assertEqual(key, '192.0.2.54/32')
        par = self.model.pt4.parent(key)
        self.assertEqual(par, '192.0.2.52/30')

        # Add another entity
        self.set_nm(_make_entity('192.0.2.2'))
        self.assertEqual(len(self.model.nms), 4)
        self.assertEqual(len(self.model.nms), len(self.model.pt4))

        # Delete an entity
        self.del_nm(_make_entity('192.0.2.54'))
        self.assertEqual(len(self.model.nms), 3)
        key = self.model.pt4.get_key('192.0.2.54')
        self.assertEqual(key, '192.0.2.52/30')
