from typing import Dict

import cv2
import socket
import numpy as np

from common import InterlacedPacket, from_bytes_int

# bind any IP address
UDP_IP = ""
UDP_PORT = 5005

sock = socket.socket(socket.AF_INET, # Internet
                     socket.SOCK_DGRAM) # UDP
sock.bind((UDP_IP, UDP_PORT))

HEIGHT = 480
WIDTH = 640

frames: Dict[str, np.ndarray] = {}

while True:
    # break the array down into 16-bit chunks, then transmit them as UDP packets
    for i in range(0, HEIGHT, 16):
        for j in range(0, WIDTH, 16):
            data, addr = sock.recvfrom(InterlacedPacket.size)  # buffer size is 768 bytes
            # print("received packet from", addr)
            pkt = from_bytes_int(data)

            uuid = str(pkt.uuid)
            x = pkt.x
            y = pkt.y
            arr = pkt.array

            if uuid not in frames.keys():
                frames[uuid] = frame = np.ndarray((HEIGHT, WIDTH, 3), dtype=np.uint8)

            if pkt.even:
                frames[uuid][y+1:y+32:2, x:x+16] = arr
            else:
                frames[uuid][y:y + 32:2, x:x + 16] = arr

    # Display the resulting frame
    for id in frames.keys():
        cv2.imshow(id, frames[id])

    # Break the loop if 'q' key is pressed
    if cv2.waitKey(1) == ord('q'):
        break

# Release the capture and close all windows
cv2.destroyAllWindows()