onemillioncheckboxes/swarm/manager.py

520 lines
18 KiB
Python

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: int = 0
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)
self.animations = []
self.animation_functions = []
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"].startswith("~"):
continue
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"] == "rect":
with Image.new(
"LA",
(elem["w"], elem["h"]),
((255 if elem["fill"] else 0), 255),
) as im:
self.put_image(elem["x"], elem["y"], im)
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),
elem.get("offset", 0)
)
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)
elif elem["type"] == "minimap":
minimap_x, minimap_y = elem["x"], elem["y"]
minimap_w, minimap_h = elem["w"], elem["h"]
minimap_state = {
"last_update": time.time()
}
def update_minimap():
if self.shmem is None:
return
if (time.time() - minimap_state["last_update"]) < 5:
return
for y in range(minimap_h):
for x in range(minimap_w):
world_x = int(x * 1000 / minimap_w)
world_y = int(y * 1000 / minimap_h)
world_i = world_x + world_y * 1000
byte, mask = world_i >> 3, 0x80 >> (world_i & 7)
state = bool(self.shmem.buf[byte] & mask)
px, py = x + minimap_x, y + minimap_y
self.set_index(px + py * 1000, state)
minimap_state["last_update"] = time.time()
self.animation_functions.append(update_minimap)
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 = ndx >> 3, ndx & 7
self.shmem.buf[OFFSET_STATE + byte] |= 0x80 >> bit
for ndx in bits_off:
byte, bit = ndx >> 3, ndx & 7
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[:])
print("# Dumping state.png")
with Image.new("RGB", (1000, 1000), 0) as im:
for i in range(1000000):
y, x = divmod(i, 1000)
byte, bit = i >> 3, i & 7
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")
print("# Dumping avoid.png")
with Image.new("L", (1000, 1000), 0) as im:
for i in range(1000000):
y, x = divmod(i, 1000)
byte, bit = i >> 3, i & 7
im.putpixel(
(x, y),
255 if (buf[OFFSET_AVOID + byte] << bit) & 0x80 else 0,
)
im.save("avoid.png")
print(f"# Animations {len(self.animations)}")
for anim in self.animations:
frame = int(time.time() / anim.spf + anim.offset) % len(anim.frames)
print(
str.format(
"X: {:4d} Y: {:4d} Frames: {:3d} S/F: {:4.1f} Frame: {:d}",
anim.x,
anim.y,
len(anim.frames),
anim.spf,
frame,
)
)
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 + animation.offset) % 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 = ndx >> 3, ndx & 7
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 = index >> 3, index & 7
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 = index >> 3, index & 7
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 = index >> 3, index & 7
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, offset: int = 0
):
self.animations.append(Animation(x, y, frames, spf, offset))
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:]))