import asyncio
import websockets
from threading import Thread
from queue import Queue
from time import sleep
import json
from ws.connection import IConnection


class WsClient:
	def __init__(self, listener: IConnection):
		self.listener = listener
		self.loop = asyncio.new_event_loop()
		self.stop = None
		self.bg_thread = None

	def bg_run(self, uri):
		print("bg_run: started")
		asyncio.set_event_loop(self.loop)
		fut = asyncio.ensure_future(self.run_client(uri), loop=self.loop)
		self.loop.run_until_complete(fut)
		self.stop = None
		print("bg_run: terminated")

	def connect(self, uri):
		if self.bg_thread is not None and self.bg_thread.is_alive():
			return

		self.stop: asyncio.Future[None] = self.loop.create_future()
		self.bg_thread = Thread(target=self.bg_run, args=(uri,))
		self.bg_thread.setDaemon(True)
		self.bg_thread.start()

	def disconnect(self):
		if self.stop is not None:
			self.loop.call_soon_threadsafe(self.stop.set_result, None)
			self.bg_thread.join()

	async def run_client(self, uri):
		try:
			websocket = await websockets.connect(uri, loop=self.loop)
		except Exception as exc:
			print(f"Failed to connect to {uri}: {exc}.")
			return

		print("handler: got connection to {}".format(websocket.remote_address))
		self.listener.on_connect()
		try:
			path = "/"
			consumer_task = asyncio.ensure_future(self.handler_recv(websocket, path), loop=self.loop)
			producer_task = asyncio.ensure_future(self.handler_send(websocket, path), loop=self.loop)
			done, pending = await asyncio.wait([consumer_task, producer_task, self.stop], return_when=asyncio.FIRST_COMPLETED,)
			for task in pending:
				task.cancel()
		finally:
			print("handler: lost connection from {}".format(websocket.remote_address))
			await websocket.close()
			self.listener.on_disconnect()

	async def handler_recv(self, websocket, path):
		while True:
			try:
				data = await websocket.recv()
				await self.listener.on_recv(data)
			except:
				print("handler_recv: lost connection from {}".format(websocket.remote_address))
				return

	async def handler_send(self, websocket, path):
		while True:
			data = await self.listener.on_send()
			if data is not None:
				try:
					await websocket.send(json.dumps(data))
				except:
					print("handler_send: lost connection from {}".format(websocket.remote_address))
					break


class PeriodicUpdater(IConnection):
	def __init__(self):
		self.ws_client = WsClient(self)

		self.state = Queue()
		self.thread_cancel = True
		self.thread = None

	def connect(self, uri):
		self.ws_client.connect(uri)

	def disconnect(self):
		self.ws_client.disconnect()

	def on_connect(self):
		self.thread_cancel = False
		self.thread = Thread(target=self.run)
		self.thread.setDaemon(True)
		self.thread.start()

	def on_disconnect(self):
		self.thread_cancel = True
		self.thread.join(1)

	async def on_send(self):
		return await asyncio.get_event_loop().run_in_executor(None, self.state.get)

	async def on_recv(self, data):
		print("R: {}".format(data))

	def send(self, data):
		self.state.put_nowait(data)

	def run(self):
		count = 0
		while not self.thread_cancel:
			self.send({"Value": count})
			count += 1
			sleep(0.5)


if __name__ == '__main__':

	ws = PeriodicUpdater()
	ws.connect(uri="ws://localhost:8765")

	sleep(5)
	ws.disconnect()
	sleep(1)
	ws.connect(uri="ws://localhost:8765")

	while True:
		sleep(1)