"""Drive several robots at the same time, on top of academy.py.

academy.py's moves block: robot_a.drive() returns once robot_a has stopped, so
robot_b can only start after it. Here, wait=False makes drive(), spin() and
wheels() return at once while a background thread keeps the robot moving, and
wait() blocks until the move is over. Without wait=False they block, as in
academy.py.

    import academy
    import academy_parallel

    academy.connect("http://localhost:8000")
    robot_a, robot_b = academy_parallel.swarm(["000000", "000001"])
    robot_a.drive(speed=100, seconds=2, wait=False)
    robot_b.spin(speed=50, seconds=2, wait=False)
    academy_parallel.wait(robot_a, robot_b)
"""

import atexit
import threading
import time

import requests

import academy

PERIOD = 0.1  # s between resends; a robot stops by itself ~0.5 s after the last one

_bots = []


class Bot(academy.Bot):
    """An academy robot whose drive(), spin() and wheels() can run in the background.

    One daemon thread per robot resends the running move every PERIOD; a new
    move replaces the running one, and stop() cancels it.
    """

    def __init__(self, address):
        super().__init__(address)
        self._lock = threading.RLock()
        self._move = None  # (wheel_velocity body, end time) of the running move
        self._done = threading.Event()
        self._done.set()
        self._thread = None
        _bots.append(self)

    def drive(self, speed=100, seconds=1, wait=True):
        """Drive straight at `speed` mm/s (negative backs up) for `seconds`, then stop."""
        self.wheels(speed, speed, seconds, wait)

    def spin(self, speed=100, seconds=1, wait=True):
        """Turn on the spot, clockwise on the map (negative: anticlockwise), then stop."""
        self.wheels(speed, -speed, seconds, wait)

    def wheels(self, left, right, seconds, wait=True):
        """Each wheel's speed in mm/s for `seconds`, then stop. Replaces a running move."""
        body = {"left_mm_s": left, "right_mm_s": right}
        with self._lock:
            self._move = (body, time.time() + seconds)
            self._done.clear()
            self._put("wheel_velocity", body)
            if self._thread is None:
                self._thread = threading.Thread(target=self._run, daemon=True)
                self._thread.start()
        if wait:
            self.wait()

    def wait(self):
        """Block until the running move is over."""
        while not self._done.wait(PERIOD):
            pass

    def stop(self):
        """Cancel the running move and set both wheels to 0."""
        with self._lock:
            self._move = None
            try:
                super().stop()
            finally:
                self._done.set()

    def goto(self, x, y, threshold=60):
        """As academy's goto(), after cancelling a running move."""
        if self._move is not None:
            self.stop()
        super().goto(x, y, threshold)

    def _run(self):
        while True:
            with self._lock:
                if self._move is not None:
                    body, end = self._move
                    try:
                        if time.time() < end:
                            self._put("wheel_velocity", body)
                        else:
                            self.stop()
                    except requests.RequestException:
                        pass  # try again next period; the robot stops by itself meanwhile
                    if time.time() >= end:
                        self._move = None
                        self._done.set()
            time.sleep(PERIOD)


def bot(label):
    """Your robot, from its label, with non-blocking moves."""
    return Bot(academy.bot(label).address)


def swarm(labels):
    """Your robots, from their labels, with non-blocking moves."""
    return [bot(label) for label in labels]


def wait(*robots):
    """Block until every robot's running move is over."""
    for robot in robots:
        robot.wait()


@atexit.register
def _stop_all():
    # Runs on a normal exit and after Ctrl+C, so no robot is left driving
    for robot in _bots:
        if robot._thread is not None:
            try:
                robot.stop()
            except requests.RequestException:
                pass