import numpy as np
import base64
import json
import zlib
from imshow import network
SocketServer = network.SocketServer
SocketMessage = network.SocketMessage

class ImageMessage(SocketMessage):
	def __init__(self, img=None, name = 'Image', wait=0):
		super(ImageMessage, self).__init__()
		self.img = img
		self.wait = wait
		self.name = name
		self.type = 'image'

	def encode(self):
		img_bytes = self.img.tobytes()
		img_bytes = base64.b64encode(img_bytes).decode('utf-8')
		img_shape = self.img.shape
		self.data['img'] = img_bytes
		self.data['dtype'] = self.img.dtype.char
		self.data['shape'] = img_shape
		self.data['wait'] = self.wait
		self.data['name'] = self.name

	def decode(self):
		byte = base64.b64decode(self.data['img'])
		img = np.frombuffer(byte, dtype=np.dtype(self.data['dtype']))
		img = img.reshape(self.data['shape'])
		self.img = img
		self.wait = self.data['wait']
		self.name = str(self.data['name'])