"""WsRpc for Shelly."""

from __future__ import annotations

import asyncio
import base64
import hashlib
import logging
import secrets
import socket
from asyncio import Task, tasks
from collections.abc import Callable, Coroutine, Iterable
from dataclasses import dataclass
from enum import Enum, auto
from http import HTTPStatus
from typing import TYPE_CHECKING, Any

from aiohttp import (
    ClientSession,
    ClientWebSocketResponse,
    WSMessage,
    WSMsgType,
    client_exceptions,
)
from aiohttp.web import (
    Application,
    AppRunner,
    BaseRequest,
    TCPSite,
    WebSocketResponse,
    get,
)
from yarl import URL

from ..const import (
    DEFAULT_HTTP_PORT,
    NOTIFY_WS_CLOSED,
    WS_API_URL,
    WS_HEARTBEAT,
)
from ..exceptions import (
    ConnectionClosed,
    DeviceConnectionError,
    DeviceConnectionTimeoutError,
    InvalidAuthError,
    InvalidMessage,
    RpcCallError,
)
from ..json import json_bytes, json_loads

_LOGGER = logging.getLogger(__name__)

BUFFER_SIZE = 1024 * 64


class RPCSource(Enum):
    """RPC message source."""

    CLIENT = auto()
    SERVER = auto()


def _receive_json_or_raise(msg: WSMessage) -> dict[str, Any]:
    """Receive json or raise."""
    if msg.type is WSMsgType.TEXT:
        try:
            data: dict[str, Any] = json_loads(msg.data)
        except ValueError as err:
            raise InvalidMessage(f"Received invalid JSON: {msg.data}") from err
        return data

    if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):
        raise ConnectionClosed("Connection was closed.")

    if msg.type is WSMsgType.ERROR:
        raise InvalidMessage("Received message error")

    raise InvalidMessage(f"Received non-Text message: {msg.type}")


def hex_hash(message: str) -> str:
    """Get hex representation of sha256 hash of string."""
    return hashlib.sha256(message.encode("utf-8")).hexdigest()


def _current_task_cancelled() -> bool:
    """Return whether the currently running task has been cancelled."""
    return bool((current_task := asyncio.current_task()) and current_task.cancelled())


HA2 = hex_hash("dummy_method:dummy_uri")


@dataclass
class AuthData:
    """RPC Auth data class."""

    realm: str
    username: str
    password: str
    nonce: str = ""
    algorithm: str = "SHA-256"
    nc: int = 0

    def __post_init__(self) -> None:
        """Call after initialization."""
        self.ha1 = hex_hash(f"{self.username}:{self.realm}:{self.password}")

    def update_challenge(self, auth_challenge: dict[str, Any]) -> None:
        """Update authentication challenge data from server."""
        self.nonce = auth_challenge["nonce"]
        # The only algorithm currently supported is SHA-256
        if self.algorithm != auth_challenge["algorithm"]:
            raise InvalidAuthError(
                f"Unsupported auth algorithm: {auth_challenge['algorithm']}"
            )

        # Shelly WallDisplay sends nc as a string but it should be an integer
        self.nc = int(auth_challenge.get("nc", 1))

    def get_auth(self) -> dict[str, Any]:
        """Get auth for RPC calls with current nc value."""
        random_bytes = secrets.token_bytes(16)
        cnonce = base64.b64encode(random_bytes).decode("ascii")

        # Calculate response hash: SHA256(HA1:nonce:nc:cnonce:qop:HA2)
        hashed = hex_hash(f"{self.ha1}:{self.nonce}:{self.nc}:{cnonce}:auth:{HA2}")

        frame = {
            "realm": self.realm,
            "username": self.username,
            "nonce": self.nonce,
            "nc": self.nc,
            "cnonce": cnonce,
            "response": hashed,
            "algorithm": self.algorithm,
        }

        # Increment nc after each auth calculation
        self.nc += 1

        return frame


@dataclass
class SessionData:
    """SessionData (src/dst/auth) class."""

    src: str | None
    dst: str | None
    auth_data: AuthData | None


class RPCCall:
    """RPCCall class."""

    __slots__ = (
        "auth_data",
        "call_id",
        "dst",
        "method",
        "params",
        "resolve",
        "result",
        "src",
    )

    def __init__(
        self,
        call_id: int,
        method: str,
        params: dict[str, Any] | None,
        session: SessionData,
        resolve: asyncio.Future[dict[str, Any]],
    ) -> None:
        """Initialize RPC class."""
        self.auth_data = session.auth_data
        self.call_id = call_id
        self.params = params
        self.method = method
        self.src = session.src
        self.dst = session.dst
        self.resolve = resolve
        self.result: dict[str, Any] = {}

    def __repr__(self) -> str:
        """Return representation of the call."""
        return (
            "<RPCCall "
            f"method={self.method} "
            f"params={self.params} "
            f"call_id={self.call_id} "
            f"result={self.result}"
            ">"
        )

    def build_request_frame(self) -> dict[str, Any]:
        """Request frame."""
        msg: dict[str, Any] = {
            "id": self.call_id,
            "method": self.method,
            "src": self.src,
        }
        if self.params is not None:
            msg["params"] = self.params
        if self.dst is not None:
            msg["dst"] = self.dst
        # Calculate auth on-demand with current nc
        if self.auth_data is not None and self.auth_data.nonce:
            msg["auth"] = self.auth_data.get_auth()
        return msg


class WsBase:
    """Base class for WebSocket handlers."""

    def __init__(self) -> None:
        """Initialize WsBase class."""
        self._background_tasks: set[Task] = set()

    def _create_and_track_task(self, func: Coroutine) -> None:
        """Create and and hold strong reference to the task."""
        task = asyncio.create_task(func)
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)


class WsRPC(WsBase):
    """WsRPC class."""

    def __init__(
        self,
        ip_address: str,
        on_notification: Callable[[RPCSource, str, dict | None], None],
        port: int = DEFAULT_HTTP_PORT,
    ) -> None:
        """Initialize WsRPC class."""
        super().__init__()
        self._ip_address = ip_address
        self._port = port
        self._on_notification = on_notification
        self._rx_task: tasks.Task[None] | None = None
        self._client: ClientWebSocketResponse | None = None
        self._calls: dict[int, RPCCall] = {}
        self._call_id = 0
        self._session = SessionData(f"aios-{id(self)}", None, None)
        self._loop = asyncio.get_running_loop()

    @property
    def _next_id(self) -> int:
        self._call_id += 1
        return self._call_id

    async def connect(self, aiohttp_session: ClientSession) -> None:
        """Connect to device."""
        if self.connected:
            raise RuntimeError("Already connected")

        _LOGGER.debug("Trying to connect to device at %s", self._ip_address)
        try:
            self._client = await aiohttp_session.ws_connect(
                URL.build(
                    scheme="http", host=self._ip_address, port=self._port, path="/rpc"
                ),
                heartbeat=WS_HEARTBEAT,
            )
        except (
            client_exceptions.WSServerHandshakeError,
            client_exceptions.ClientError,
        ) as err:
            raise DeviceConnectionError(err) from err

        # Try to reduce the pressure on shelly device as it measures
        # ram in bytes and we measure ram in megabytes.
        sock: socket.socket = self._client.get_extra_info("socket")
        try:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, BUFFER_SIZE)
        except OSError as err:
            _LOGGER.warning(
                "%s:%s: Failed to set socket receive buffer size: %s",
                self._ip_address,
                self._port,
                err,
            )

        self._rx_task = asyncio.create_task(self._rx_msgs())

        _LOGGER.info("Connected to %s", self._ip_address)

    async def disconnect(self) -> None:
        """Disconnect all sessions."""
        if self._rx_task is not None:
            self._rx_task.cancel()
            try:
                await self._rx_task
            except asyncio.CancelledError:
                if _current_task_cancelled():
                    raise
            self._rx_task = None
        if self._client is None:
            return

        await self._client.close()

    def set_auth_data(self, realm: str, username: str, password: str) -> None:
        """Set authentication data and generate session auth."""
        self._session.auth_data = AuthData(realm, username, password)

    async def _handle_call(self, frame_id: str) -> None:
        if TYPE_CHECKING:
            assert self._client

        await self._send_json(
            {
                "id": frame_id,
                "src": self._session.src,
                "error": {"code": 500, "message": "Not Implemented"},
            }
        )

    def handle_frame(self, source: RPCSource, frame: dict[str, Any]) -> None:
        """Handle RPC frame."""
        if peer_src := frame.get("src"):
            if self._session.dst is not None and peer_src != self._session.dst:
                _LOGGER.warning(
                    "Remote src changed: %s -> %s", self._session.dst, peer_src
                )
            self._session.dst = peer_src

        frame_id = frame.get("id")

        if method := frame.get("method"):
            # peer is invoking a method
            params = frame.get("params")
            if frame_id:
                # and expects a response
                _LOGGER.debug("handle call for frame_id: %s", frame_id)
                self._create_and_track_task(self._handle_call(frame_id))
            else:
                # this is a notification
                _LOGGER.debug("Notification: %s %s", method, params)
                try:
                    self._on_notification(source, method, params)
                except Exception as err:
                    _LOGGER.exception(
                        "Error handling notification frame: %s", frame, exc_info=err
                    )

            return

        if frame_id:
            # looks like a response
            if (call := self._calls.pop(frame_id, None)) is None:
                _LOGGER.warning(
                    "Response from (%s:%s) for an unknown request id: %s: %s",
                    self._ip_address,
                    self._port,
                    frame_id,
                    method,
                )
                return

            if not call.resolve.cancelled():
                call.resolve.set_result(frame)

            return

        _LOGGER.warning("Invalid frame: %s", frame)

    async def _rx_msgs(self) -> None:
        if TYPE_CHECKING:
            assert self._client

        try:
            while True:
                try:
                    msg = await self._client.receive()
                    frame = _receive_json_or_raise(msg)
                    _LOGGER.debug(
                        "recv(%s:%s): %s", self._ip_address, self._port, frame
                    )
                except InvalidMessage as err:
                    _LOGGER.error(
                        "Invalid Message from host %s:%s: %s",
                        self._ip_address,
                        self._port,
                        err,
                    )
                except (ConnectionClosed, client_exceptions.ClientConnectionResetError):
                    _LOGGER.debug(
                        "Connection issue with device %s:%s",
                        self._ip_address,
                        self._port,
                    )
                    break
                except Exception:
                    _LOGGER.exception("Unexpected error while receiving message")
                    raise

                if self._client.closed:
                    break

                self.handle_frame(RPCSource.CLIENT, frame)
        finally:
            _LOGGER.debug(
                "Websocket client connection from %s:%s closed",
                self._ip_address,
                self._port,
            )

            for call_item in self._calls.values():
                if not call_item.resolve.done():
                    call_item.resolve.set_exception(DeviceConnectionError(call_item))
            self._calls.clear()

            client = self._client
            self._client = None
            self._on_notification(RPCSource.CLIENT, NOTIFY_WS_CLOSED, None)

            # Close last since the await can yield
            # to the event loop and we want to minimize
            # race conditions
            if not client.closed:
                await client.close()

    @property
    def connected(self) -> bool:
        """Return if we're currently connected."""
        return self._client is not None and not self._client.closed

    async def call(
        self,
        method: str,
        params: dict[str, Any] | None = None,
        timeout: float = 10.0,
    ) -> dict[str, Any]:
        """Websocket RPC call."""
        return (await self.calls([(method, params)], timeout))[0]

    def _raise_for_unrecoverable_errors(
        self, resp: dict[str, Any], allow_auth_retry: bool
    ) -> None:
        """Raise for unrecoverable errors."""
        try:
            error = resp["error"]
            code = error["code"]
            msg = error["message"]
        except KeyError as err:
            raise RpcCallError(0, f"bad response: {resp}") from err

        if code != HTTPStatus.UNAUTHORIZED.value:
            raise RpcCallError(code, msg)

        if allow_auth_retry and self._session.auth_data is not None:
            return

        raise InvalidAuthError(msg)

    async def _rpc_call_with_auth_retry(
        self,
        method: str,
        params: dict[str, Any] | None = None,
        timeout: float = 10.0,
        allow_auth_retry: bool = True,
    ) -> dict[str, Any]:
        """Websocket RPC call with authentication retry."""
        if self._client is None:
            raise RuntimeError("Not connected")

        try:
            async with asyncio.timeout(timeout):
                call_id = self._next_id
                call = RPCCall(
                    call_id,
                    method,
                    params,
                    self._session,
                    self._loop.create_future(),
                )
                self._calls[call_id] = call
                await self._send_json(call.build_request_frame())

                # Wait response
                response = await call.resolve
                if "result" not in response:
                    self._raise_for_unrecoverable_errors(response, allow_auth_retry)
                    auth_challenge = json_loads(response["error"]["message"])
                    if TYPE_CHECKING:
                        assert self._session.auth_data
                    self._session.auth_data.update_challenge(auth_challenge)
                    return await self._rpc_call_with_auth_retry(
                        method,
                        params,
                        timeout,
                        allow_auth_retry=False,
                    )
                call.result = response["result"]
        except TimeoutError as exc:
            call.resolve.cancel()
            try:
                await call.resolve
            except asyncio.CancelledError:
                if _current_task_cancelled():
                    raise
            # Ensure the call is removed from the calls dict
            # on failure
            self._calls.pop(call.call_id, None)
            raise DeviceConnectionTimeoutError(call) from exc

        if _LOGGER.isEnabledFor(logging.DEBUG):
            _LOGGER.debug(
                "result(%s:%s): %s(%s) -> %s",
                self._ip_address,
                self._port,
                call.method,
                call.params,
                call.result,
            )

        return call.result

    async def calls(
        self, calls: Iterable[tuple[str, dict[str, Any] | None]], timeout: float = 10.0
    ) -> list[dict[str, Any]]:
        """Websocket RPC calls."""
        if self._session.auth_data is not None:
            # sequential: each call uses fresh nc from prior challenge,
            # first call seeds the nonce via 401 challenge
            results = []
            deadline = self._loop.time() + timeout
            for method, params in calls:
                remaining = deadline - self._loop.time()
                results.append(
                    await self._rpc_call_with_auth_retry(method, params, remaining)
                )
            return results

        results = await self._rpc_calls(calls, timeout)
        return [call.result for call in results]

    async def _rpc_calls(
        self, rpc_calls: Iterable[tuple[str, dict[str, Any] | None]], timeout: float
    ) -> list[RPCCall]:
        """Websocket RPC call.

        calls is a tuple of tuples of
        (
           (method, params),
           ...

        """
        if self._client is None:
            raise RuntimeError("Not connected")

        sent_calls: list[RPCCall] = []
        loop = self._loop
        future: asyncio.Future[dict[str, Any]]

        try:
            async with asyncio.timeout(timeout):
                for method, params in rpc_calls:
                    call_id = self._next_id
                    future = loop.create_future()
                    call = RPCCall(call_id, method, params, self._session, future)
                    sent_calls.append(call)
                    self._calls[call_id] = call
                    await self._send_json(call.build_request_frame())

                # Wait for all the responses
                for call in sent_calls:
                    response = await call.resolve
                    if "result" not in response:
                        self._raise_for_unrecoverable_errors(
                            response, allow_auth_retry=False
                        )
                    else:
                        call.result = response["result"]
        except TimeoutError as exc:
            for call in sent_calls:
                call.resolve.cancel()
                try:
                    await call.resolve
                except asyncio.CancelledError:
                    if _current_task_cancelled():
                        raise
                # Ensure the call is removed from the calls dict
                # on failure
                self._calls.pop(call.call_id, None)
            raise DeviceConnectionTimeoutError(sent_calls) from exc

        if _LOGGER.isEnabledFor(logging.DEBUG):
            for call in sent_calls:
                _LOGGER.debug(
                    "result(%s:%s): %s(%s) -> %s",
                    self._ip_address,
                    self._port,
                    call.method,
                    call.params,
                    call.result,
                )

        return sent_calls

    async def _send_json(self, data: dict[str, Any]) -> None:
        """Send json frame to device."""
        _LOGGER.debug("send(%s:%s): %s", self._ip_address, self._port, data)

        if TYPE_CHECKING:
            assert self._client

        await self._client.send_frame(json_bytes(data), WSMsgType.TEXT)


class WsServer(WsBase):
    """WsServer class."""

    def __init__(self) -> None:
        """Initialize WsServer class."""
        super().__init__()
        self._runner: AppRunner | None = None
        self.subscriptions: dict[str, Callable] = {}

    async def initialize(self, port: int, api_url: str = WS_API_URL) -> None:
        """Initialize the websocket server, used only in standalone mode."""
        app = Application()
        app.add_routes([get(api_url, self.websocket_handler)])
        self._runner = AppRunner(app)
        await self._runner.setup()
        site = TCPSite(self._runner, port=port)
        await site.start()

    def close(self) -> None:
        """Stop the websocket server."""
        if self._runner is not None:
            self._create_and_track_task(self._runner.cleanup())

    async def websocket_handler(self, request: BaseRequest) -> WebSocketResponse:
        """Handle connections from sleeping devices."""
        ip = request.remote
        _LOGGER.debug("Websocket server connection from %s starting", ip)
        ws_res = WebSocketResponse(protocols=["json-rpc"])
        await ws_res.prepare(request)
        _LOGGER.debug("Websocket server connection from %s ready", ip)

        async for msg in ws_res:
            try:
                frame = _receive_json_or_raise(msg)
                _LOGGER.debug("recv(%s): %s", ip, frame)
            except ConnectionClosed:
                await ws_res.close()
            except InvalidMessage as err:
                _LOGGER.debug("Invalid Message from host %s: %s", ip, err)
            else:
                try:
                    device_id = frame["src"].split("-")[1].upper()
                except (KeyError, IndexError) as err:
                    _LOGGER.debug("Invalid device id from host %s: %s", ip, err)
                    continue

                if device_id in self.subscriptions:
                    _LOGGER.debug(
                        "Calling WsRPC message update for device id %s", device_id
                    )
                    self.subscriptions[device_id](frame)
                    continue

                if ip in self.subscriptions:
                    _LOGGER.debug("Calling WsRPC message update for host %s", ip)
                    self.subscriptions[ip](frame)

        _LOGGER.debug("Websocket server connection from %s closed", ip)
        return ws_res

    def subscribe_updates(self, ip: str, message_received: Callable) -> Callable:
        """Subscribe to received updates."""
        _LOGGER.debug("Adding device %s to WsServer message subscriptions", ip)
        self.subscriptions[ip] = message_received
        return lambda: self.subscriptions.pop(ip)