import asyncio
from ws.connection import IConnection
import json


class MsgIo:
	def __init__(self, key, send, send_from_thread):
		self.key = key
		self.sender = send
		self.send_from_thread_cb = send_from_thread
		self.receiver = None

	def set_recv_handler(self, handler):
		self.receiver = handler

	async def on_recv(self, data):
		print("MessageHandler {}".format(data))
		if self.receiver is not None:
			await self.receiver(data[self.key])

	async def send(self, data):
		return await self.sender({self.key: data})

	def send_from_thread(self, data):
		return self.send_from_thread_cb({self.key: data})

	def can_key(self, key):
		return self.key == key

	def get_key(self):
		return self.key


class MessageDispatcher(IConnection):
	def __init__(self, auto_subscribe=False, loop=None):
		self.msg_handlers: MsgIo = []
		self.state = None
		self.state = asyncio.Queue(loop=loop)
		self.auto_subscribe = auto_subscribe
		self.loop = loop

	def msgio_get(self, key):
		obj = MsgIo(key, self.send, self.send_from_thread)
		if key not in self.msg_handlers:
			self.msg_handlers.append(obj)
			return obj
		return None

	async def on_connect(self):
		if self.auto_subscribe:
			for handler in self.msg_handlers:
				key = handler.get_key()
				print("Would subscribe {}".format(key))
				await self.send({'+': key})

	async def on_disconnect(self):
		pass

	async def on_recv(self, data):
		d = json.loads(data)
		for key_req in d.keys():
			for h in self.msg_handlers:
				if h.can_key(key_req):
					await h.on_recv(d)

	async def on_send(self):
		return await self.state.get()

	async def send(self, data):
		await self.state.put(data)

	def send_from_thread(self, data):
		asyncio.ensure_future(self.send(data), loop=self.loop)