import asyncio
import websockets
import abc


class WsServer:
	def __init__(self, loop=None):
		self.loop = loop

	def run_forever(self):
		if self.loop is None:
			asyncio.get_event_loop().run_forever()
		else:
			self.loop.run_forever()

	async def listen(self, host, port):
		return await websockets.serve(self.handler, host, port)

	def disconnect(self):
		pass

	async def handler(self, websocket, path):
		await self.register(websocket)
		try:
			consumer_task = asyncio.ensure_future(self.handler_recv(websocket, path))
			producer_task = asyncio.ensure_future(self.handler_send(websocket, path))
			done, pending = await asyncio.wait([consumer_task, producer_task], return_when=asyncio.FIRST_COMPLETED,)
			for task in pending:
				task.cancel()
		finally:
			await self.unregister(websocket)

	@abc.abstractmethod
	async def handler_recv(self, websocket, path):
		pass

	@abc.abstractmethod
	async def handler_send(self, websocket, path):
		pass