import atexit
from threading import Thread
from rpcserver import RPCServer


class RPCServerThread(Thread):

    def __init__(self, *args, **kwargs):
        super().__init__(daemon=True) # atexit seems to only work for deamon threads
        self.server = RPCServer(*args, **kwargs)
        self.thread_shutdown = None


    def run(self):
        atexit.register(self.wait_for_stop)
        self.server.serve_forever()


    # BaseServer.shutdown docs say:
    # This must be called while serve_forever() is running in another thread,
    # or it will deadlock.

    def wait_for_stop(self):
        if self.thread_shutdown is None:
            self.stop()
        self.thread_shutdown.join()
        print("RPC server stopped")

    def stop(self):
        self.thread_shutdown = t = Thread(target=self.shutdown)
        t.start()

    def shutdown(self):
        self.server.shutdown()
        self.server.server_close()





if __name__ == "__main__":
    from time import sleep

    def test():
        print("test")
        return "test"

    rst = RPCServerThread("localhost", 8000)
    rst.server.register_function(test)
    rst.server.register_function(rst.stop)
    rst.start()

    while rst.is_alive():
        sleep(1)