onemillioncheckboxes/swarm/manager.py

from multiprocessing.shared_memory import SharedMemory
from typing import Callable, NamedTuple, Optional
import asyncio
import socketio
import aiohttp
from PIL import (
    Image,
    ImageFont,
    ImageDraw,
    ImageFilter,
    ImageSequence,
    ImageChops,
)
from base64 import b64decode
import signal
import os
import time
import json
import datetime


class Animation(NamedTuple):
    x: int
    y: int
    frames: list[Image.Image]
    spf: float


OFFSET_STATE = 0
OFFSET_AVOID = 125000
OFFSET_CANVAS = 250000
OFFSET_MASK = 375000


class Manager:
    def __init__(self, settings_path: str):
        self.shmem: Optional[SharedMemory] = None
        self.shmem_name = "omcb-bot"

        self.settings_path = settings_path

        self.base = "https://onemillioncheckboxes.com"
        self.last_update = 0

        self.bits_toggled_on = 0
        self.bits_toggled_off = 0
        self.last_printout = 0

        self.fonts: dict[tuple[str, int], ImageFont.FreeTypeFont] = {}
        self.default_font_size = 8
        self.animations: list[Animation] = []
        self.animation_functions: list[Callable] = []

    def reload_config(self):
        with open(self.settings_path, "r") as fp:
            settings = json.load(fp)

        assert self.shmem is not None

        print("Resetting shmem...")
        self.shmem.buf[OFFSET_AVOID:] = bytes(500000 - OFFSET_AVOID)

        if fontconfig := settings.get("default_font"):
            self.default_font_size = int(fontconfig.get("size", 8))
            self.fonts["default", self.default_font_size] = self.get_font(
                fontconfig["path"], self.default_font_size
            )

        for avoid in settings.get("avoid", []):
            if avoid["type"] == "rect":
                self.add_avoid_rect(
                    avoid["x"], avoid["y"], avoid["w"], avoid["h"]
                )
            elif avoid["type"] == "range":
                self.add_avoid_range(
                    range(avoid["start"], avoid["stop"], avoid.get("step", 1))
                )
            elif avoid["type"] == "image":
                with Image.open(avoid["path"]).convert("LA") as im:
                    assert im.width == 1000 and im.height == 1000
                    pixels = im.load()
                    assert pixels is not None
                    for y in range(im.height):
                        for x in range(im.width):
                            _, a = pixels[(x, y)]  # type: ignore
                            if a > 128:
                                self.add_avoid_index(x + y * 1000)
            else:
                raise ValueError(f"invalid avoid: {avoid}")
            print("AVOID", avoid)

        for elem in settings.get("elements", []):
            if elem["type"].startswith("~"):
                continue
            elif elem["type"] == "text":
                self.put_text(
                    elem["x"],
                    elem["y"],
                    elem["text"],
                    elem.get("font", "default"),
                    elem.get("size", self.default_font_size),
                    elem.get("negative", False),
                    elem.get("padding", 2),
                )
            elif elem["type"] in ("text_anim", "text_animation"):
                frames: list[Image.Image] = []
                for line in elem["lines"]:
                    if isinstance(line, str):
                        frames.append(
                            self.make_text_image(
                                line,
                                elem.get("font", "default"),
                                elem.get("size", 8),
                                elem.get("negative", False),
                                elem.get("padding", 2),
                            )
                        )
                    elif isinstance(line, dict):
                        frames.append(
                            self.make_text_image(
                                line["text"],
                                line.get("font", elem.get("font", "default")),
                                line.get(
                                    "size",
                                    elem.get("size", self.default_font_size),
                                ),
                                line.get(
                                    "negative", elem.get("negative", False)
                                ),
                                line.get("padding", elem.get("padding", 2)),
                            )
                        )
                    else:
                        raise TypeError(f"invalid line: {line}")
                self.put_animation(
                    elem["x"], elem["y"], frames, elem.get("spf", 10)
                )
            elif elem["type"] in ("anim", "animation"):
                with Image.open(elem["path"]) as anim:
                    self.put_animation(
                        elem["x"],
                        elem["y"],
                        [
                            frame.convert("LA")
                            for frame in ImageSequence.Iterator(anim)
                        ],
                        elem.get("spf", 10),
                    )
            elif elem["type"] == "image":
                with Image.open(elem["path"]).convert("LA") as im:
                    self.put_image(elem["x"], elem["y"], im)
            elif elem["type"] == "time":
                pos_x, pos_y = elem["x"], elem["y"]

                time_format = elem.get("format", "%Y%m%dT%H%M%S UTC")
                time_font = elem.get("font", "default")
                time_font_size = elem.get("size", self.default_font_size)
                time_negative = elem.get("negative", False)
                time_outline = elem.get("outline", 2)

                def update_time():
                    now = datetime.datetime.now(datetime.timezone.utc)
                    self.put_image(
                        pos_x,
                        pos_y,
                        self.make_text_image(
                            now.strftime(time_format),
                            time_font,
                            time_font_size,
                            time_negative,
                            time_outline,
                        ),
                    )

                self.animation_functions.append(update_time)
            elif elem["type"] == "shrek":
                shrek_x, shrek_y = elem["x"], elem["y"]

                shrek_font = self.get_font(
                    elem.get("font", "default"),
                    elem.get("size", self.default_font_size),
                )

                shrek_spf = elem.get("spf", 10)

                with open(elem["path"], "r") as fp:
                    lyrics = list(map(str.strip, fp))

                def update_shrek():
                    with Image.new("LA", (325, 10), (0, 255)) as im:
                        draw = ImageDraw.Draw(im)
                        draw.rectangle((0, 0, 325, 10), fill=(0, 255))
                        now = datetime.datetime.now(datetime.timezone.utc)
                        line = lyrics[
                            int(now.timestamp() / shrek_spf) % len(lyrics)
                        ]
                        draw.text(
                            (2, -1), line, font=shrek_font, fill=(255, 255)
                        )
                        self.put_image(shrek_x, shrek_y, im)

                self.animation_functions.append(update_shrek)
            else:
                raise TypeError(f"invalid element: {elem}")
            print("ADD", elem)

    async def listener(self):
        sio = socketio.AsyncClient()
        sio.on("connect", self.on_connect)
        sio.on("batched_bit_toggles", self.on_batched_bit_toggles)
        sio.on("full_state", self.on_full_state)

        await sio.connect(self.base.replace("http", "ws"))
        await sio.wait()

    def update_shmem(self, state: bytes):
        assert self.shmem is not None
        self.shmem.buf[OFFSET_STATE : OFFSET_STATE + 125000] = state

    async def on_connect(self):
        async with aiohttp.ClientSession() as http:
            async with http.get(f"{self.base}/api/initial-state") as req:
                data = await req.json()
                buffer = b64decode(data["full_state"].encode() + b"=")
                self.update_shmem(buffer)
                self.last_update = data["timestamp"]

    async def on_full_state(self, data):
        assert self.shmem is not None
        buffer = b64decode(data["full_state"].encode() + b"=")
        self.update_shmem(buffer)
        self.last_update = data["timestamp"]

    async def on_batched_bit_toggles(self, data):
        assert self.shmem is not None
        bits_on, bits_off, timestamp = data
        if timestamp < self.last_update:
            print("old update, ignoring")

        self.last_update = timestamp

        self.bits_toggled_on += len(bits_on)
        self.bits_toggled_off = len(bits_off)

        for ndx in bits_on:
            byte, bit = divmod(ndx, 8)
            self.shmem.buf[OFFSET_STATE + byte] |= 0x80 >> bit
        for ndx in bits_off:
            byte, bit = divmod(ndx, 8)
            self.shmem.buf[OFFSET_STATE + byte] &= 0xFF ^ (0x80 >> bit)

        since_last_printout = time.time() - self.last_printout
        if since_last_printout >= 5:
            self.last_printout = time.time()
            print()
            print(
                f"Toggled on: {self.bits_toggled_on / since_last_printout}/s"
            )
            print(
                f"Toggled off: {self.bits_toggled_off / since_last_printout}/s"
            )
            self.bits_toggled_on = self.bits_toggled_off = 0

    def on_sigusr1(self, signum, frame):
        assert self.shmem is not None
        print("Caught SIGUSR1, dumping state")
        buf = bytes(self.shmem.buf[:])
        with Image.new("RGB", (1000, 1000), 0) as im:
            for i in range(1000000):
                y, x = divmod(i, 1000)
                byte, bit = divmod(i, 8)
                im.putpixel(
                    (x, y),
                    (
                        255 if (buf[OFFSET_MASK + byte] << bit) & 0x80 else 0,
                        (
                            255
                            if (buf[OFFSET_CANVAS + byte] << bit) & 0x80
                            else 0
                        ),
                        255 if (buf[OFFSET_STATE + byte] << bit) & 0x80 else 0,
                    ),
                )
            im.save("state.png")
        with Image.new("L", (1000, 1000), 0) as im:
            for i in range(1000000):
                y, x = divmod(i, 1000)
                byte, bit = divmod(i, 8)
                im.putpixel(
                    (x, y),
                    255 if (buf[OFFSET_AVOID + byte] << bit) & 0x80 else 0,
                )
            im.save("avoid.png")
        print("Dump done")

    def on_sigusr2(self, signum, frame):
        print("Reloading config")
        self.reload_config()

    async def animator(self):
        while True:
            for animation in self.animations:
                frame = int(time.time() / animation.spf) % len(
                    animation.frames
                )
                self.put_image(
                    animation.x, animation.y, animation.frames[frame]
                )
            for func in self.animation_functions:
                func()
            await asyncio.sleep(0.1)

    def add_avoid_range(self, rng: range):
        assert self.shmem is not None
        for ndx in rng:
            byte, bit = divmod(ndx, 8)
            self.shmem.buf[OFFSET_AVOID + byte] |= 0x80 >> bit

    def add_avoid_rect(self, sx: int, sy: int, w: int, h: int):
        for y in range(sy, sy + h):
            ox = y * 1000
            self.add_avoid_range(range(sx + ox, sx + w + ox))

    def add_avoid_index(self, index: int):
        assert self.shmem is not None
        assert 0 <= index < 1000000
        byte, bit = divmod(index, 8)
        self.shmem.buf[OFFSET_AVOID + byte] |= 0x80 >> bit

    def add_avoid_image(self, im: Image.Image):
        assert im.width == 1000
        assert im.height == 1000
        pixels = im.load()
        assert pixels is not None
        for i in range(1000000):
            y, x = divmod(i, 1000)
            la = pixels[(x, y)]
            assert isinstance(la, (tuple, list)) and len(la) == 2
            if la[1] > 128:
                self.add_avoid_index(i)

    def set_index(self, index: int, value: bool):
        assert 0 <= index <= 1000000
        assert self.shmem is not None
        byte, bit = divmod(index, 8)
        self.shmem.buf[OFFSET_MASK + byte] |= 0x80 >> bit
        if value:
            self.shmem.buf[OFFSET_CANVAS + byte] |= 0x80 >> bit
        else:
            self.shmem.buf[OFFSET_CANVAS + byte] &= 0xFF ^ (0x80 >> bit)

    def clear_index(self, index: int):
        assert 0 <= index <= 1000000
        assert self.shmem is not None
        byte, bit = divmod(index, 8)
        self.shmem.buf[OFFSET_MASK + byte] &= 0xFF ^ (0x80 >> bit)
        self.shmem.buf[OFFSET_CANVAS + byte] &= 0xFF ^ (0x80 >> bit)

    def put_image(self, ox: int, oy: int, im: Image.Image):
        pixels = im.load()
        assert pixels is not None
        for y in range(im.height):
            for x in range(im.width):
                la: tuple[int, int] = pixels[(x, y)]  # type: ignore
                if la[1]:
                    self.set_index(x + ox + (y + oy) * 1000, la[0] > 0)

    def put_text(
        self,
        x: int,
        y: int,
        text: str,
        font_name: str = "default",
        font_size: int = 8,
        negative: bool = False,
        outline: int = 2,
    ):
        self.put_image(
            x,
            y,
            self.make_text_image(
                text, font_name, font_size, negative, outline
            ),
        )

    def make_text_image(
        self,
        text: str,
        font_name: str = "default",
        font_size: int = 8,
        negative: bool = False,
        outline: int = 2,
    ):
        font = self.get_font(font_name, font_size)
        left, top, right, bottom = font.getbbox(text, anchor="lt")
        with Image.new(
            "RGBA",
            (int(right - left) + outline * 2, int(bottom - top) + outline * 2),
            0,
        ) as im:
            draw = ImageDraw.Draw(im)
            draw.rectangle((0, 0, im.width, im.height), (0, 0, 0, 0))
            draw.text(
                (left + outline, top + outline),
                text,
                font=font,
                fill=(255, 255, 255, 255),
                anchor="lt",
            )

            alpha = im.convert("L").filter(
                ImageFilter.MaxFilter(outline * 2 + 1)
            )
            im.putalpha(alpha)
            if negative:
                im = ImageChops.invert(im)
            im.putalpha(alpha)  # ty PIL
            return im.convert("LA")

    def put_animation(
        self, x: int, y: int, frames: list[Image.Image], spf: float = 10
    ):
        self.animations.append(Animation(x, y, frames, spf))

    def get_font(
        self, font_name: str, font_size: int
    ) -> ImageFont.FreeTypeFont:
        if font := self.fonts.get((font_name, font_size)):
            return font
        print("FONT", font_name, font_size)
        font = ImageFont.truetype(font_name, font_size)
        self.fonts[font_name, font_size] = font
        return font

    async def __aenter__(self):
        print("Acquiring shared memory")
        self.shmem = SharedMemory(self.shmem_name, True, 500000)
        print("Loading config...")
        try:
            self.reload_config()
        except Exception:
            self.shmem.close()
            self.shmem.unlink()
            raise
        return self

    async def __aexit__(self, a, b, c):
        if self.shmem:
            print("cleaning up shmem")
            self.shmem.close()
            self.shmem.unlink()


async def main(cfg_path: str = "./settings.json", *_):
    print(f"PID: {os.getpid()}")
    async with Manager(cfg_path) as mgr:
        signal.signal(signal.SIGUSR1, mgr.on_sigusr1)
        signal.signal(signal.SIGUSR2, mgr.on_sigusr2)

        print("Listening...")
        await asyncio.gather(mgr.listener(), mgr.animator())


if __name__ == "__main__":
    from sys import argv

    asyncio.run(main(*argv[1:]))