"""
This submodule contains the WebSocket class, which handles the websocket connection to Lavalink.
"""
import asyncio
import logging
import sys
import traceback
from typing import TYPE_CHECKING, Any, Dict, Union

import aiohttp
from disnake.ext import commands

if TYPE_CHECKING:
    from wavelink import Node

from .backoff import ExponentialBackoff
from .events import *
from .stats import Stats
from .utils import MISSING

logger = logging.getLogger(__name__)


class WebSocket:
    """
    This class handles the websocket connection of a Lavalink node.
    """

    def __init__(self, **attrs):
        """
        Initializes a new websocket connection.
        """
        self._closed = None
        self._node: Node = attrs.get("node")
        self.client = self._node.client
        self.bot: commands.Bot = self.client.bot
        self.host = attrs.get("host")
        self.port = attrs.get("port")
        self.password = attrs.get("password")
        self.shard_count = attrs.get("shard_count")
        self.user_id = attrs.get("user_id")
        self.secure: bool = attrs.get("secure")
        self._dumps = attrs.get("dumps")

        self._websocket = MISSING
        self._last_exc = MISSING
        self._task = MISSING

    @property
    def headers(self):
        return {
            "Authorization": self.password,
            "Num-Shards": str(self.shard_count),
            "User-Id": str(self.user_id),
            "Client-Name": "wavelink",
            "Resume-Key": self._node.resume_key,
        }

    @property
    def is_connected(self) -> bool:
        return self._websocket is not MISSING and not self._websocket.closed

    @property
    def needs_to_be_reconnected(self) -> bool:
        return self._websocket is not MISSING and self._websocket.closed

    async def connect(self):
        """
        |coro|

        This method connects to the Lavalink websocket.
        """
        await self.bot.wait_until_ready()
        if self.is_connected:
            assert isinstance(self._websocket, aiohttp.ClientWebSocketResponse)
            await self._websocket.close(code=1006, message=b"WaveLink: Attempting reconnection.")

        try:
            if self.secure is True:
                uri = f"wss://{self.host}:{self.port}"
            else:
                uri = f"ws://{self.host}:{self.port}"

            if not self.is_connected:
                self._websocket = await self._node.session.ws_connect(
                    uri, headers=self.headers, heartbeat=self._node.heartbeat
                )

        except Exception as error:
            self._last_exc = error
            self._node.available = False

            if isinstance(error, aiohttp.WSServerHandshakeError) and error.status == 401:
                print(f"\nAuthorization Failed for Node:: {self._node}\n", file=sys.stderr)
            else:
                logger.error(f"WEBSOCKET | Connection Failure:: {error}")
                traceback.print_exception(type(error), error, error.__traceback__, file=sys.stderr)

        if not self._task:
            self._task = self.bot.loop.create_task(self.listen())

        self._last_exc = None
        self._closed = False
        self._node.available = True

        if self.is_connected:
            await self.client.dispatch_listeners("on_node_ready", self._node)
            logger.debug(f"Establishing Websocket Connection...{self._node.__repr__()}")

        if self.needs_to_be_reconnected:
            await self.client.dispatch_listeners("on_node_ready", self._node)
            logger.debug(f"Re-establishing Websocket Connection...{self._node.__repr__()}")
            resume = {
                "op": "configureResuming",
                "key": f"{self._node.resume_key}",
                "timeout": 60,
            }
            await self.send(**resume)

    async def listen(self):
        backoff = ExponentialBackoff(base=1, maximum_time=60, maximum_tries=None)

        while True:
            assert isinstance(self._websocket, aiohttp.ClientWebSocketResponse)
            msg = await self._websocket.receive()

            if msg.type is aiohttp.WSMsgType.CLOSED:
                logger.debug(f"WEBSOCKET | Close data: {msg.extra}")

                self._closed = True
                retry = backoff.calculate()

                logger.warning(f"\nWEBSOCKET | Connection closed:: Retrying connection in <{retry}> seconds\n")

                await asyncio.sleep(retry)
                if not self.is_connected:
                    self.bot.loop.create_task(self.connect())
            else:
                logger.debug(f"WEBSOCKET | Received Payload:: <{msg.data}>")
                if msg.data == 1011:
                    # Lavalink encountered an internal error which can not be fixed...
                    # Consider updating Lavalink...
                    logger.error(
                        "Internal Lavalink Error encountered. Terminating WaveLink without retries."
                        "Consider updating your Lavalink Server."
                    )

                    sys.exit(1)
                self.bot.loop.create_task(self.process_data(msg.json()))

    async def process_data(self, data: Dict[str, Any]):
        op = data.get("op", None)
        if not op:
            return

        if op == "stats":
            self._node.stats = Stats(data=data)
            return
        if op == "event":
            try:
                data["player"] = self._node.players[int(data["guildId"])]
            except KeyError:
                return

            listener, payload = self._get_event_payload(data["type"], data)

            logger.debug(f"WEBSOCKET | op: event:: {data}")

            # Dispatch node event/player hooks
            try:
                await self._node.on_event(payload)
            except Exception as e:
                traceback.print_exception(type(e), e, e.__traceback__, file=sys.stderr)

            # Dispatch listeners
            await self.client.dispatch_listeners(listener, self._node, payload)
            return

        elif op == "playerUpdate":
            logger.debug(f"WEBSOCKET | op: playerUpdate:: {data}")
            try:
                await self._node.players[int(data["guildId"])].update_state(data)
            except KeyError:
                pass
            return

    @staticmethod
    def _get_event_payload(
        name: str, data
    ) -> Union[
        tuple[str, TrackEnd],
        tuple[str, TrackStart],
        tuple[str, TrackException],
        tuple[str, TrackStuck],
        tuple[str, WebsocketClosed],
    ]:
        if name == "TrackEndEvent":
            return "on_track_end", TrackEnd(data)
        elif name == "TrackStartEvent":
            return "on_track_start", TrackStart(data)
        elif name == "TrackExceptionEvent":
            return "on_track_exception", TrackException(data)
        elif name == "TrackStuckEvent":
            return "on_track_stuck", TrackStuck(data)
        elif name == "WebSocketClosedEvent":
            return "on_websocket_closed", WebsocketClosed(data)

    async def send(self, **data):
        if self.is_connected:
            logger.debug(f"WEBSOCKET | Sending Payload:: {data}")
            data_str = self._dumps(data)
            if isinstance(data_str, bytes):
                # Some JSON libraries serialize to bytes
                # Yet Lavalink does not support binary websockets
                # So we need to decode. In the future, maybe
                # self._websocket.send_bytes could be used
                # if Lavalink ever implements it
                data_str = data_str.decode("utf-8")
            await self._websocket.send_str(data_str)
            return

"""
This submodule contains the WebSocket class, which handles the websocket connection to Lavalink.
"""
import asyncio
import logging
import sys
import traceback
from typing import TYPE_CHECKING, Any, Dict, Union

import aiohttp
from disnake.ext import commands

if TYPE_CHECKING:
    from wavelink import Node

from .backoff import ExponentialBackoff
from .events import *
from .stats import Stats
from .utils import MISSING

logger = logging.getLogger(__name__)


class WebSocket:
    """
    This class handles the websocket connection of a Lavalink node.
    """

    def __init__(self, **attrs):
        """
        Initializes a new websocket connection.
        """
        self._closed = None
        self._node: Node = attrs.get("node")
        self.client = self._node.client
        self.bot: commands.Bot = self.client.bot
        self.host = attrs.get("host")
        self.port = attrs.get("port")
        self.password = attrs.get("password")
        self.shard_count = attrs.get("shard_count")
        self.user_id = attrs.get("user_id")
        self.secure: bool = attrs.get("secure")
        self._dumps = attrs.get("dumps")

        self._websocket = MISSING
        self._last_exc = MISSING
        self._task = MISSING

    @property
    def headers(self):
        return {
            "Authorization": self.password,
            "Num-Shards": str(self.shard_count),
            "User-Id": str(self.user_id),
            "Client-Name": "wavelink",
            "Resume-Key": self._node.resume_key,
        }

    @property
    def is_connected(self) -> bool:
        return self._websocket is not MISSING and not self._websocket.closed

    @property
    def needs_to_be_reconnected(self) -> bool:
        return self._websocket is not MISSING and self._websocket.closed

    async def connect(self):
        """
        |coro|

        This method connects to the Lavalink websocket.
        """
        await self.bot.wait_until_ready()
        if self.is_connected:
            assert isinstance(self._websocket, aiohttp.ClientWebSocketResponse)
            await self._websocket.close(code=1006, message=b"WaveLink: Attempting reconnection.")

        try:
            if self.secure is True:
                uri = f"wss://{self.host}:{self.port}"
            else:
                uri = f"ws://{self.host}:{self.port}"

            if not self.is_connected:
                self._websocket = await self._node.session.ws_connect(
                    uri, headers=self.headers, heartbeat=self._node.heartbeat
                )

        except Exception as error:
            self._last_exc = error
            self._node.available = False

            if isinstance(error, aiohttp.WSServerHandshakeError) and error.status == 401:
                print(f"\nAuthorization Failed for Node:: {self._node}\n", file=sys.stderr)
            else:
                logger.error(f"WEBSOCKET | Connection Failure:: {error}")
                traceback.print_exception(type(error), error, error.__traceback__, file=sys.stderr)

        if not self._task:
            self._task = self.bot.loop.create_task(self.listen())

        self._last_exc = None
        self._closed = False
        self._node.available = True

        if self.is_connected:
            await self.client.dispatch_listeners("on_node_ready", self._node)
            logger.debug(f"Establishing Websocket Connection...{self._node.__repr__()}")

        if self.needs_to_be_reconnected:
            await self.client.dispatch_listeners("on_node_ready", self._node)
            logger.debug(f"Re-establishing Websocket Connection...{self._node.__repr__()}")
            resume = {
                "op": "configureResuming",
                "key": f"{self._node.resume_key}",
                "timeout": 60,
            }
            await self.send(**resume)

    async def listen(self):
        backoff = ExponentialBackoff(base=1, maximum_time=60, maximum_tries=None)

        while True:
            assert isinstance(self._websocket, aiohttp.ClientWebSocketResponse)
            msg = await self._websocket.receive()

            if msg.type is aiohttp.WSMsgType.CLOSED:
                logger.debug(f"WEBSOCKET | Close data: {msg.extra}")

                self._closed = True
                retry = backoff.calculate()

                logger.warning(f"\nWEBSOCKET | Connection closed:: Retrying connection in <{retry}> seconds\n")

                await asyncio.sleep(retry)
                if not self.is_connected:
                    self.bot.loop.create_task(self.connect())
            else:
                logger.debug(f"WEBSOCKET | Received Payload:: <{msg.data}>")
                if msg.data == 1011:
                    # Lavalink encountered an internal error which can not be fixed...
                    # Consider updating Lavalink...
                    logger.error(
                        "Internal Lavalink Error encountered. Terminating WaveLink without retries."
                        "Consider updating your Lavalink Server."
                    )

                    sys.exit(1)
                self.bot.loop.create_task(self.process_data(msg.json()))

    async def process_data(self, data: Dict[str, Any]):
        op = data.get("op", None)
        if not op:
            return

        if op == "stats":
            self._node.stats = Stats(data=data)
            return
        if op == "event":
            try:
                data["player"] = self._node.players[int(data["guildId"])]
            except KeyError:
                return

            listener, payload = self._get_event_payload(data["type"], data)

            logger.debug(f"WEBSOCKET | op: event:: {data}")

            # Dispatch node event/player hooks
            try:
                await self._node.on_event(payload)
            except Exception as e:
                traceback.print_exception(type(e), e, e.__traceback__, file=sys.stderr)

            # Dispatch listeners
            await self.client.dispatch_listeners(listener, self._node, payload)
            return

        elif op == "playerUpdate":
            logger.debug(f"WEBSOCKET | op: playerUpdate:: {data}")
            try:
                await self._node.players[int(data["guildId"])].update_state(data)
            except KeyError:
                pass
            return

    @staticmethod
    def _get_event_payload(
        name: str, data
    ) -> Union[
        tuple[str, TrackEnd],
        tuple[str, TrackStart],
        tuple[str, TrackException],
        tuple[str, TrackStuck],
        tuple[str, WebsocketClosed],
    ]:
        if name == "TrackEndEvent":
            return "on_track_end", TrackEnd(data)
        elif name == "TrackStartEvent":
            return "on_track_start", TrackStart(data)
        elif name == "TrackExceptionEvent":
            return "on_track_exception", TrackException(data)
        elif name == "TrackStuckEvent":
            return "on_track_stuck", TrackStuck(data)
        elif name == "WebSocketClosedEvent":
            return "on_websocket_closed", WebsocketClosed(data)

    async def send(self, **data):
        if self.is_connected:
            logger.debug(f"WEBSOCKET | Sending Payload:: {data}")
            data_str = self._dumps(data)
            if isinstance(data_str, bytes):
                # Some JSON libraries serialize to bytes
                # Yet Lavalink does not support binary websockets
                # So we need to decode. In the future, maybe
                # self._websocket.send_bytes could be used
                # if Lavalink ever implements it
                data_str = data_str.decode("utf-8")
            await self._websocket.send_str(data_str)
            return

This submodule contains the WebSocket class, which handles the websocket connection to Lavalink. This class handles the websocket connection of a Lavalink node. Initializes a new websocket connection. |coro|          This method connects to the Lavalink websocket. # Lavalink encountered an internal error which can not be fixed... # Consider updating Lavalink... # Dispatch node event/player hooks # Dispatch listeners # Some JSON libraries serialize to bytes # Yet Lavalink does not support binary websockets # So we need to decode. In the future, maybe # self._websocket.send_bytes could be used # if Lavalink ever implements it