import asyncio
import json
import logging
from typing import Any, Dict, Optional, Tuple, Union

from graphql import ExecutionResult, print_ast

from ..graphql_request import GraphQLRequest
from .common.adapters.websockets import WebSocketsAdapter
from .common.base import SubscriptionTransportBase
from .exceptions import (
    TransportConnectionFailed,
    TransportProtocolError,
    TransportQueryError,
    TransportServerError,
)

log = logging.getLogger(__name__)


class Subscription:
    """Records listener_id and unsubscribe query_id for a subscription."""

    def __init__(self, query_id: int) -> None:
        self.listener_id: int = query_id
        self.unsubscribe_id: Optional[int] = None


class PhoenixChannelWebsocketsTransport(SubscriptionTransportBase):
    """The PhoenixChannelWebsocketsTransport is an async transport
    which allows you to execute queries and subscriptions against an `Absinthe`_
    backend using the `Phoenix`_ framework `channels`_.

    .. _Absinthe: http://absinthe-graphql.org
    .. _Phoenix: https://www.phoenixframework.org
    .. _channels: https://hexdocs.pm/phoenix/Phoenix.Channel.html#content
    """

    def __init__(
        self,
        url: str,
        *,
        channel_name: str = "__absinthe__:control",
        heartbeat_interval: float = 30,
        ack_timeout: Optional[Union[int, float]] = 10,
        **kwargs: Any,
    ) -> None:
        """Initialize the transport with the given parameters.

        :param url: The server URL.'.
        :param channel_name: Channel on the server this transport will join.
            The default for Absinthe servers is "__absinthe__:control"
        :param heartbeat_interval: Interval in second between each heartbeat messages
            sent by the client
        :param ack_timeout: Timeout in seconds to wait for the reply message
            from the server.
        """
        self.channel_name: str = channel_name
        self.heartbeat_interval: float = heartbeat_interval
        self.heartbeat_task: Optional[asyncio.Future] = None
        self.subscriptions: Dict[str, Subscription] = {}
        self.ack_timeout: Optional[Union[int, float]] = ack_timeout

        # Instanciate a WebSocketAdapter to indicate the use
        # of the websockets dependency for this transport
        ws_adapter_args = {}
        for ws_arg in ["headers", "ssl", "connect_args"]:
            try:
                ws_adapter_args[ws_arg] = kwargs.pop(ws_arg)
            except KeyError:
                pass

        self.adapter: WebSocketsAdapter = WebSocketsAdapter(
            url=url,
            **ws_adapter_args,
        )

        # Initialize the generic SubscriptionTransportBase parent class
        super().__init__(
            adapter=self.adapter,
            **kwargs,
        )

    async def _initialize(self) -> None:
        """Join the specified channel and wait for the connection ACK.

        If the answer is not a connection_ack message, we will return an Exception.
        """

        query_id = self.next_query_id
        self.next_query_id += 1

        init_message = json.dumps(
            {
                "topic": self.channel_name,
                "event": "phx_join",
                "payload": {},
                "ref": query_id,
            }
        )

        await self._send(init_message)

        # Wait for the connection_ack message or raise a TimeoutError
        init_answer = await asyncio.wait_for(self._receive(), self.ack_timeout)

        answer_type, answer_id, execution_result = self._parse_answer(init_answer)

        if answer_type != "reply":
            raise TransportProtocolError(
                "Websocket server did not return a connection ack"
            )

        async def heartbeat_coro():
            while True:
                await asyncio.sleep(self.heartbeat_interval)
                try:
                    query_id = self.next_query_id
                    self.next_query_id += 1

                    await self._send(
                        json.dumps(
                            {
                                "topic": "phoenix",
                                "event": "heartbeat",
                                "payload": {},
                                "ref": query_id,
                            }
                        )
                    )
                except TransportConnectionFailed:  # pragma: no cover
                    return

        self.heartbeat_task = asyncio.ensure_future(heartbeat_coro())

    async def _send_stop_message(self, query_id: int) -> None:
        """Send an 'unsubscribe' message to the Phoenix Channel referencing
        the listener's query_id, saving the query_id of the message.

        The server should afterwards return a 'phx_reply' message with
        the same query_id and subscription_id of the 'unsubscribe' request.
        """
        subscription_id = self._find_existing_subscription(query_id)

        unsubscribe_query_id = self.next_query_id
        self.next_query_id += 1

        # Save the ref so it can be matched in the reply
        self.subscriptions[subscription_id].unsubscribe_id = unsubscribe_query_id
        unsubscribe_message = json.dumps(
            {
                "topic": self.channel_name,
                "event": "unsubscribe",
                "payload": {"subscriptionId": subscription_id},
                "ref": unsubscribe_query_id,
            }
        )

        await self._send(unsubscribe_message)

    async def _stop_listener(self, query_id: int) -> None:
        await self._send_stop_message(query_id)

    async def _send_connection_terminate_message(self) -> None:
        """Send a phx_leave message to disconnect from the provided channel."""

        query_id = self.next_query_id
        self.next_query_id += 1

        connection_terminate_message = json.dumps(
            {
                "topic": self.channel_name,
                "event": "phx_leave",
                "payload": {},
                "ref": query_id,
            }
        )

        await self._send(connection_terminate_message)

    async def _connection_terminate(self):
        await self._send_connection_terminate_message()

    async def _send_query(
        self,
        request: GraphQLRequest,
    ) -> int:
        """Send a query to the provided websocket connection.

        We use an incremented id to reference the query.

        Returns the used id for this query.
        """

        query_id = self.next_query_id
        self.next_query_id += 1

        query_str = json.dumps(
            {
                "topic": self.channel_name,
                "event": "doc",
                "payload": {
                    "query": print_ast(request.document),
                    "variables": request.variable_values or {},
                },
                "ref": query_id,
            }
        )

        await self._send(query_str)

        return query_id

    def _parse_answer(
        self, answer: str
    ) -> Tuple[str, Optional[int], Optional[ExecutionResult]]:
        """Parse the answer received from the server

        Returns a list consisting of:
            - the answer_type (between:
              'data', 'reply', 'complete', 'close')
            - the answer id (Integer) if received or None
            - an execution Result if the answer_type is 'data' or None
        """

        event: str = ""
        answer_id: Optional[int] = None
        answer_type: str = ""
        execution_result: Optional[ExecutionResult] = None
        subscription_id: Optional[str] = None

        def _get_value(d: Any, key: str, label: str) -> Any:
            if not isinstance(d, dict):
                raise ValueError(f"{label} is not a dict")

            return d.get(key)

        def _required_value(d: Any, key: str, label: str) -> Any:
            value = _get_value(d, key, label)
            if value is None:
                raise ValueError(f"null {key} in {label}")

            return value

        def _required_subscription_id(
            d: Any, label: str, must_exist: bool = False, must_not_exist: bool = False
        ) -> str:
            subscription_id = str(_required_value(d, "subscriptionId", label))
            if must_exist and (subscription_id not in self.subscriptions):
                raise ValueError("unregistered subscriptionId")
            if must_not_exist and (subscription_id in self.subscriptions):
                raise ValueError("previously registered subscriptionId")

            return subscription_id

        def _validate_data_response(d: Any, label: str) -> dict:
            """Make sure query, mutation or subscription answer conforms.
            The GraphQL spec says only three keys are permitted.
            """
            if not isinstance(d, dict):
                raise ValueError(f"{label} is not a dict")

            keys = set(d.keys())
            invalid = keys - {"data", "errors", "extensions"}
            if len(invalid) > 0:
                raise ValueError(
                    f"{label} contains invalid items: " + ", ".join(invalid)
                )
            return d

        try:
            json_answer = json.loads(answer)

            event = str(_required_value(json_answer, "event", "answer"))

            if event == "subscription:data":
                payload = _required_value(json_answer, "payload", "answer")

                subscription_id = _required_subscription_id(
                    payload, "payload", must_exist=True
                )

                result = _validate_data_response(payload.get("result"), "result")

                answer_type = "data"

                subscription = self.subscriptions[subscription_id]
                answer_id = subscription.listener_id

                execution_result = ExecutionResult(
                    data=result.get("data"),
                    errors=result.get("errors"),
                    extensions=result.get("extensions"),
                )

            elif event == "phx_reply":

                # Will generate a ValueError if 'ref' is not there
                # or if it is not an integer
                answer_id = int(_required_value(json_answer, "ref", "answer"))

                payload = _required_value(json_answer, "payload", "answer")

                status = _get_value(payload, "status", "payload")

                if status == "ok":
                    answer_type = "reply"

                    if answer_id in self.listeners:
                        response = _required_value(payload, "response", "payload")

                        if isinstance(response, dict) and "subscriptionId" in response:

                            # Subscription answer
                            subscription_id = _required_subscription_id(
                                response, "response", must_not_exist=True
                            )

                            self.subscriptions[subscription_id] = Subscription(
                                answer_id
                            )

                        else:
                            # Query or mutation answer
                            # GraphQL spec says only three keys are permitted
                            response = _validate_data_response(response, "response")

                            answer_type = "data"

                            execution_result = ExecutionResult(
                                data=response.get("data"),
                                errors=response.get("errors"),
                                extensions=response.get("extensions"),
                            )
                    else:
                        (
                            registered_subscription_id,
                            listener_id,
                        ) = self._find_subscription(answer_id)
                        if registered_subscription_id is not None:
                            # Unsubscription answer
                            response = _required_value(payload, "response", "payload")
                            subscription_id = _required_subscription_id(
                                response, "response"
                            )

                            if subscription_id != registered_subscription_id:
                                raise ValueError("subscription id does not match")

                            answer_type = "complete"

                            answer_id = listener_id

                elif status == "error":
                    response = payload.get("response")

                    if isinstance(response, dict):
                        if "errors" in response:
                            raise TransportQueryError(
                                str(response.get("errors")), query_id=answer_id
                            )
                        elif "reason" in response:
                            raise TransportQueryError(
                                str(response.get("reason")), query_id=answer_id
                            )
                    raise TransportQueryError("reply error", query_id=answer_id)

                elif status == "timeout":
                    raise TransportQueryError("reply timeout", query_id=answer_id)

                # In case of missing or unrecognized status, just continue

            elif event == "phx_error":
                # Sent if the channel has crashed
                # answer_id will be the "join_ref" for the channel
                # answer_id = int(json_answer.get("ref"))
                raise TransportServerError("Server error")
            elif event == "phx_close":
                answer_type = "close"
            else:
                raise ValueError("unrecognized event")

        except ValueError as e:
            log.error(f"Error parsing answer '{answer}': {e!r}")
            raise TransportProtocolError(
                f"Server did not return a GraphQL result: {e!s}"
            ) from e

        return answer_type, answer_id, execution_result

    async def _handle_answer(
        self,
        answer_type: str,
        answer_id: Optional[int],
        execution_result: Optional[ExecutionResult],
    ) -> None:
        if answer_type == "close":
            pass
        else:
            await super()._handle_answer(answer_type, answer_id, execution_result)

    def _remove_listener(self, query_id: int) -> None:
        """If the listener was a subscription, remove that information."""
        try:
            subscription_id = self._find_existing_subscription(query_id)
            del self.subscriptions[subscription_id]
        except Exception:
            pass
        super()._remove_listener(query_id)

    def _find_subscription(self, query_id: int) -> Tuple[Optional[str], int]:
        """Perform a reverse lookup to find the subscription id matching
        a listener's query_id.
        """
        for subscription_id, subscription in self.subscriptions.items():
            if query_id == subscription.listener_id:
                return subscription_id, query_id
            if query_id == subscription.unsubscribe_id:
                return subscription_id, subscription.listener_id
        return None, query_id

    def _find_existing_subscription(self, query_id: int) -> str:
        """Perform a reverse lookup to find the subscription id matching
        a listener's query_id.
        """
        subscription_id, _listener_id = self._find_subscription(query_id)

        if subscription_id is None:
            raise TransportProtocolError(
                f"No subscription registered for listener {query_id}"
            )
        return subscription_id

    async def _close_coro(self, e: Exception, clean_close: bool = True) -> None:
        if self.heartbeat_task is not None:
            self.heartbeat_task.cancel()

        await super()._close_coro(e, clean_close)
