from base64 import b64decode
import tkinter as tk
from PIL import Image, ImageTk
from socketio import SimpleClient, exceptions
from threading import Thread
from requests import get

COLORS = [(0x33, 0x33, 0x66), (0x96, 0x96, 0xe0)]
COLORS_UNPACKED = COLORS[0] + ((0, 0, 0) * 254) + COLORS[1]


class App(tk.Tk):
    def __init__(self, url: str = "https://onemillioncheckboxes.com") -> None:
        self.url = url
        super().__init__()
        self.title("1M pixels")

        self.sio = SimpleClient()

        self._canvas = tk.Canvas(self)
        self._canvas.config(width=1000, height=1000, borderwidth=0, highlightthickness=0)
        self._canvas.pack()
        self.config(width=1000, height=1000, borderwidth=0, highlightthickness=0)

        self._running = False

        self._image = Image.new("RGB", (1000, 1000), COLORS[0])

        self._canvas.bind("<Motion>", self._on_mouse_move)
        self._canvas.bind("<Button>", self._on_mouse_click)
        self.bind("<KeyPress>", self._on_key_down)


    def _on_mouse_move(self, event: tk.Event):
        self.title(f"1M pixels: {event.x}, {event.y}")

    def _on_mouse_click(self, event: tk.Event):
        x, y = event.x, event.y
        self.sio.emit("toggle_bit", { "index": x + y * 1000 })

    def _on_key_down(self, event: tk.Event):
        # <KeyPress event state=Control keysym=r keycode=27 char='\x12' x=538 y=556>
        if event.keysym == "r":
            print("FULL REFRESH")
            with get(f"{self.url}/api/initial-state") as req:
                buffer = b64decode(req.json()["full_state"].encode() + b"=")
                img = Image.frombytes("1", (1000, 1000), buffer).convert("P")
                img.putpalette(COLORS_UNPACKED)
                self._image.paste(img.convert("RGB"))
            print("FULL REFRESH DONE")

    def _reader(self):
        while self._running:
            try:
                name, data = self.sio.receive()
            except exceptions.TimeoutError:
                print("Timeout")
                continue

            if name == "batched_bit_toggles":
                bits_on, bits_off, timestamp = data
                for ndx in bits_on:
                    y, x = divmod(ndx, 1000)
                    self._image.putpixel((x, y), COLORS[1])
                for ndx in bits_off:
                    y, x = divmod(ndx, 1000)
                    self._image.putpixel((x, y), COLORS[0])
                self._tk_image.paste(self._image)
                print("partial update", len(bits_on), len(bits_off))

            elif name == "full_state":
                print("full update")
                buffer = b64decode(data["full_state"].encode() + b"=")
                img = Image.frombytes("1", (1000, 1000), buffer).convert("P")
                img.putpalette(COLORS_UNPACKED)
                self._image.paste(img.convert("RGB"))
                self._tk_image.paste(self._image)

            else:
                print(name, data)

        self.sio.disconnect()
        print("_reader exited")

    def _close(self):
        self._running = False
        print("waiting for reader to close")
        self._reader_thr.join()
        self.destroy()

    def start(self):
        self.sio.connect(f"{self.url}/socket.io")

        with get(f"{self.url}/api/initial-state") as req:
            buffer = b64decode(req.json()["full_state"].encode() + b"=")
            img = Image.frombytes("1", (1000, 1000), buffer).convert("P")
            img.putpalette(COLORS_UNPACKED)
            self._image.paste(img.convert("RGB"))
            self._tk_image = ImageTk.PhotoImage(self._image)
            self._tk_image_id = self._canvas.create_image(0, 0, anchor="nw", image=self._tk_image)
        self._running = True

        self._reader_thr = Thread(target=self._reader, name="Reader thread")
        self._reader_thr.start()

        self.protocol("WM_DELETE_WINDOW", self._close)
        self.mainloop()

if __name__ == "__main__":
    app = App()
    app.start()