import asyncio
from sender import SenderNode
from receiver import ReceiverNode
from libp2p.peer.peerinfo import info_from_p2p_addr
from tests.utils import cleanup


async def create_receivers(num_receivers):
    receivers = []
    for i in range(num_receivers):
        receivers.append(await ReceiverNode.create())
    return receivers

async def connect(node1, node2):
    # node1 connects to node2
    addr = node2.get_addrs()[0]
    info = info_from_p2p_addr(addr)
    await node1.connect(info)

async def create_topology(adjacency_map, sender, receivers):
    # Create network

    # Connect senders to receivers
    for target_num in adjacency_map["sender"]:
        await connect(sender.libp2p_node, receivers[target_num].libp2p_node)

    # Connect receivers to other receivers
    for source_num in adjacency_map:
        if source_num != "sender":
            target_nums = adjacency_map[source_num]
            for target_num in target_nums:
                await connect(receivers[source_num].libp2p_node, \
                    receivers[target_num].libp2p_node)

async def main():
    # Create sender
    print("Sender created")
    sender = await SenderNode.create()

    # Create receivers
    print("Receivers created")
    # sender_id = sender.libp2p_node.get_id()
    receivers = await create_receivers(2)

    # Define topology
    topology = {
        "sender": [0],
        0: [1]
    }

    # Create network topology
    await create_topology(topology, sender, receivers)
    print("Topology created")

    # Perform throughput test
    # 1) Start receivers 
    sender_info = info_from_p2p_addr(sender.libp2p_node.get_addrs()[0])
    for receiver in receivers:
        print("Starting receiving")
        asyncio.ensure_future(receiver.start_receiving(sender_info))
    await asyncio.sleep(0.5)

    # 2) Start sending messages and perform throughput test
    print("Performing test")
    await sender.perform_test(2, 5)

    await cleanup()



if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())
    loop.close()