import logging
logging.basicConfig(level=logging.DEBUG)
import time

import zmq

from robot_interface.endpoints.actuation_receiver import ActuationReceiver
from robot_interface.endpoints.main_receiver import MainReceiver
from robot_interface.endpoints.video_sender import VideoSender
from robot_interface.state import state


def main_loop(context):
    """
    Run the main loop, handling all incoming requests like pings, negotiation, actuation, etc.

    :param context: The ZeroMQ context to use.
    :type context: zmq.Context
    """
    # When creating sockets, remember to add them to the `sockets` list of the state to ensure they're deinitialized
    main_receiver = MainReceiver(context)
    state.sockets.append(main_receiver)
    actuation_receiver = ActuationReceiver(context)
    state.sockets.append(actuation_receiver)

    video_sender = VideoSender(context)
    state.sockets.append(video_sender)

    video_sender.start_video_rcv()

    # Sockets that can run on the main thread. These sockets' endpoints should not block for long (say 50 ms at most).
    receivers = [main_receiver, actuation_receiver]

    poller = zmq.Poller()
    for receiver in receivers:
        poller.register(receiver.socket, zmq.POLLIN)

    logging.debug("Starting main loop.")

    while True:
        if state.exit_event.is_set(): break
        socks = dict(poller.poll(100))

        for receiver in receivers:
            if receiver.socket not in socks: continue

            start_time = time.time()

            message = receiver.socket.recv_json()
            response = receiver.handle_message(message)
            if receiver.socket.getsockopt(zmq.TYPE) == zmq.REP:
                receiver.socket.send_json(response)

            time_spent_ms = (time.time() - start_time) * 1000
            if time_spent_ms > 50:
                logging.warn("Endpoint \"%s\" took too long (%.2f ms) on the main thread.", receiver.name, time_spent_ms)


def main():
    context = zmq.Context()

    state.initialize()

    try:
        main_loop(context)
    except KeyboardInterrupt:
        logging.info("User interrupted.")
    finally:
        state.deinitialize()
        context.term()


if __name__ == "__main__":
    main()