import asyncio
import websockets.exceptions
from ws.server.ws_server import WsServer
from ws.user import User, UserSet, update
from ws.connection import IConnection


class WsServerMultiUser(WsServer):
	def __init__(self, loop=None, listener: IConnection = None):
		WsServer.__init__(self, loop)
		self.listener = listener
		self.USERS = UserSet()
		self.global_state = {}
		self.pump_task = None

	async def listen(self, host, port):
		# A single task drains the dispatcher's queue regardless of how many
		# clients are connected, so state updates never pile up in the queue
		# while nobody is listening. A reconnecting client resyncs from
		# global_state instead of replaying a backlog (see handler_recv).
		self.pump_task = asyncio.ensure_future(self.state_pump())
		return await super().listen(host, port)

	async def state_pump(self):
		while True:
			data = await self.listener.on_send()
			self.global_state = update(self.global_state, data)
			await self.notify_state(data)

	async def notify_state(self, data):
		if self.USERS:  # asyncio.wait doesn't accept an empty list
			await asyncio.gather(*[user.send(data) for user in self.USERS])

	async def notify_users(self):
		if self.USERS:  # asyncio.wait doesn't accept an empty list
			message = {"info": {"type": "users", "count": len(self.USERS)}}
			await asyncio.gather(*[user.send(message) for user in self.USERS])

	async def register(self, websocket):
		usr = User(websocket)
		usr.path_add("info")
		self.USERS.add(websocket, usr)
		await self.notify_users()

	async def unregister(self, websocket):
		self.USERS.remove(websocket)
		await self.notify_users()

	async def handler_recv(self, websocket, path):
		try:
			while True:
				data = await websocket.recv()
				usr = self.USERS.get(websocket)
				processed = await usr.process(data)
				if processed:
					await usr.send(self.global_state)
				else:
					await self.listener.on_recv(data)
		except websockets.exceptions.ConnectionClosed:
			pass