import asyncio
from typing import NewType, Optional
from socketio import AsyncClient, AsyncSimpleClient
from aiohttp import ClientSession
from aiohttp_socks import ProxyConnector
from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageSequence, ImageChops
from base64 import b64decode
from random import choice
from json import load
from time import time as time_now


PixelMap = NewType("PixelMap", dict[int, bool])
Animation = NewType("Animation", tuple[list[PixelMap], float])


class AsyncBotManager:
    def __init__(self, base: str = "https://onemillioncheckboxes.com"):
        self.base = base

        self.canvas = Image.new("1", (1000, 1000))
        self.font = ImageFont.load_default(8)

        self.difference: PixelMap = PixelMap({})
        self._last_update = 0
        self._shutdown: bool = False
        self.avoid: set[int] = set()

        self.animations: list[Animation] = []
        self._change_count = 0
        self._last_change_printout = time_now()

    @staticmethod
    def get_text_image(text: str, font: ImageFont.ImageFont | ImageFont.FreeTypeFont) -> Image.Image:
        with Image.new("LA", (int(font.getlength(text) + 12), 16)) as im:
            draw = ImageDraw.Draw(im)
            draw.rectangle((0, 0, im.width, im.height), (0, 0))
            draw.text((6, 5), text, font=font, fill=(255, 0))

            alpha = im.convert("L").filter(ImageFilter.MaxFilter(3))
            im.putalpha(alpha)
            return im.copy()

    def put_text(self, x: int, y: int, text: str):
        self.put_image(x, y, self.get_text_image(text, self.font))

    def put_image(self, ox: int, oy: int, im: Image.Image):
        for y in range(im.height):
            for x in range(im.width):
                l, a = im.getpixel((x, y))  # type: ignore
                index = x + ox + (y + oy) * 1000
                if a:
                    self.put_bit(index, l > 0)

    def put_pixel(self, x: int, y: int, val: bool):
        self.put_bit(x + y * 1000, val)

    def put_bit(self, index: int, value: bool):
        if index not in self.avoid:
            self.difference[index] = value

    def add_animation(
        self, ox: int, oy: int, frames: list[Image.Image], spf: float = 1
    ) -> None:
        animation = Animation(([], spf))
        alpha = Image.new("L", frames[0].size, 0)
        for frame in frames:
            frame_la = frame.convert("LA")
            frame_alpha = frame_la.getchannel("A")
            alpha = ImageChops.add(alpha, frame_alpha)

        for frame in frames:
            frame_la = frame.convert("LA")
            pixelmap = PixelMap({})
            for y in range(frame_la.height):
                for x in range(frame_la.width):
                    l, a = frame_la.getpixel((x, y))  # type: ignore
                    ga: int = alpha.getpixel((x, y))  # type: ignore
                    index = x + ox + (y + oy) * 1000
                    if (a > 128 or ga > 128) and index not in self.avoid:
                        pixelmap[index] = l > 128 if a > 128 else True
            animation[0].append(pixelmap)
        self.animations.append(animation)

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

    def add_avoid_range(self, start: int, stop: int, step: int = 1) -> None:
        self.avoid |= set(range(start, stop, step))

    def add_avoid_index(self, *indices: int) -> None:
        self.avoid |= set(indices)

    def get_difference_image(self) -> Image.Image:
        with Image.new("LA", (1000, 1000), 0) as im:
            for index, expected in self.difference.items():
                y, x = divmod(index, 1000)
                im.putpixel((x, y), (255 if expected else 0, 255))
            return im.copy()

    def get_avoid_image(self) -> Image.Image:
        with Image.new("RGB", (1000, 1000), (0, 255, 0)) as im:
            for index in self.avoid:
                y, x = divmod(index, 1000)
                im.putpixel((x, y), (255, 0, 0))
            return im.copy()

    async def listener(self) -> None:
        async with 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.canvas.paste(Image.frombytes("1", (1000, 1000), buffer))
                self._last_update = data["timestamp"]
            async with AsyncSimpleClient(http_session=http) as sio:
                await sio.connect(f"{self.base}/socket.io")
                while not self._shutdown:
                    event, data = await sio.receive()
                    if event == "full_state":
                        buffer = b64decode(data["full_state"].encode() + b"=")
                        image = Image.frombytes("1", (1000, 1000), buffer)
                        self.canvas.paste(image)
                        image.close()
                        self._last_update = data["timestamp"]
                    elif event == "batched_bit_toggles":
                        bits_on, bits_off, timestamp = data
                        if timestamp < self._last_update:
                            print("SKIPPING UPDATES: TOO OLD")
                        self._last_update = timestamp
                        for ndx in bits_on:
                            y, x = divmod(ndx, 1000)
                            self.canvas.putpixel((x, y), 255)
                        for ndx in bits_off:
                            y, x = divmod(ndx, 1000)
                            self.canvas.putpixel((x, y), 0)
                    else:
                        print("unknown event", event, data)
                    now = time_now()
                    if (now - self._last_change_printout) > 1:
                        cps = self._change_count / (now - self._last_change_printout)
                        print(f"Speed: {cps:.2f}/s")
                        self._change_count = 0
                        self._last_change_printout = now
                    for pixmaps, spf in self.animations:
                        frame_index = int(now / spf)
                        self.difference.update(
                            pixmaps[frame_index % len(pixmaps)]
                        )

    async def writer(
        self,
        bot_index: int,
        proxy_url: Optional[str] = None,
        delay: float = 0.25,
    ):
        proxy = ProxyConnector.from_url(proxy_url) if proxy_url else None
        async with ClientSession(connector=proxy) as http:
            sio = AsyncClient(http_session=http)
            await sio.connect(f"{self.base}/socket.io")
            while not self._shutdown:
                diff = list(self.difference.items())
                for _ in range(100):
                    index, expected = choice(diff)
                    y, x = divmod(index, 1000)
                    current = self.canvas.getpixel((x, y)) > 0  # type: ignore
                    if current != expected:
                        # print(f"[{bot_index:2d}] swap {x:3d} {y:3d}")
                        self._change_count += 1
                        await sio.emit("toggle_bit", {"index": index})
                        await asyncio.sleep(delay)
                        break
                await asyncio.sleep(0.1)

    async def __aenter__(self):
        self._listener_task = asyncio.create_task(self.listener())
        return self

    async def __aexit__(self, a, b, c):
        self._shutdown = True
        await self._listener_task


async def amain() -> None:
    with open("settings.json", "r") as fp:
        settings = load(fp)

    async with AsyncBotManager() as mgr:
        mgr.font = ImageFont.truetype(settings["font"], 8)
        for avoid in settings["avoid"]:
            if avoid["type"] == "rect":
                mgr.add_avoid_rect(
                    avoid["x"], avoid["y"], avoid["w"], avoid["h"]
                )
            elif avoid["type"] == "range":
                mgr.add_avoid_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
                    for y in range(im.height):
                        for x in range(im.width):
                            l, a = im.getpixel((x, y))  # type: ignore
                            if a > 128:
                                mgr.add_avoid_index(x + y * 1000)
        for elem in settings["elements"]:
            if elem["type"] == "text":
                mgr.put_text(elem["x"], elem["y"], elem["text"])
            elif elem["type"] == "text_anim":
                frames: list[Image.Image] = []
                for text in elem["lines"]:
                    frames.append(mgr.get_text_image(text, mgr.font))
                mgr.add_animation(elem["x"], elem["y"], frames, elem["spf"])
                for frame in frames:
                    frame.close()
            elif elem["type"] == "image":
                with Image.open(elem["path"]).convert("LA") as im:
                    mgr.put_image(elem["x"], elem["y"], im)
            elif elem["type"] == "animation":
                with Image.open(elem["path"]) as anim:
                    frames: list[Image.Image] = []
                    for frame in ImageSequence.Iterator(anim):
                        frames.append(frame.copy())
                    mgr.add_animation(elem["x"], elem["y"], frames, elem["spf"])
                    for frame in frames:
                        frame.close()
            elif elem["type"] == "rgb111":
                ox, oy = elem["x"], elem["y"]
                with Image.open(elem["path"]).convert("RGBA") as im:
                    for y in range(im.height):
                        for x in range(im.width):
                            r, g, b, a = im.getpixel((x, y))  # type: ignore
                            if a < 128:
                                continue

                            start_ndx = (x + ox + (y + oy) * 577) * 3
                            mgr.put_bit(start_ndx, r > 128)
                            mgr.put_bit(start_ndx + 1, g > 128)
                            mgr.put_bit(start_ndx + 2, b > 128)
            elif elem["type"] == "rgb565":
                ox, oy = elem["x"], elem["y"]
                with Image.open(elem["path"]).convert("RGBA") as im:
                    for y in range(im.height):
                        for x in range(im.width):
                            r, g, b, a = im.getpixel((x, y))  # type: ignore
                            if a < 128:
                                continue

                            offset: int = (x + ox + (y + oy) * 250) * 16

                            # pack rgb888 to rgb565
                            rgb: int = (r >> 3) << 11
                            rgb |= (g >> 2) << 5
                            rgb |= (b >> 3)

                            # write to a bitstream
                            for i in range(16):
                                mgr.put_bit(
                                    offset + i, ((rgb << i) & 0x8000) > 0
                                )
            elif elem["type"] == "rect":
                for y in range(elem["y"], elem["y"] + elem["h"]):
                    for x in range(elem["x"], elem["x"] + elem["w"]):
                        mgr.put_pixel(x, y, elem["fill"])
            elif elem["type"] == "blob":
                with open(elem["path"], "rb") as fp:
                    offset = elem["offset"]
                    length = elem.get("length", 1000000)
                    written = 0
                    while (char := fp.read(1)):
                        byte = char[0]
                        if written > length:
                            break
                        for i in range(8):
                            if written > length:
                                break
                            mgr.put_bit(offset, bool((byte >> (7 - i)) & 1))
                            written += 1
                            offset += 1

        mgr.get_difference_image().save("result.png")
        mgr.get_avoid_image().save("avoid.png")

        print("Starting writers...")
        if n_proxies := len(settings["proxies"]):
            await asyncio.gather(
                *[
                    mgr.writer(
                        i,
                        settings["proxies"][i % n_proxies],
                        settings["delay"],
                    )
                    for i in range(settings["n_bots"])
                ],
                return_exceptions=True,
            )
        else:
            await asyncio.gather(
                *[mgr.writer(i) for i in range(settings["n_bots"])],
                return_exceptions=True,
            )


if __name__ == "__main__":
    asyncio.run(amain())