import asyncio
import pygame
import aiomqtt
import struct

WIDTH = 320
HEIGHT = 320

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()


async def on_message(client, userdata, message):
    payload = bytes(message.payload)
    x, y, w, h = struct.unpack("HHHH", payload[:8])
    buf = payload[8:]
    for dy in range(h):
        for dx in range(w):
            off = w * dy + dx
            r, g, b = buf[off * 3 : off * 3 + 3]
            screen.set_at((x + dx, y + dy), (r, g, b))


async def main():
    async with aiomqtt.Client("mqtt2.bigasterisk.com") as client:
        client.on_message = on_message
        await client.subscribe('display/squib/updates')

        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    raise SystemExit

            await client.loop_read()
            pygame.display.flip()
            clock.tick(60)


asyncio.run(main())