from typing import Callable, Optional
from mlagents_envs.communicator_objects.unity_output_pb2 import UnityOutputProto
from mlagents_envs.communicator_objects.unity_input_pb2 import UnityInputProto


# Function to call while waiting for a connection timeout.
# This should raise an exception if it needs to break from waiting for the timeout.
PollCallback = Callable[[], None]


class Communicator:
    def __init__(self, worker_id=0, base_port=5005):
        """
        Python side of the communication. Must be used in pair with the right Unity Communicator equivalent.

        :int worker_id: Offset from base_port. Used for training multiple environments simultaneously.
        :int base_port: Baseline port number to connect to Unity environment over. worker_id increments over this.
        """

    def initialize(
        self, inputs: UnityInputProto, poll_callback: Optional[PollCallback] = None
    ) -> UnityOutputProto:
        """
        Used to exchange initialization parameters between Python and the Environment
        :param inputs: The initialization input that will be sent to the environment.
        :param poll_callback: Optional callback to be used while polling the connection.
        :return: UnityOutput: The initialization output sent by Unity
        """

    def exchange(
        self, inputs: UnityInputProto, poll_callback: Optional[PollCallback] = None
    ) -> Optional[UnityOutputProto]:
        """
        Used to send an input and receive an output from the Environment
        :param inputs: The UnityInput that needs to be sent the Environment
        :param poll_callback: Optional callback to be used while polling the connection.
        :return: The UnityOutputs generated by the Environment
        """

    def close(self):
        """
        Sends a shutdown signal to the unity environment, and closes the connection.
        """