from collections import defaultdict

from .base import Router

class EpidemicRouter(Router):
    def __init__(self, node, **kwargs):
        super().__init__(node, **kwargs)

        # maps a packet to a set of nodes that the packet has been sent to
        self.sent = defaultdict(set)

    def __call__(self, packet):
        # get list of currently encountered nodes that do not have the packet
        targets = [
            met
            for met in self.node.neighbours
            if met not in self.sent[packet]
        ]

        # return list of nodes to send packet to
        # (targets, remove from buffer, reason)
        if targets:
            return targets, 'epidemic', False
        else:
            return None, None, False

    def on_send_success(self, target, packet):
        self.sent[packet].add(target)