from socket import socket
from threading import Thread
from zlib import compress

from mss import mss
import time

Width = 1920
Height = 1080


def screenshot(conn):
    with mss() as sct:
        rect = {'top': 0, 'left':0, 'width': Width, 'height': Height}

        while 'recording':
            img = sct.grab(rect)

            pixels = compress(img.rgb, 6)

            size = len(pixels)
            size_len = (size.bit_length() + 7) // 8
            sock.sendto(bytes([size_len]))

            size_bytes = size.to_bytes(size_len, 'big')
            sock.sendto(size_bytes)

            sock.sentto(pixels)
            time.sleep(.25)

def main(host='127.0.0.1', port=5000):
    sock = socket()
    sock.bind((host, port))
    sock.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        sock.listen(5)
        print('Server Started')
        while 'connected':
            _, addr = sock.recvfrom()
            print('Client connected IP:', addr)
            thread = Thread(target=screenshot, args=(addr,))
            thread.start()
    finally:
        pass

if __name__ == '__main__':
    main()