import asyncio
import math
from centrifuge import Client, SubscriptionEventHandler, PublicationContext

TIMESTEP_MS = 1000
ROBOT_COUNT = 40
SPAWN_RADIUS = 0.9

class WorldControllerBase:
    def __init__(self):
        client = Client('ws://localhost:9120/connection/websocket')

        class EventHandler(SubscriptionEventHandler):
            async def on_publication(_, ctx: PublicationContext) -> None:
                print(f"tick {ctx.pub.data}")
                self._last_tick_received = ctx.pub.data
                self._process_tick(ctx.pub.data)

        sync = client.new_subscription('sync', EventHandler())
        asyncio.ensure_future(sync.subscribe())

        self.client = client
        self._dt_ms = TIMESTEP_MS
        self._last_tick_received = -1
        self._user_task = None

        fut = asyncio.ensure_future(client.connect())
        fut.add_done_callback(lambda _: asyncio.ensure_future(self.on_start()))

    async def on_start(self):
        # implemented by the user
        pass

    async def on_tick(self, _tick):
        # implemented by the user
        pass

    async def rpc(self, method, args):
        result = await self.client.rpc(method, args)
        return result.data

    def _process_tick(self, tick):
        if self._user_task and self._last_tick_received > 0:
            return # busy

        # send sync
        next_tick = tick + self._dt_ms * 1_000
        sync = self.client.get_subscription('sync')
        asyncio.ensure_future(sync.publish(next_tick))
        # print(f"send sync {next_tick}")

        def on_task_done(_):
            self._user_task = None
            if self._last_tick_received >= next_tick:
                self._process_tick(next_tick)

        self._user_task = asyncio.ensure_future(self.on_tick(tick))
        self._user_task.add_done_callback(on_task_done)

class WorldController(WorldControllerBase):
    async def on_start(self):
        self._robots = []
        await self.rpc('session.restart', None)

        commands = []
        for i in range(ROBOT_COUNT):
            robot = VehicleController(i)
            self._robots.append(robot)
            commands.append(robot.spawn(self))
        commands.append(self.rpc('session.run', None))
        await asyncio.gather(*commands)

    async def on_tick(self, _tick):
        commands = []
        for robot in self._robots:
            commands.append(robot.control(self))
        await asyncio.gather(*commands)
        self._robots = [robot for robot in self._robots if robot.error is None]

class VehicleController:
    def __init__(self, robot_id):
        self.takeoff = True
        self.robot_id = robot_id
        self.entity = None
        self.error = None

    async def spawn(self, world):
        angle = (2 * math.pi / ROBOT_COUNT) * self.robot_id

        self.entity = await world.rpc('map.spawn_uav', {
            'robot_id': self.robot_id,
            'position': {
                'x': SPAWN_RADIUS * math.cos(angle),
                'y': SPAWN_RADIUS * math.sin(angle),
                'z': 0,
            },
        })

    async def control(self, world):
        try:
            TAKEOFF_ALTITUDE = 2

            position = await world.rpc(f'object_{self.entity}.position', None)

            if self.takeoff and position['z'] > TAKEOFF_ALTITUDE - 0.05:
                self.takeoff = False

            if self.takeoff:
                await world.rpc(f'object_{self.entity}.attitude_control', {
                    'z': TAKEOFF_ALTITUDE,
                    'vxy': 0,
                    'zx': 0,
                    'yz': 0,
                })
                return

            distance_from_zero = math.sqrt(position['x'] ** 2 + position['y'] ** 2)
            angle = math.atan2(position['y'], position['x'])

            # await world.rpc(f'object_{self.entity}.velocity_control', {
            #     'z': position['z'],
            #     'vxy': 0,
            #     'vx': math.cos(angle) * SPAWN_RADIUS,
            #     'vy': math.sin(angle) * SPAWN_RADIUS,
            # })

            await world.rpc(f'object_{self.entity}.position_control', {
                'z': position['z'],
                'x': math.cos(angle + 0.7) * distance_from_zero * 1.5,
                'y': math.sin(angle + 0.7) * distance_from_zero * 1.5,
                'xy': 0,
            })
        except Exception as e:
            self.error = e
            print(f"vehicle {self.robot_id} error: {e}")

WorldController()
asyncio.get_event_loop().run_forever()
