#!/usr/bin/python3
import socket
import PIL
from PIL import Image, ImageColor
import argparse
import multiprocessing
import random
from pprint import pprint
from sys import exit

sizex = 800    
sizey = 600

threads = []

old_framebuffer = Image.new("RGBA", (sizex, sizey), None)
framebuffer = Image.new("RGBA", (sizex, sizey), (0,0,0,0))

parser = argparse.ArgumentParser()
mode = parser.add_mutually_exclusive_group(required=True)
mode.add_argument("--color")
mode.add_argument("--image", metavar="FILENAME")
parser.add_argument("--host", default="fsi-hase.fsmi.uni-karlsruhe.de")
parser.add_argument("--port", "-p", type=int, default="3141")
parser.add_argument("--yoffset", default=19, type=int)
parser.add_argument("--xoffset", default=0, type=int)
effect = parser.add_mutually_exclusive_group(required=False)
effect.add_argument("--repeat", action="store_true")
effect.add_argument("--animate", type=int, const=5, nargs='?', metavar="STEPSIZE")
parser.add_argument("--bounce", action="store_true")
parser.add_argument("--debug", action="store_true")
size = parser.add_mutually_exclusive_group(required=False)
size.add_argument("--size", default=[800, 600], type=int, nargs=2)
size.add_argument("--autodetect", action="store_true")
parser.add_argument("--threads", type=int, default="5")

args = parser.parse_args()

pprint(args)

sizex = args.size[0]
sizey = args.size[1]

dirtybox = (0,0,sizex,sizey)

XSPLIT = args.threads
YSPLIT = args.threads

if args.bounce and not args.animate:
    print("--bounce requires --animate")
    exit(1)

class DrawThread(multiprocessing.Process):
    def run(self):
        xfrom = self._kwargs["xfrom"]
        yfrom = self._kwargs["yfrom"]
        xto = self._kwargs["xto"]
        yto = self._kwargs["yto"]

        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((args.host, args.port))
        buf = bytearray()

        global framebuffer
        if args.debug:
            for x in range(xfrom,xto):
                for y in range(yfrom,yto):
                    if framebuffer.getpixel((x,y))[3] != 0:
                        value = b"%02X%02X%02X%02X"%self._kwargs["color"]
                        buf += b"PX %d %d %b\n"%(x+args.xoffset,y+args.yoffset,value)
        else:
            for x in range(xfrom,xto):
                for y in range(yfrom,yto):
                    if framebuffer.getpixel((x,y))[3] != 0:
                        value = b"%02X%02X%02X%02X"%framebuffer.getpixel((x,y))
                        buf += b"PX %d %d %b\n"%(x+args.xoffset,y+args.yoffset,value)

        s.send(buf)
        while args.repeat:
            s.send(buf)

        s.close()

def draw():
    global threads

    for thread in threads:
        thread.join()

    threads = []

    dirtysizex = (dirtybox[2]-dirtybox[0])
    dirtysizey = (dirtybox[3]-dirtybox[1])

    chunkxsize = dirtysizex//XSPLIT
    chunkysize = dirtysizey//YSPLIT

    remainderx = dirtysizex % XSPLIT
    remaindery = dirtysizey % YSPLIT

    offsetx = dirtybox[0]
    offsety = dirtybox[1]

    xchunks = []
    ychunks = []

    allocatedx = 0

    remainingx = (dirtybox[2]-dirtybox[0])
    remainingy = (dirtybox[3]-dirtybox[0])

    for i in range(dirtysizex//chunkxsize):
        start = offsetx + (i*chunkxsize)
        stop = offsetx + ((i+1)*chunkxsize)
        xchunks.append((start,stop))

    if remainderx != 0:
        xchunks.append(((offsetx+dirtysizex)-remainderx, offsetx+remainderx))

    for i in range(dirtysizey//chunkysize):
        start = offsety + (i*chunkysize)
        stop = offsety + ((i+1)*chunkysize)
        ychunks.append((start,stop))

    if remaindery != 0:
        ychunks.append(((offsety+dirtysizey)-remaindery, offsety+remaindery))

    for xchunk in xchunks:
        for ychunk in ychunks:
            if args.debug:
                color = (random.randint(0,255), random.randint(0,255), random.randint(0,255), 255)
                threads.append(DrawThread(kwargs={"xfrom":xchunk[0], "xto":xchunk[1], "yfrom":ychunk[0], "yto":ychunk[1], "color": color}))
            else:
                threads.append(DrawThread(kwargs={"xfrom":xchunk[0], "xto":xchunk[1], "yfrom":ychunk[0], "yto":ychunk[1]}))

    for thread in threads:
        thread.start()

if args.color:
    framebuffer.paste(ImageColor.getrgb(args.color), (0,0,sizex,sizey))
    draw()

elif args.image:
    xfactor = 1
    yfactor = 1

    coordinates = (0,0)
    image = Image.open(str(args.image))
    framebuffer.paste(image)
    if args.animate:
        dirtybox = (*coordinates, min(image.size[0]+coordinates[0], sizex), min(image.size[1]+coordinates[1], sizey))
    else:
        dirtybox = (0,0,min(image.size[0], sizex),min(image.size[1], sizey))
    draw()

    while args.animate:
        coordinates = (coordinates[0] + args.animate * xfactor, coordinates[1] + args.animate * yfactor)
        if coordinates[0] >= sizex or coordinates[1] >= sizey:
            break
        if args.bounce:
            if coordinates[0] + image.size[0] >= sizex:
                xfactor *= -1
            if coordinates[1] + image.size[1] >= sizey:
                yfactor *= -1
            if coordinates[0] <= 0:
                xfactor *= -1
            if coordinates[1] <= 0:
                yfactor *= -1
        dirtybox = (*coordinates, min(image.size[0]+coordinates[0], sizex), min(image.size[1]+coordinates[1], sizey))
        framebuffer = Image.new("RGBA", (sizex, sizey), (0,0,0,0))
        framebuffer.paste(image, coordinates)
        draw()