from fl_networking_tools import ImageViewer
import socket
import pickle

received_pixels = expected_pixels = 0
last_pixel_updated = (-1, -1)
max_x = max_y = 0

def get_pixel_data():
    global last_pixel_updated, max_x, max_y, received_pixels, expected_pixels

    while True:
        #unpack data
        data, client_addr = udp_server.recvfrom(1024)
        msg = pickle.loads(data)
        pos = msg[0]
        rgba = msg[1]

        #compare to expected
        if (pos[0] > max_x):
            max_x = pos[0]
            expected_pixels = (max_x+1) * (max_y+1)
        if (pos[1] > max_y):
            max_y = pos[1]
            expected_pixels = (max_x+1) * (max_y+1)
        received_pixels += 1

        # update viewer
        viewer.put_pixel(pos, rgba)
        viewer.text = f"Pixels not received = {expected_pixels - received_pixels}, expected = {expected_pixels}, received = {received_pixels}, max_x = {max_x}, max_y = {max_y}"
        last_pixel_updated = pos

udp_server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
udp_server.bind(("0.0.0.0", 20001))

viewer = ImageViewer()
viewer.start(get_pixel_data)