import cosocket

_channels = {}

class until_message_sent(object):
    def __init__(self, channel_name, message):
        self.channel_name = channel_name
        self.message = message

    def execute(self, dispatcher):
        if self.channel_name in _channels:
            receivers = _channels[self.channel_name].values()
            del _channels[self.channel_name]
            for receiver in receivers:
                receiver.continue_from_yield(self.message)
        dispatcher.continue_from_yield()

class _until_message_received(object):
    def __init__(self, channel_name, timeout = cosocket.DEFAULT_TIMEOUT):
        self.channel_name = channel_name
        self._timeout = timeout
        self._fd = None

    def execute(self, dispatcher):
        if self.channel_name not in _channels:
            _channels[self.channel_name] = {}
        self._fd = dispatcher.socket.fileno()
        _channels[self.channel_name][self._fd] = dispatcher
        dispatcher.set_timeout(self._timeout)

    def finalize(self):
        if (self.channel_name in _channels and
            self._fd in _channels[self.channel_name]):
            del _channels[self.channel_name][self._fd]
            if not _channels[self.channel_name]:
                del _channels[self.channel_name]

def until_message_received(channel_name):
    instruction = _until_message_received(channel_name)
    try:
        message = yield instruction
        yield cosocket.return_value(message)
    finally:
        instruction.finalize()