import asyncio
from centrifuge import Client, SubscriptionEventHandler, PublicationContext

async def find_vehicle_by_id(client, vehicle_id):
    response = await client.rpc('script.eval', {
        "code": """
            function find_vehicle_by_id(vehicle_id)
                local found = sim.map.find_all({
                    groups = { "vehicle" },
                })
                for idx = 1, #found do
                    if sim.object.get_robot_id(found[idx].entity) == vehicle_id then
                        return found[idx].entity
                    end
                end
                return nil
            end

            return find_vehicle_by_id(ARGS.id)
        """,
        "args": {
            "id": vehicle_id,
        }
    })

    return response.data

async def main():
    client = Client('ws://localhost:9120/connection/websocket')
    await client.connect()

    vehicle = await find_vehicle_by_id(client, 1)
    if not vehicle:
        print('No vehicle found')
        return

    class EventHandler(SubscriptionEventHandler):
        async def on_publication(self, ctx: PublicationContext) -> None:
            print("drone position:", ctx.pub.data)

    sub = client.new_subscription(f'object_{vehicle}.position', EventHandler())
    await sub.subscribe()


asyncio.ensure_future(main())
loop = asyncio.get_event_loop()
loop.run_forever()
