import asyncio
import websockets
import abc


class WsServer:
	def __init__(self, loop=None):
		if loop is None:
			self.loop = asyncio.get_event_loop()
		else:
			self.loop = loop

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

	def disconnect(self):
		pass

	async def handler(self, websocket, path):
		print("handler: got connection from {}".format(websocket.remote_address))
		await self.register(websocket)
		try:
			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], return_when=asyncio.FIRST_COMPLETED,)
			for task in pending:
				task.cancel()
		finally:
			print("handler: lost connection from {}".format(websocket.remote_address))
			await self.unregister(websocket)

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

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