import tornado.websocket
import tornado.web
import tornado.ioloop
import os
from PIL import Image
import numpy as np
from blend_modes import blend_modes
from io import BytesIO

web_dir = os.path.join(os.path.split(__file__)[0], 'www')
output_dir = os.path.join(os.path.split(__file__)[0], 'output')

# This is our WebSocketHandler - it handles the messages
# from the tornado server
class WebSocketHandler(tornado.websocket.WebSocketHandler):
	connections = set()

	# the client connected
	def open(self):
		self.connections.add(self)
		print ("New client connected")

	# the client sent the message
	def on_message(self, message):
		[con.write_message(message) for con in self.connections]


	# client disconnected
	def on_close(self):
		self.connections.remove(self)
		print ("Client disconnected")

class CompositeHandler(tornado.web.RequestHandler):
	def get(self, img1, img2):
		# images can only be in the output dir
		img1 = os.path.join(output_dir, os.path.basename(img1))
		img2 = os.path.join(output_dir, os.path.basename(img2))

		i1 = np.asarray(Image.open(img1).resize((200,200)).convert('RGBA'), dtype=float)
		i2 = np.asarray(Image.open(img2).resize((200,200)).convert('RGBA'), dtype=float)

		diffI = blend_modes.difference(i1,i2, opacity=1)
		noAlpha = diffI[:,:,0:3]
		maxValue = np.max(noAlpha)
		# make sure a black image wont crash the system
		factor = 255/maxValue if maxValue>0 else 0
		composite = Image.fromarray(np.uint8(noAlpha * factor))
		byteImgIO = BytesIO()
		composite.save(byteImgIO, format="png")
		byteImgIO.seek(0)
		# print(len(byteImgIO.read()))
		self.write(byteImgIO.read())
		self.set_header("Content-type",  "image/png")


application = tornado.web.Application([
	(r"/ws", WebSocketHandler),
	(r"/output/(.*)", tornado.web.StaticFileHandler, {"path": output_dir}),
	(r"/composite/(.*)/(.*)", CompositeHandler),
	(r"/(.*)", tornado.web.StaticFileHandler, {"path": web_dir, "default_filename": 'index.html'}),
],debug=True)

application.listen(8888)
tornado.ioloop.IOLoop.instance().start()