import time

class ShortBuffer(object):
    def __init__(self, state):
        self.state = state
        self.buffer = {}

    def embargo(self, message):
        if message.message_hash not in self.buffer.keys():
            self.buffer[message.message_hash] = {
                'received': time.time(),
                'message': message,
                'low_bounce_count': message.bounces,
                'closest_peers': [message.peer],
                'reporting_peers': [message.peer]
            }
        else:
            embargoed_message = self.buffer[message.message_hash]
            if message.bounces < embargoed_message['low_bounce_count']:
                embargoed_message['low_bounce_count'] = message.bounces
                embargoed_message['closest_peers'] = [message.peer]
            elif message.bounces == embargoed_message['low_bounce_count']:
                embargoed_message['closest_peers'].append(message.peer)
            else:
                # not interested in the message because the bounce count
                # is higher than what we've already got and we just want the
                # list of peers closest to the originator
                pass
            embargoed_message['reporting_peers'].append(message.peer)

    def flush(self):
        current_time = time.time()
        messages = []
        for message_with_stats in self.buffer.values():
            if (message_with_stats['received'] <
                    (current_time - int(self.state.get_knob('short_buffer_expiration_seconds')))):
                messages.append(message_with_stats)
                del self.buffer[message_with_stats['message'].message_hash]
        return sorted(messages, key=lambda m: m['message'].timestamp)

    def has(self, message_hash):
        return self.buffer.get(message_hash)

    def drop(self, message_hash):
        del self.buffer[message_hash]