"""
Created by Epic at 9/5/20
"""

from asyncio import Event, Lock, AbstractEventLoop, sleep
from aiohttp.client_exceptions import ClientConnectorError
from aiohttp import WSMessage, WSMsgType
from logging import getLogger
from sys import platform
from ujson import loads, dumps
from time import time
from .exceptions import GatewayUnavailable


class DefaultShard:
	def __init__(self, shard_id, client, loop: AbstractEventLoop):
		self.id = shard_id
		self.client = client
		self.loop = loop

		self.ws = None
		self.gateway_url = None
		self.logger = getLogger(f"speedcord.shard.{self.id}")
		self.connected = Event(loop=self.loop)  # Some bots might wanna know which shards is online at all times

		self.received_heartbeat_ack = True
		self.heartbeat_interval = None
		self.heartbeat_count = None
		self.failed_heartbeats = 0
		self.session_id = None
		self.last_event_id = None  # This gets modified by gateway.py

		self.gateway_send_lock = Lock(loop=self.loop)
		self.gateway_send_limit = 120
		self.gateway_send_per = 60
		self.gateway_send_left = self.gateway_send_limit
		self.gateway_send_reset = time() + self.gateway_send_per

		# Default events
		self.client.opcode_dispatcher.register(10, self.handle_hello)
		self.client.opcode_dispatcher.register(11, self.handle_heartbeat_ack)
		self.client.opcode_dispatcher.register(9, self.handle_invalid_session)

		self.client.event_dispatcher.register("READY", self.handle_ready)

	async def connect(self, gateway_url):
		if self.ws is not None:
			if not self.ws.closed:
				await self.ws.close()
			self.ws = None
		self.gateway_url = gateway_url
		try:
			self.ws = await self.client.http.create_ws(gateway_url, compression=0)
		except ClientConnectorError:
			await self.client.close()
			raise GatewayUnavailable() from None
		self.loop.create_task(self.read_loop())
		self.connected.set()
		if self.session_id is None:
			async with self.client.connection_lock:
				self.client.remaining_connections -= 1
				if self.client.remaining_connections <= 1:
					self.logger.info("Max connections reached!")
					gateway_url, shard_count, _, connections_reset_after = await self.client.get_gateway()
					await sleep(connections_reset_after / 1000)
					gateway_url, shard_count, \
					self.client.remaining_connections, connections_reset_after = await self.client.get_gateway()
				await self.identify()
		else:
			await self.resume()

	async def close(self):
		if self.ws is not None and not self.ws.closed:
			await self.ws.close()
		self.connected.clear()

	async def read_loop(self):
		message: WSMessage  # Fix typehinting
		async for message in self.ws:
			if message.type == WSMsgType.TEXT:
				await self.client.gateway_handler.on_receive(message.json(loads=loads), self)
			elif message.type in [WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED]:
				self.logger.warning(
					f"WebSocket is closing! Details: {message.json()}. Close code: {self.ws.close_code}")
			else:
				self.logger.warning("Unknown message type: " + str(type(message)))

	async def send(self, data: dict):
		async with self.gateway_send_lock:
			current_time = time()
			if current_time >= self.gateway_send_reset:
				self.gateway_send_reset = current_time + self.gateway_send_per
				self.gateway_send_left = self.gateway_send_limit
			if self.gateway_send_left == 0:
				sleep_for = self.gateway_send_reset - current_time
				self.logger.debug(f"Gateway ratelimited! Sleeping for {sleep_for}s")
				await sleep(self.gateway_send_reset - current_time)
			self.logger.debug("Data sent: " + str(data))
			await self.ws.send_json(data, dumps=dumps)

	async def identify(self):
		await self.send({
			"op": 2,
			"d": {
				"token": self.client.token,
				"properties": {
					"$os": platform,
					"$browser": "SpeedCord",
					"$device": "SpeedCord"
				},
				"intents": self.client.intents,
				"shard": (self.id, self.client.shard_count)
			}
		})

	async def resume(self):
		await self.send({
			"op": 6,
			"d": {
				"token": self.client.token,
				"session_id": self.session_id,
				"seq": self.last_event_id
			}
		})

	async def heartbeat_loop(self):
		while self.connected.is_set():
			if not self.received_heartbeat_ack:
				self.failed_heartbeats += 1
				self.logger.info(
					"WebSocket did not respond to a heartbeat! Failed attempts: " + str(self.failed_heartbeats))
				if self.failed_heartbeats > 2:
					self.logger.warning("Gateway stopped responding, reconnecting!")
					await self.close()
					await self.connect(self.gateway_url)
					return
			self.received_heartbeat_ack = False
			await self.send({
				"op": 1,
				"d": self.heartbeat_count
			})
			if self.heartbeat_count is not None:
				self.heartbeat_count += 1
			else:
				self.heartbeat_count = 0
			await sleep(self.heartbeat_interval)

	async def handle_hello(self, data, shard):
		if shard.id != self.id:
			return
		self.received_heartbeat_ack = True
		self.heartbeat_interval = data["d"]["heartbeat_interval"] / 1000
		self.loop.create_task(self.heartbeat_loop())
		self.logger.debug("Started heartbeat loop")

	async def handle_heartbeat_ack(self, data, shard):
		if shard.id != self.id:
			return
		self.received_heartbeat_ack = True
		self.failed_heartbeats = 0

	async def handle_ready(self, data, shard):
		if shard.id != self.id:
			return
		self.session_id = data["session_id"]

	async def handle_invalid_session(self, data, shard):
		if shard.id != self.id:
			return
		if not data.get("d", False):
			# Session is no longer valid, create a new session
			self.session_id = None
		await self.close()
		await self.connect(self.gateway_url)
