import Pyro4
import logging

Pyro4.config.REQUIRE_EXPOSE = True
log = logging.getLogger(__name__)


class CallbackServer(object):

    def __init__(self, listen_port, remote_port):
        self.daemon = Pyro4.Daemon(host='localhost', port=listen_port,
                                   nathost='localhost', natport=remote_port,
                                   unixsocket=None)
        self.daemon.register(self)

    def __enter__(self):
        def serve():
            self.daemon.requestLoop()
        from threading import Thread
        self.thread = Thread(target=serve)
        log.debug('Starting callback server')
        self.thread.start()
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        log.debug('Shutting down callback server')
        self.daemon.shutdown()
        self.thread.join()

    @Pyro4.expose
    def handle_log(self, pickled_record):
        import pickle
        record = pickle.loads(pickled_record)
        record.extra = getattr(record, 'extra', {})
        record.extra['source'] = 'remote'
        log.handle(record)