import logging
import asyncio
import aiohttp

from ..log import LOG_NOTICE

#

L = logging.getLogger(__name__)

#


class WebSocketFactory(object):

	'''
	wsfactory = asab.web.WebSocketFactory(self)
	websvc.WebApp.router.add_get('/api/ws', wsfactory)
	wsfactory.on_message = ...


	Javascript example:

	<script type="text/javascript">
		var connection = new WebSocket(((window.location.protocol === "https:") ? "wss://" : "ws://") + window.location.host + "/api/ws");
		connection.onmessage = function (e) {
			...
		};
	</script>

	'''


	def __init__(self, app, *, timeout=10.0, protocols=(), compress=True, max_msg_size=4194304):
		self.App = app
		self.Counter = 0
		self.WebSockets = {}

		self.Timeout = timeout
		self.Protocols = protocols
		self.Compress = compress
		self.MaxMsgSize = max_msg_size
		self.__doc__ = "Provides WebSocket connection. Requires WebSocket upgrade."

		app.PubSub.subscribe("Application.stop!", self._on_app_stop)


	async def _on_app_stop(self, message_type, counter):
		# Clean up during application exit
		await self.close_all(code=aiohttp.WSCloseCode.GOING_AWAY, message='Server shutdown')


	async def close_all(self, *, code=aiohttp.WSCloseCode.OK, message=b''):
		'''
		Close all websockets
		'''
		await asyncio.gather(
			*[ws.close(code=code, message=message) for ws in self.WebSockets.values()],
			return_exceptions=True
		)


	async def ping_all(self):
		'''
		Ping all connected websockets
		'''
		await asyncio.gather(
			*[ws.ping() for ws in self.WebSockets.values()],
			return_exceptions=True
		)


	async def send_str_all(self, data, compress=None):
		'''
		Send string to all connected websockets
		'''
		await asyncio.gather(
			*[ws.send_str(data, compress=compress) for ws in self.WebSockets.values()],
			return_exceptions=True
		)

	async def send_json_all(self, data, compress=None):
		'''
		Send json to all connected websockets
		'''
		await asyncio.gather(
			*[ws.send_json(data, compress=compress) for ws in self.WebSockets.values()],
			return_exceptions=True
		)


	def get(self, wsid):
		'''
		Get a websocket using its wsid.

		IMPORTANT: This method can return None if the websocket has been already closed.
		'''
		return self.WebSockets.get(wsid)


	async def __call__(self, request):
		ws = aiohttp.web.WebSocketResponse(
			timeout=self.Timeout,
			protocols=self.Protocols,
			compress=self.Compress,
			max_msg_size=self.MaxMsgSize,
		)

		await ws.prepare(request)

		L.log(LOG_NOTICE, "Websocket connection accepted")

		wsid = await self.register(ws, request)

		try:
			self.WebSockets[wsid] = ws
			await self.on_connect(ws, wsid)

			async for msg in ws:
				if msg.type == aiohttp.WSMsgType.ERROR:
					break
				else:
					await self.on_message(ws, msg, wsid)

		except asyncio.CancelledError:
			await ws.close()

		finally:
			del self.WebSockets[wsid]
			L.log(LOG_NOTICE, "Websocket connection closed")
			await self.on_close(ws, wsid)

		return ws


	async def register(self, ws, request):
		'''
		Must return the unique identifier of the websocket `wsid`.
		The default implementation is the unique counter.
		You may overload this to provide custom unique identifier
		'''
		self.Counter += 1
		return self.Counter


	async def on_connect(self, websocket, wsid):
		'''
		Override this method to implement action on websocket connection
		'''
		pass


	async def on_message(self, websocket, message, wsid):
		'''
		Override this method to receive messages from client over the websocket
		'''
		pass


	async def on_close(self, websocket, wsid):
		'''
		Override this method to receive a notification that client closed the websocket connection
		'''
		pass
