"""View to accept incoming websocket connection."""
from __future__ import annotations

import asyncio
from collections import deque
from collections.abc import Callable
import datetime as dt
import logging
from typing import TYPE_CHECKING, Any, Final

from aiohttp import WSMsgType, web

from homeassistant.components.http import HomeAssistantView
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.event import async_call_later
from homeassistant.util.json import json_loads

from .auth import AuthPhase, auth_required_message
from .const import (
    DATA_CONNECTIONS,
    MAX_PENDING_MSG,
    PENDING_MSG_PEAK,
    PENDING_MSG_PEAK_TIME,
    SIGNAL_WEBSOCKET_CONNECTED,
    SIGNAL_WEBSOCKET_DISCONNECTED,
    URL,
)
from .error import Disconnect
from .messages import message_to_json
from .util import describe_request

if TYPE_CHECKING:
    from .connection import ActiveConnection


_WS_LOGGER: Final = logging.getLogger(f"{__name__}.connection")


class WebsocketAPIView(HomeAssistantView):
    """View to serve a websockets endpoint."""

    name: str = "websocketapi"
    url: str = URL
    requires_auth: bool = False

    async def get(self, request: web.Request) -> web.WebSocketResponse:
        """Handle an incoming websocket connection."""
        return await WebSocketHandler(request.app["hass"], request).async_handle()


class WebSocketAdapter(logging.LoggerAdapter):
    """Add connection id to websocket messages."""

    def process(self, msg: str, kwargs: Any) -> tuple[str, Any]:
        """Add connid to websocket log messages."""
        assert self.extra is not None
        return f'[{self.extra["connid"]}] {msg}', kwargs


class WebSocketHandler:
    """Handle an active websocket client connection."""

    __slots__ = (
        "_hass",
        "_request",
        "_wsock",
        "_handle_task",
        "_writer_task",
        "_closing",
        "_authenticated",
        "_logger",
        "_peak_checker_unsub",
        "_connection",
        "_message_queue",
        "_ready_future",
    )

    def __init__(self, hass: HomeAssistant, request: web.Request) -> None:
        """Initialize an active connection."""
        self._hass = hass
        self._request: web.Request = request
        self._wsock = web.WebSocketResponse(heartbeat=55)
        self._handle_task: asyncio.Task | None = None
        self._writer_task: asyncio.Task | None = None
        self._closing: bool = False
        self._authenticated: bool = False
        self._logger = WebSocketAdapter(_WS_LOGGER, {"connid": id(self)})
        self._peak_checker_unsub: Callable[[], None] | None = None
        self._connection: ActiveConnection | None = None

        # The WebSocketHandler has a single consumer and path
        # to where messages are queued. This allows the implementation
        # to use a deque and an asyncio.Future to avoid the overhead of
        # an asyncio.Queue.
        self._message_queue: deque[str | None] = deque()
        self._ready_future: asyncio.Future[None] | None = None

    def __repr__(self) -> str:
        """Return the representation."""
        return (
            "<WebSocketHandler "
            f"closing={self._closing} "
            f"authenticated={self._authenticated} "
            f"description={self.description}>"
        )

    @property
    def description(self) -> str:
        """Return a description of the connection."""
        if connection := self._connection:
            return connection.get_description(self._request)
        if request := self._request:
            return describe_request(request)
        return "finished connection"

    async def _writer(self) -> None:
        """Write outgoing messages."""
        # Variables are set locally to avoid lookups in the loop
        message_queue = self._message_queue
        logger = self._logger
        wsock = self._wsock
        send_str = wsock.send_str
        loop = self._hass.loop
        debug = logger.debug
        is_enabled_for = logger.isEnabledFor
        logging_debug = logging.DEBUG
        # Exceptions if Socket disconnected or cancelled by connection handler
        try:
            while not wsock.closed:
                if (messages_remaining := len(message_queue)) == 0:
                    self._ready_future = loop.create_future()
                    await self._ready_future
                    messages_remaining = len(message_queue)

                # A None message is used to signal the end of the connection
                if (message := message_queue.popleft()) is None:
                    return

                debug_enabled = is_enabled_for(logging_debug)
                messages_remaining -= 1

                if (
                    not messages_remaining
                    or not (connection := self._connection)
                    or not connection.can_coalesce
                ):
                    if debug_enabled:
                        debug("%s: Sending %s", self.description, message)
                    await send_str(message)
                    continue

                messages: list[str] = [message]
                while messages_remaining:
                    # A None message is used to signal the end of the connection
                    if (message := message_queue.popleft()) is None:
                        return
                    messages.append(message)
                    messages_remaining -= 1

                joined_messages = ",".join(messages)
                coalesced_messages = f"[{joined_messages}]"
                if debug_enabled:
                    debug("%s: Sending %s", self.description, coalesced_messages)
                await send_str(coalesced_messages)
        except asyncio.CancelledError:
            debug("%s: Writer cancelled", self.description)
            raise
        except (RuntimeError, ConnectionResetError) as ex:
            debug("%s: Unexpected error in writer: %s", self.description, ex)
        finally:
            debug("%s: Writer done", self.description)
            # Clean up the peak checker when we shut down the writer
            self._cancel_peak_checker()

    @callback
    def _cancel_peak_checker(self) -> None:
        """Cancel the peak checker."""
        if self._peak_checker_unsub is not None:
            self._peak_checker_unsub()
            self._peak_checker_unsub = None

    @callback
    def _send_message(self, message: str | dict[str, Any]) -> None:
        """Send a message to the client.

        Closes connection if the client is not reading the messages.

        Async friendly.
        """
        if self._closing:
            # Connection is cancelled, don't flood logs about exceeding
            # max pending messages.
            return

        if isinstance(message, dict):
            message = message_to_json(message)

        message_queue = self._message_queue
        queue_size_before_add = len(message_queue)
        if queue_size_before_add >= MAX_PENDING_MSG:
            self._logger.error(
                (
                    "%s: Client unable to keep up with pending messages. Reached %s pending"
                    " messages. The system's load is too high or an integration is"
                    " misbehaving; Last message was: %s"
                ),
                self.description,
                MAX_PENDING_MSG,
                message,
            )
            self._cancel()
            return

        message_queue.append(message)
        ready_future = self._ready_future
        if ready_future and not ready_future.done():
            ready_future.set_result(None)

        peak_checker_active = self._peak_checker_unsub is not None

        if queue_size_before_add <= PENDING_MSG_PEAK:
            if peak_checker_active:
                self._cancel_peak_checker()
            return

        if not peak_checker_active:
            self._peak_checker_unsub = async_call_later(
                self._hass, PENDING_MSG_PEAK_TIME, self._check_write_peak
            )

    @callback
    def _check_write_peak(self, _utc_time: dt.datetime) -> None:
        """Check that we are no longer above the write peak."""
        self._peak_checker_unsub = None

        if len(self._message_queue) < PENDING_MSG_PEAK:
            return

        self._logger.error(
            (
                "%s: Client unable to keep up with pending messages. Stayed over %s for %s"
                " seconds. The system's load is too high or an integration is"
                " misbehaving; Last message was: %s"
            ),
            self.description,
            PENDING_MSG_PEAK,
            PENDING_MSG_PEAK_TIME,
            self._message_queue[-1],
        )
        self._cancel()

    @callback
    def _cancel(self) -> None:
        """Cancel the connection."""
        self._closing = True
        self._cancel_peak_checker()
        if self._handle_task is not None:
            self._handle_task.cancel()
        if self._writer_task is not None:
            self._writer_task.cancel()

    async def async_handle(self) -> web.WebSocketResponse:
        """Handle a websocket response."""
        request = self._request
        wsock = self._wsock
        logger = self._logger
        debug = logger.debug
        hass = self._hass
        is_enabled_for = logger.isEnabledFor
        logging_debug = logging.DEBUG

        try:
            async with asyncio.timeout(10):
                await wsock.prepare(request)
        except asyncio.TimeoutError:
            self._logger.warning("Timeout preparing request from %s", request.remote)
            return wsock

        debug("%s: Connected from %s", self.description, request.remote)
        self._handle_task = asyncio.current_task()

        @callback
        def handle_hass_stop(event: Event) -> None:
            """Cancel this connection."""
            self._cancel()

        unsub_stop = hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, handle_hass_stop)

        # As the webserver is now started before the start
        # event we do not want to block for websocket responses
        self._writer_task = asyncio.create_task(self._writer())

        auth = AuthPhase(logger, hass, self._send_message, self._cancel, request)
        connection = None
        disconnect_warn = None

        try:
            self._send_message(auth_required_message())

            # Auth Phase
            try:
                async with asyncio.timeout(10):
                    msg = await wsock.receive()
            except asyncio.TimeoutError as err:
                disconnect_warn = "Did not receive auth message within 10 seconds"
                raise Disconnect from err

            if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):
                raise Disconnect

            if msg.type != WSMsgType.TEXT:
                disconnect_warn = "Received non-Text message."
                raise Disconnect

            try:
                auth_msg_data = json_loads(msg.data)
            except ValueError as err:
                disconnect_warn = "Received invalid JSON."
                raise Disconnect from err

            if is_enabled_for(logging_debug):
                debug("%s: Received %s", self.description, auth_msg_data)
            connection = await auth.async_handle(auth_msg_data)
            self._connection = connection
            hass.data[DATA_CONNECTIONS] = hass.data.get(DATA_CONNECTIONS, 0) + 1
            async_dispatcher_send(hass, SIGNAL_WEBSOCKET_CONNECTED)

            self._authenticated = True
            #
            #
            # Our websocket implementation is backed by a deque
            #
            # As back-pressure builds, the queue will back up and use more memory
            # until we disconnect the client when the queue size reaches
            # MAX_PENDING_MSG. When we are generating a high volume of websocket messages,
            # we hit a bottleneck in aiohttp where it will wait for
            # the buffer to drain before sending the next message and messages
            # start backing up in the queue.
            #
            # https://github.com/aio-libs/aiohttp/issues/1367 added drains
            # to the websocket writer to handle malicious clients and network issues.
            # The drain causes multiple problems for us since the buffer cannot be
            # drained fast enough when we deliver a high volume or large messages:
            #
            # - We end up disconnecting the client. The client will then reconnect,
            # and the cycle repeats itself, which results in a significant amount of
            # CPU usage.
            #
            # - Messages latency increases because messages cannot be moved into
            # the TCP buffer because it is blocked waiting for the drain to happen because
            # of the low default limit of 16KiB. By increasing the limit, we instead
            # rely on the underlying TCP buffer and stack to deliver the messages which
            # can typically happen much faster.
            #
            # After the auth phase is completed, and we are not concerned about
            # the user being a malicious client, we set the limit to force a drain
            # to 1MiB. 1MiB is the maximum expected size of the serialized entity
            # registry, which is the largest message we usually send.
            #
            # https://github.com/aio-libs/aiohttp/commit/b3c80ee3f7d5d8f0b8bc27afe52e4d46621eaf99
            # added a way to set the limit, but there is no way to actually
            # reach the code to set the limit, so we have to set it directly.
            #
            wsock._writer._limit = 2**20  # type: ignore[union-attr] # pylint: disable=protected-access
            async_handle_str = connection.async_handle
            async_handle_binary = connection.async_handle_binary

            # Command phase
            while not wsock.closed:
                msg = await wsock.receive()

                if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):
                    break

                if msg.type == WSMsgType.BINARY:
                    if len(msg.data) < 1:
                        disconnect_warn = "Received invalid binary message."
                        break
                    handler = msg.data[0]
                    payload = msg.data[1:]
                    async_handle_binary(handler, payload)
                    continue

                if msg.type != WSMsgType.TEXT:
                    disconnect_warn = "Received non-Text message."
                    break

                try:
                    command_msg_data = json_loads(msg.data)
                except ValueError:
                    disconnect_warn = "Received invalid JSON."
                    break

                if is_enabled_for(logging_debug):
                    debug("%s: Received %s", self.description, command_msg_data)

                if not isinstance(command_msg_data, list):
                    async_handle_str(command_msg_data)
                    continue

                for split_msg in command_msg_data:
                    async_handle_str(split_msg)

        except asyncio.CancelledError:
            debug("%s: Connection cancelled", self.description)
            raise

        except Disconnect as ex:
            debug("%s: Connection closed by client: %s", self.description, ex)

        except Exception:  # pylint: disable=broad-except
            self._logger.exception(
                "%s: Unexpected error inside websocket API", self.description
            )

        finally:
            unsub_stop()

            self._cancel_peak_checker()

            if connection is not None:
                connection.async_handle_close()

            self._closing = True

            self._message_queue.append(None)
            if self._ready_future and not self._ready_future.done():
                self._ready_future.set_result(None)

            # If the writer gets canceled we still need to close the websocket
            # so we have another finally block to make sure we close the websocket
            # if the writer gets canceled.
            try:
                await self._writer_task
            finally:
                try:
                    # Make sure all error messages are written before closing
                    await wsock.close()
                finally:
                    if disconnect_warn is None:
                        debug("%s: Disconnected", self.description)
                    else:
                        self._logger.warning(
                            "%s: Disconnected: %s", self.description, disconnect_warn
                        )

                    if connection is not None:
                        hass.data[DATA_CONNECTIONS] -= 1
                        self._connection = None

                    async_dispatcher_send(hass, SIGNAL_WEBSOCKET_DISCONNECTED)

                    # Break reference cycles to make sure GC can happen sooner
                    self._wsock = None  # type: ignore[assignment]
                    self._request = None  # type: ignore[assignment]
                    self._hass = None  # type: ignore[assignment]
                    self._logger = None  # type: ignore[assignment]
                    self._message_queue = None  # type: ignore[assignment]
                    self._handle_task = None
                    self._writer_task = None
                    self._ready_future = None

        return wsock