from __future__ import annotations

import logging
from contextlib import aclosing, asynccontextmanager
from functools import partial
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncContextManager,
    AsyncIterable,
    Awaitable,
    Callable,
    Optional,
    Type,
    TypeVar,
    Union,
    cast,
)

import anyio
from anyio import CancelScope, ClosedResourceError, EndOfStream, Event, Lock
from anyio.abc import TaskStatus
from anyio.lowlevel import checkpoint

from serena.enums import ExchangeType
from serena.exc import (
    AMQPStateError,
    InvalidPayloadTypeError,
    MessageReturnedError,
    UnexpectedCloseError,
)
from serena.frame import BodyFrame, Frame
from serena.message import AMQPEnvelope, AMQPMessage
from serena.mixin import ChannelLike
from serena.payloads.header import BasicHeader, ContentHeaderFrame, ContentHeaderPayload
from serena.payloads.method import (
    BasicAckPayload,
    BasicCancelPayload,
    BasicConsumeOkPayload,
    BasicConsumePayload,
    BasicDeliverPayload,
    BasicGetEmptyPayload,
    BasicGetOkPayload,
    BasicGetPayload,
    BasicNackPayload,
    BasicPublishPayload,
    BasicRejectPayload,
    BasicReturnPayload,
    ChannelClosePayload,
    ChannelOpenOkPayload,
    ExchangeBindOkPayload,
    ExchangeBindPayload,
    ExchangeDeclareOkPayload,
    ExchangeDeclarePayload,
    ExchangeDeleteOkPayload,
    ExchangeDeletePayload,
    ExchangeUnBindOkPayload,
    ExchangeUnBindPayload,
    MethodFrame,
    MethodPayload,
    QueueBindOkPayload,
    QueueBindPayload,
    QueueDeclareOkPayload,
    QueueDeclarePayload,
    QueueDeleteOkPayload,
    QueueDeletePayload,
    QueuePurgeOkPayload,
    QueuePurgePayload,
    QueueUnbindOkPayload,
    QueueUnbindPayload,
)

if TYPE_CHECKING:
    from serena.connection import AMQPConnection

_PAYLOAD = TypeVar("_PAYLOAD", bound=MethodPayload)

logger = logging.getLogger(__name__)


# noinspection PyProtectedMember
class Channel(ChannelLike):
    """
    A wrapper around an AMQP channel.
    """

    def __init__(self, channel_id: int, connection: AMQPConnection, stream_buffer_size: int):
        """
        :param channel_id: The ID of this channel.
        :param connection: The AMQP connection object to send data on.
        :param stream_buffer_size: The buffer size for the streams.
        """

        self._connection = connection
        self._channel_id = channel_id

        # Why is there both an `_open` and a `_closed`?
        # `_open` is True between ChannelOpenOk, and ChannelCloseOk, and is mostly just used for
        #  internal bookkeeping.
        # `_closed` is False when the ``async with`` block exits.

        self._open = False
        self._closed = False
        self._close_event = Event()

        # no buffer as these are for events that should return immediately
        self._send, self._receive = anyio.create_memory_object_stream(0)

        self._delivery_send, self._delivery_receive = anyio.create_memory_object_stream(
            max_buffer_size=stream_buffer_size
        )

        self._close_info: Optional[ChannelClosePayload] = None
        self._lock = Lock()

        self._is_consuming = False

        # internal state used by the connection
        # server requested a flow stop
        self._server_flow_stopped = False

        # used to count acks
        self._message_counter = 0

    def __str__(self):
        return f"<Channel id={self.id} buffered={self.current_buffer_size}>"

    __repr__ = __str__

    @property
    def id(self) -> int:
        """
        Returns the ID of this channel.
        """

        return self._channel_id

    @property
    def open(self) -> bool:
        """
        Returns if this channel is open or not.
        """

        return not self._closed

    @property
    def max_buffer_size(self) -> int:
        """
        Returns the maximum number of frames buffered in this channel. Used internally.
        """

        return int(self._delivery_send.statistics().max_buffer_size)

    @property
    def current_buffer_size(self) -> int:
        """
        Returns the current number of frames buffered in this channel. Used internally.
        """

        return self._delivery_send.statistics().current_buffer_used

    def _check_closed(self):
        """
        Checks if the channel is closed.
        """

        if self._closed:
            # todo: switch to our own exception?
            raise ClosedResourceError("This channel is closed")

    async def _close(self, payload: ChannelClosePayload):
        """
        Closes this channel.
        """

        self._close_info = payload
        await self._send.aclose()
        await self._delivery_send.aclose()
        self._closed = True

        # noinspection PyAsyncCall
        self._close_event.set()

        # aclose doesn't seem to checkpoint...
        await checkpoint()

    def _enqueue_regular(self, frame: MethodFrame):
        """
        Enqueues a regular method frame.
        """

        self._send.send_nowait(frame)

    async def _enqueue_delivery(self, frame: Frame):
        """
        Enqueues a delivery frame.
        """

        await self._delivery_send.send(frame)

    async def _receive_delivery_message(self):
        """
        Receives a single delivery message. This will reassemble a full message into its constituent
        frames.
        """

        method = None
        headers = None
        body = b""

        while True:
            # check for successful reassembly first
            if headers is not None and len(body) >= headers.payload.full_size:
                await checkpoint()
                # hehe payload.payload
                return AMQPMessage(
                    channel=self,
                    envelope=AMQPEnvelope.of(method.payload),
                    header=headers.payload.payload,
                    body=body,
                )

            try:
                next_frame = await self._delivery_receive.receive()
            except EndOfStream:
                if self._close_info is None:
                    raise AMQPStateError("Channel was closed improperly")

                raise UnexpectedCloseError.of(self._close_info) from None

            if method is None:
                if not isinstance(next_frame, MethodFrame):
                    raise AMQPStateError(f"Expected a method frame, got {next_frame} instead")

                if isinstance(next_frame.payload, (BasicGetOkPayload, BasicDeliverPayload)):
                    method = next_frame

                elif isinstance(next_frame.payload, BasicGetEmptyPayload):
                    return None

                else:
                    raise AMQPStateError(
                        f"Expected Basic.Deliver or Basic.Get-Ok, got {next_frame.payload} instead"
                    )

            elif headers is None:
                if not isinstance(next_frame, ContentHeaderFrame):
                    raise AMQPStateError(f"Expected a header frame, got {next_frame} instead")

                # explicit type hint as pycharm incorrectly infers based on the previous if check
                payload: ContentHeaderPayload = next_frame.payload  # type: ignore
                if payload.class_id != method.payload.klass:
                    raise AMQPStateError(
                        f"Class mismatch ({payload.class_id} != {method.payload.klass})"
                    )

                headers = next_frame

            else:
                if not isinstance(next_frame, BodyFrame):
                    raise AMQPStateError(f"Expected a body frame, got {next_frame} instead")

                body += next_frame.data

    async def _receive_frame(self) -> MethodFrame:
        """
        Receives a single frame from the channel.
        """

        self._check_closed()

        try:
            frame = await self._receive.receive()
        except EndOfStream:
            if self._close_info is None:
                raise AMQPStateError("Channel was closed improperly")

            raise UnexpectedCloseError.of(self._close_info) from None

        # notify connection code that we received an event so it can unblock the channel
        # noinspection PyAsyncCall
        return frame

    async def _wait_until_open(self):
        """
        Waits until the channel is open.
        """

        frame = await self._receive.receive()
        if not isinstance(frame, MethodFrame):
            raise AMQPStateError(f"Expected MethodFrame, got {frame}")

        if not isinstance(frame.payload, ChannelOpenOkPayload):
            raise InvalidPayloadTypeError(ChannelOpenOkPayload, frame.payload)

        self._open = True

    async def _send_single_frame(self, payload: MethodPayload):
        """
        Sends a single frame to the connection. This won't wait for a response.
        """

        async with self._lock:
            await self._connection._send_method_frame(self._channel_id, payload)

    async def _send_and_receive(
        self, callable: Callable[[], Awaitable[None]], type: Type[_PAYLOAD] = None
    ) -> MethodFrame[_PAYLOAD]:
        """
        Sends and receives a payload.

        :param callable: A callable to do the actual sending.
        :param type: The type of the expected payload to return.
        :return: A :class:`.MethodFrame` that was returned as a result.
        """

        returned_frame: MethodFrame[_PAYLOAD] = None  # type: ignore

        async def _closure(task_status: TaskStatus):
            task_status.started()

            nonlocal returned_frame
            returned_frame = await self._receive_frame()

        async with anyio.create_task_group() as nursery:
            await nursery.start(_closure)

            await callable()

        if type is not None and not isinstance(returned_frame.payload, type):
            raise InvalidPayloadTypeError(type, returned_frame.payload)

        return cast(MethodFrame[type], returned_frame)

    async def _send_and_receive_frame(
        self, payload: MethodPayload, type: Type[_PAYLOAD] = None
    ) -> MethodFrame[_PAYLOAD]:
        """
        Sends and receives a method payload.

        :param payload: The :class:`.MethodPayload` to send.
        :param type: The type of the expected payload to return.
        :return: A :class:`.MethodFrame` that was returned as a result.
        """

        async with self._lock:
            fn = partial(self._connection._send_method_frame, self._channel_id, payload)
            return await self._send_and_receive(fn, type=type)

    async def wait_until_closed(self):
        """
        Waits until the channel is closed.
        """

        await self._close_event.wait()

    ## METHODS ##
    async def exchange_declare(
        self,
        name: str,
        type: Union[ExchangeType, str],
        *,
        passive: bool = False,
        durable: bool = False,
        auto_delete: bool = False,
        internal: bool = False,
        **arguments: Any,
    ) -> str:
        """
        Declares a new exchange.

        :param name: The name of the exchange. Must not be empty.
        :param type: The type of the exchange to create.
        :param passive: If True, the server will return a DeclareOk if the exchange exists, and
                        an error if it doesn't. This can be used to inspect server state without
                        modification.
        :param durable: If True, then the declared exchange will survive a server restart.
        :param auto_delete: If True, then the declared exchange will be automatically deleted
                            when all queues have finished using it.
        :param internal: If True, then the exchange may not be used directly by publishers.
        :param arguments: A dictionary of implementation-specific arguments.
        :return: The name of the exchange, as it exists on the server.
        """

        if isinstance(type, ExchangeType):
            type = type.value

        payload = ExchangeDeclarePayload(
            reserved_1=0,
            name=name,
            exchange_type=type,
            passive=passive,
            durable=durable,
            auto_delete=auto_delete,
            internal=internal,
            no_wait=False,
            arguments=arguments,
        )

        await self._send_and_receive_frame(payload, ExchangeDeclareOkPayload)
        return payload.name

    async def exchange_delete(
        self,
        name: str,
        *,
        if_unused: bool = False,
    ) -> None:
        """
        Deletes an exchange.

        :param name: The name of the exchange to delete.
        :param if_unused: If True, then the exchange will only be deleted if it has no queue
                          bindings.
        :return: Nothing.
        """

        payload = ExchangeDeletePayload(
            reserved_1=0,
            name=name,
            if_unused=if_unused,
            no_wait=False,
        )

        await self._send_and_receive_frame(payload, ExchangeDeleteOkPayload)

    async def exchange_bind(
        self,
        destination: str,
        source: str,
        routing_key: str,
        **arguments: Any,
    ) -> None:
        """
        Binds an exchange to another exchange. This is a
        `RabbitMQ extension <https://www.rabbitmq.com/e2e.html>`__ and may not be supported in other
        AMQP implementations.

        :param destination: The name of the destination exchange to bind. A blank name means the
                            default exchange.
        :param source: The name of the source exchange to bind. A blank name means the default
                       exchange.
        :param routing_key: The routing key for the exchange binding.
        :param arguments: A dictionary of implementation-specific arguments.
        :return: Nothing.
        """

        payload = ExchangeBindPayload(
            reserved_1=0,
            destination_name=destination,
            source_name=source,
            routing_key=routing_key,
            no_wait=False,
            arguments=arguments,
        )

        await self._send_and_receive_frame(payload, ExchangeBindOkPayload)

    async def exchange_unbind(
        self,
        destination: str,
        source: str,
        routing_key: str,
        **arguments: Any,
    ) -> None:
        """
        Unbinds an exchange from another exchange. This is a
        `RabbitMQ extension <https://www.rabbitmq.com/e2e.html>`__ and may not be supported in other
        AMQP implementations.

        :param destination: The name of the destination exchange to unbind. A blank name means the
                            default exchange.
        :param source: The name of the source exchange to unbind. A blank name means the default
                       exchange.
        :param routing_key: The routing key for the exchange binding that is being unbinded.
        :param arguments: A dictionary of implementation-specific arguments.
        :return: Nothing.
        """

        payload = ExchangeUnBindPayload(
            reserved_1=0,
            destination_name=destination,
            source_name=source,
            routing_key=routing_key,
            no_wait=False,
            arguments=arguments,
        )

        await self._send_and_receive_frame(payload, ExchangeUnBindOkPayload)

    async def queue_declare(
        self,
        name: str = "",
        *,
        passive: bool = False,
        durable: bool = False,
        exclusive: bool = False,
        auto_delete: bool = False,
        **arguments: Any,
    ) -> QueueDeclareOkPayload:
        """
        Declares a queue.

        :param name: The name of the queue. If blank, a name will be automatically generated by
                     the server and returned.
        :param passive: If True, the server will return a DeclareOk if the queue exists, and an
                        error if it doesn't. This can be used to inspect server state without
                        modification.
        :param durable: If True, the queue being created will persist past server restarts.
        :param exclusive: If True, this queue will only belong to this connection, and will be
                          automatically deleted when the connection closes. Best combined with an
                          automatically generated queue name.
        :param auto_delete: If True, this queue will be automatically deleted after all consumers
                            have finished. The queue will never be deleted before the first consumer
                            starts.
        :param arguments: Optional server implementation-specific arguments.
        :return: The :class:`.QueueDeclareOkPayload` the server returned.
        """

        self._check_closed()

        payload = QueueDeclarePayload(
            reserved_1=0,
            name=name or "",
            passive=passive,
            durable=durable,
            exclusive=exclusive,
            auto_delete=auto_delete,
            no_wait=False,
            arguments=arguments or {},
        )

        result = await self._send_and_receive_frame(payload, QueueDeclareOkPayload)
        return result.payload

    async def queue_bind(
        self,
        queue_name: str,
        exchange_name: str,
        routing_key: str,
        **arguments: Any,
    ) -> None:
        """
        Binds a queue to an exchange.

        :param queue_name: The queue to bind.
        :param exchange_name: The exchange to bind to.
        :param routing_key: The routing key to use when binding.
        :param arguments: Any server-specific or exchange-specific extra arguments.
        :return: Nothing.
        """

        payload = QueueBindPayload(
            reserved_1=0,
            queue_name=queue_name,
            exchange_name=exchange_name,
            routing_key=routing_key,
            no_wait=False,
            arguments=arguments,
        )

        await self._send_and_receive_frame(payload, QueueBindOkPayload)

    async def queue_delete(
        self,
        queue_name: str,
        *,
        if_empty: bool = False,
        if_unused: bool = False,
    ) -> int:
        """
        Deletes a queue.

        :param queue_name: The name of the queue to delete.
        :param if_empty: If True, the queue will only be deleted if it is empty.
        :param if_unused: If True, the queue will only be deleted if it is unused.
        :return: The number of messages deleted.
        """

        payload = QueueDeletePayload(
            reserved_1=0,
            queue_name=queue_name,
            if_empty=if_empty,
            if_unused=if_unused,
            no_wait=False,
        )

        result = await self._send_and_receive_frame(payload, QueueDeleteOkPayload)
        return result.payload.message_count

    async def queue_purge(
        self,
        queue_name: str,
    ) -> int:
        """
        Purges all messages from a queue.

        :param queue_name: The name of the queue to be purged.
        :return: The number of messages deleted.
        """

        payload = QueuePurgePayload(
            reserved_1=0,
            name=queue_name,
            no_wait=False,
        )

        result = await self._send_and_receive_frame(payload, QueuePurgeOkPayload)
        return result.payload.message_count

    async def queue_unbind(
        self,
        queue_name: str,
        exchange_name: str,
        routing_key: str,
        **arguments: Any,
    ) -> None:
        """
        Unbinds a queue from an exchange.

        :param queue_name: The name of the queue to unbind.
        :param exchange_name: The name of the exchange to unbind from.
        :param routing_key: The routing key to unbind using.
        :param arguments: Implementation-specific arguments to use.
        """

        payload = QueueUnbindPayload(
            reserved_1=0,
            queue_name=queue_name,
            exchange_name=exchange_name,
            routing_key=routing_key,
            no_wait=False,
            arguments=arguments,
        )

        await self._send_and_receive_frame(payload, QueueUnbindOkPayload)

    def basic_consume(
        self,
        queue_name: str,
        consumer_tag: str = "",
        *,
        no_local: bool = False,
        no_ack: bool = False,
        exclusive: bool = False,
        auto_ack: bool = True,
        **arguments: Any,
    ) -> AsyncContextManager[AsyncIterable[AMQPMessage]]:
        """
        Starts a basic consume operation. This returns an async context manager over an asynchronous
        iterator that yields incoming :class:`.AMQPMessage` instances.

        The channel can still be used for other operations during this operation.

        :param queue_name: The name of the queue to consume from.
        :param consumer_tag: The tag for this consume.
        :param no_local: If True, messages will not be sent to this consumer if it is on the same
                         connection that published them.
        :param no_ack: If True, messages will not be expected to be acknowledged. This can cause
                       data loss.
        :param exclusive: If True, then only this consumer can access the queue. Will fail if there
                          is another consumer already active.
        :param arguments: Implementation-specific arguments.
        :param auto_ack: If True, then messages will be automatically positively acknowledged
                         in the generator loop. Has no effect if ``no_ack`` is True. This is a
                         Serena-exclusive feature, not a protocol feature.
        """

        if self._is_consuming:
            raise AMQPStateError("Cannot start two consumers on the same channel")

        async def _agen():
            while True:
                message = await self._receive_delivery_message()
                yield message

                if auto_ack:
                    await self.basic_ack(message.envelope.delivery_tag)

        @asynccontextmanager
        async def _do():
            payload = BasicConsumePayload(
                reserved_1=0,
                queue_name=queue_name,
                consumer_tag=consumer_tag,
                no_local=no_local,
                no_ack=no_ack,
                exclusive=exclusive,
                no_wait=False,
                arguments=arguments or {},
            )

            response = await self._send_and_receive_frame(payload, BasicConsumeOkPayload)

            try:
                self._is_consuming = True
                async with aclosing(_agen()) as agen:
                    yield agen
            finally:
                cancel_payload = BasicCancelPayload(
                    consumer_tag=response.payload.consumer_tag,
                    no_wait=False,
                )

                with CancelScope(shield=True):
                    if not self._closed:
                        await self._send_and_receive_frame(cancel_payload)

                    self._is_consuming = False

        return _do()

    async def basic_publish(
        self,
        exchange_name: str,
        routing_key: str,
        body: bytes,
        *,
        header: BasicHeader = None,
        mandatory: bool = True,
        immediate: bool = False,
    ):
        """
        Publishes a message to a specific exchange.

        :param exchange_name: The name of the exchange to publish to. This can be blank to mean the
                              default exchange.
        :param routing_key: The routing key to publish to.
        :param body: The body for this payload.
        :param header: The headers to use for this message. If unset, will use the default blank
                       headers.
        :param mandatory: Iff True, the server must return a ``Return`` message if the message
                          could not be routed to a queue.
        :param immediate: Iff True, the server must return a ``Return`` message if the message could
                          not be immediately consumed.
        :raise MessageReturnedError: If the message was returned to the publisher.

        .. warning::

            The immediate flag is *not* supported in RabbitMQ 3.x, and will cause the connection
            to close.
        """

        # we have to send three manual frames before we get an ACK

        self._check_closed()

        async with self._lock:
            method_payload = BasicPublishPayload(
                reserved_1=1,
                name=exchange_name,
                routing_key=routing_key,
                mandatory=mandatory,
                immediate=immediate,
            )

            # 1) method frame
            await self._connection._send_method_frame(self._channel_id, method_payload)

            # 2) header frame
            headers = header or BasicHeader()
            if body:
                # Body is not empty, we send_and_receive with the body and not the headers.
                await self._connection._send_header_frame(
                    self._channel_id,
                    method_klass=method_payload.klass,
                    body_length=len(body),
                    headers=headers,
                )

                # 3) Body
                fn = partial(self._connection._send_body_frames, self._channel_id, body)
                response = await self._send_and_receive(fn)

            else:
                # Body is empty, we send_and_receive with the headers.
                fn = partial(
                    self._connection._send_header_frame,
                    self._channel_id,
                    method_klass=method_payload.klass,
                    body_length=len(body),
                    headers=headers,
                )

                response = await self._send_and_receive(fn)

            self._message_counter += 1

            # 4) check for Ack or Return
            payload = response.payload

            if isinstance(payload, BasicAckPayload):
                if payload.delivery_tag != self._message_counter:
                    raise AMQPStateError(
                        f"Expected Ack for delivery tag {self._message_counter}, "
                        f"but got Ack for delivery tag {payload.delivery_tag}"
                    )
                logger.trace(f"C#{self.id}: Server ACKed published message")

                return

            elif isinstance(payload, BasicReturnPayload):
                raise MessageReturnedError(
                    exchange=payload.exchange,
                    routing_key=payload.routing_key,
                    reply_code=payload.reply_code,
                    reply_text=payload.reply_text,
                )

            elif isinstance(payload, BasicNackPayload):
                raise AMQPStateError("Server NACKed message to be published")

    async def basic_ack(self, delivery_tag: int, *, multiple: bool = False):
        """
        Acknowledges AMQP messages.

        :param delivery_tag: The delivery tag of the message to acknowledge.
        :param multiple: If True, then all messages up to and including the message specified will
                         be acknowledged, not just the message specified.
        """

        payload = BasicAckPayload(delivery_tag, multiple)
        await self._send_single_frame(payload)

    async def basic_reject(self, delivery_tag: int, *, requeue: bool = True):
        """
        Rejects an AMQP message.

        .. note::

            If you are using RabbitMQ, you might want to use :meth:`~.Channel.nack` instead.

        :param delivery_tag: The delivery tag of the message to acknowledge.
        :param requeue: If True, then the rejected message will be requeued if possible.
        """

        payload = BasicRejectPayload(delivery_tag, requeue)
        await self._send_single_frame(payload)

    async def basic_nack(
        self,
        delivery_tag: int,
        *,
        multiple: bool = False,
        requeue: bool = False,
    ):
        """
        Rejects an AMQP message. This is a RabbitMQ-specific extension.

        :param delivery_tag: The delivery tag of the message to acknowledge.
        :param multiple: If True, then all messages up to and including the message specified will
                         be acknowledged, not just the message specified.
        :param requeue: If True, then the rejected message will be requeued if possible.
        """

        payload = BasicNackPayload(delivery_tag, multiple, requeue)
        await self._send_single_frame(payload)

    async def basic_get(self, queue: str, *, no_ack: bool = False) -> Optional[AMQPMessage]:
        """
        Gets a single message from a queue.

        :param queue: The queue to get the message from.
        :param no_ack: Iff not True, then messages will need to be explicitly acknowledged on
                       consumption.
        :return: A :class:`.AMQPMessage` if one existed on the queue, otherwise None.
        """

        payload = BasicGetPayload(reserved_1=0, queue_name=queue, no_ack=no_ack)
        await self._send_single_frame(payload)
        return await self._receive_delivery_message()
