import asyncio
import io
import logging
from collections import OrderedDict
from contextlib import suppress
from functools import partial
from io import BytesIO
from random import getrandbits
from types import MappingProxyType
from typing import (
    Any, Dict, List, Mapping, Optional,
    Set, Type, Union,
)
from uuid import UUID

import pamqp.frame
from pamqp import commands as spec
from pamqp.base import Frame
from pamqp.body import ContentBody
from pamqp.constants import REPLY_SUCCESS
from pamqp.header import ContentHeader

from aiormq.tools import Countdown, awaitable

from .abc import (
    AbstractChannel, AbstractConnection, ArgumentsType, ChannelFrame,
    ConfirmationFrameType, ConsumerCallback, DeliveredMessage, ExceptionType,
    FrameType, GetResultType, ReturnCallback, RpcReturnType, TimeoutType,
)
from .base import Base, task
from .exceptions import (
    AMQPChannelError, AMQPError, ChannelAccessRefused, ChannelClosed,
    ChannelInvalidStateError, ChannelLockedResource, ChannelNotFoundEntity,
    ChannelPreconditionFailed, DeliveryError, DuplicateConsumerTag,
    InvalidFrameError, MethodNotImplemented, PublishError,
)


log = logging.getLogger(__name__)


EXCEPTION_MAPPING: Mapping[int, Type[AMQPChannelError]] = MappingProxyType({
    403: ChannelAccessRefused,
    404: ChannelNotFoundEntity,
    405: ChannelLockedResource,
    406: ChannelPreconditionFailed,
})


def exception_by_code(frame: spec.Channel.Close) -> AMQPError:
    if frame.reply_code is None:
        return ChannelClosed(frame.reply_code, frame.reply_text)

    exception_class = EXCEPTION_MAPPING.get(frame.reply_code)

    if exception_class is None:
        return ChannelClosed(frame.reply_code, frame.reply_text)

    return exception_class(frame.reply_text)


class Returning(asyncio.Future):
    pass


ConfirmationType = Union[asyncio.Future, Returning]


class Channel(Base, AbstractChannel):
    # noinspection PyTypeChecker
    CONTENT_FRAME_SIZE = len(pamqp.frame.marshal(ContentBody(b""), 0))

    confirmations: Dict[int, ConfirmationType]

    def __init__(
        self,
        connector: AbstractConnection,
        number: int,
        publisher_confirms: bool = True,
        frame_buffer: Optional[int] = None,
        on_return_raises: bool = True,
    ):

        super().__init__(loop=connector.loop, parent=connector)

        self.connection = connector

        if (
            publisher_confirms and not connector.publisher_confirms
        ):  # pragma: no cover
            raise ValueError("Server does't support publisher confirms")

        self.consumers: Dict[str, ConsumerCallback] = {}
        self.confirmations = OrderedDict()
        self.message_id_delivery_tag: Dict[str, int] = dict()

        self.delivery_tag = 0

        self.getter: Optional[asyncio.Future] = None
        self.getter_lock = asyncio.Lock()

        self.frames: asyncio.Queue = asyncio.Queue(maxsize=frame_buffer or 0)

        self.max_content_size = (
            connector.connection_tune.frame_max - self.CONTENT_FRAME_SIZE
        )

        self.__lock = asyncio.Lock()
        self.number: int = number
        self.publisher_confirms = publisher_confirms
        self.rpc_frames: asyncio.Queue = asyncio.Queue(
            maxsize=frame_buffer or 0,
        )
        self.write_queue = connector.write_queue
        self.on_return_raises = on_return_raises
        self.on_return_callbacks: Set[ReturnCallback] = set()
        self._close_exception = None

        self.create_task(self._reader())

        self.__close_reply_code: int = REPLY_SUCCESS
        self.__close_reply_text: str = ""
        self.__close_class_id: int = 0
        self.__close_method_id: int = 0
        self.__close_event: asyncio.Event = asyncio.Event()

    def set_close_reason(
        self, reply_code: int = REPLY_SUCCESS,
        reply_text: str = "", class_id: int = 0, method_id: int = 0,
    ) -> None:
        self.__close_reply_code = reply_code
        self.__close_reply_text = reply_text
        self.__close_class_id = class_id
        self.__close_method_id = method_id

    @property
    def lock(self) -> asyncio.Lock:
        if self.is_closed:
            raise ChannelInvalidStateError("%r closed" % self)

        return self.__lock

    async def _get_frame(self) -> FrameType:
        weight, frame = await self.frames.get()
        self.frames.task_done()
        return frame

    def __str__(self) -> str:
        return str(self.number)

    @task
    async def rpc(
        self, frame: Frame, timeout: TimeoutType = None,
    ) -> RpcReturnType:

        if self.__close_event.is_set():
            raise ChannelInvalidStateError("Channel closed by RPC timeout")

        countdown = Countdown(timeout)
        lock = self.lock

        async with countdown.enter_context(lock):
            try:
                await countdown(
                    self.write_queue.put(
                        ChannelFrame(
                            channel_number=self.number,
                            frames=[frame],
                        ),
                    ),
                )

                if not (frame.synchronous or getattr(frame, "nowait", False)):
                    return None

                result = await countdown(self.rpc_frames.get())

                self.rpc_frames.task_done()

                if result.name not in frame.valid_responses:  # pragma: no cover
                    raise InvalidFrameError(frame)

                return result
            except (asyncio.CancelledError, asyncio.TimeoutError):
                if self.is_closed:
                    raise

                log.warning(
                    "Closing channel %r because RPC call %s cancelled",
                    self, frame,
                )

                self.__close_event.set()
                self.write_queue.put_nowait(
                    ChannelFrame(
                        channel_number=self.number,
                        frames=[
                            spec.Channel.Close(
                                class_id=0,
                                method_id=0,
                                reply_code=504,
                                reply_text=(
                                    "RPC timeout on frame {!s}".format(frame)
                                ),
                            )
                        ],
                    ),
                )

                raise

    async def open(self, timeout: TimeoutType = None) -> spec.Channel.OpenOk:
        frame: spec.Channel.OpenOk = await self.rpc(
            spec.Channel.Open(), timeout=timeout,
        )

        if self.publisher_confirms:
            await self.rpc(spec.Confirm.Select())

        if frame is None:  # pragma: no cover
            raise spec.AMQPFrameError(frame)
        return frame

    async def __get_content_frame(self) -> ContentBody:
        content_frame = await self._get_frame()
        if not isinstance(content_frame, ContentBody):
            raise ValueError(
                "Unexpected frame {!r}".format(content_frame),
                content_frame,
            )
        return content_frame

    async def _read_content(
        self,
        frame: Union[spec.Basic.Deliver, spec.Basic.Return, GetResultType],
        header: ContentHeader,
    ) -> DeliveredMessage:
        with BytesIO() as body:
            content: Optional[ContentBody] = None

            if header.body_size:
                content = await self.__get_content_frame()

            while content and body.tell() < header.body_size:
                body.write(content.value)

                if body.tell() < header.body_size:
                    content = await self.__get_content_frame()

            return DeliveredMessage(
                delivery=frame,
                header=header,
                body=body.getvalue(),
                channel=self,
            )

    async def __get_content_header(self) -> ContentHeader:
        frame: FrameType = await self._get_frame()

        if not isinstance(frame, ContentHeader):
            raise ValueError(
                "Unexpected frame {!r} instead of ContentHeader".format(frame),
                frame,
            )

        return frame

    async def _on_deliver(self, frame: spec.Basic.Deliver) -> None:
        header: ContentHeader = await self.__get_content_header()
        message = await self._read_content(frame, header)

        if frame.consumer_tag is None:
            log.warning("Frame %r has no consumer tag", frame)
            return

        consumer = self.consumers.get(frame.consumer_tag)
        if consumer is not None:
            # noinspection PyAsyncCall
            self.create_task(consumer(message))

    async def _on_get(
        self, frame: Union[spec.Basic.GetOk, spec.Basic.GetEmpty],
    ) -> None:
        message = None
        if isinstance(frame, spec.Basic.GetOk):
            header: ContentHeader = await self.__get_content_header()
            message = await self._read_content(frame, header)
        if isinstance(frame, spec.Basic.GetEmpty):
            message = DeliveredMessage(
                delivery=frame,
                header=ContentHeader(),
                body=b"",
                channel=self,
            )

        getter = getattr(self, "getter", None)

        if getter is None:
            raise RuntimeError("Getter is None")

        if getter.done():
            log.error("Got message but no active getter")
            return

        getter.set_result((frame, message))
        return

    async def _on_return(self, frame: spec.Basic.Return) -> None:
        header: ContentHeader = await self.__get_content_header()
        message = await self._read_content(frame, header)
        message_id = message.header.properties.message_id

        if message_id is None:
            log.error("message_id if None on returned message %r", message)
            return

        delivery_tag = self.message_id_delivery_tag.get(message_id)

        if delivery_tag is None:  # pragma: nocover
            log.error("Unhandled message %r returning", message)
            return

        confirmation = self.confirmations.pop(delivery_tag, None)
        if confirmation is None:  # pragma: nocover
            return

        self.confirmations[delivery_tag] = Returning()

        if self.on_return_raises:
            confirmation.set_exception(PublishError(message, frame))
            return

        for cb in self.on_return_callbacks:
            # noinspection PyBroadException
            try:
                cb(message)
            except Exception:
                log.exception("Unhandled return callback exception")

        confirmation.set_result(message)

    def _confirm_delivery(
        self, delivery_tag: Optional[int],
        frame: ConfirmationFrameType,
    ) -> None:
        if delivery_tag not in self.confirmations:
            return

        confirmation = self.confirmations.pop(delivery_tag)

        if isinstance(confirmation, Returning):
            return
        elif confirmation.done():  # pragma: nocover
            log.warning(
                "Delivery tag %r confirmed %r was ignored", delivery_tag, frame,
            )
            return
        elif isinstance(frame, spec.Basic.Ack):
            confirmation.set_result(frame)
            return

        confirmation.set_exception(
            DeliveryError(None, frame),
        )  # pragma: nocover

    async def _on_confirm(self, frame: ConfirmationFrameType) -> None:
        if not self.publisher_confirms:  # pragma: nocover
            return

        if frame.delivery_tag not in self.confirmations:
            log.error("Unexpected confirmation frame %r from broker", frame)
            return

        multiple = getattr(frame, "multiple", False)

        if multiple:
            for delivery_tag in self.confirmations.keys():
                if frame.delivery_tag >= delivery_tag:
                    # Should be called later to avoid keys copying
                    self.loop.call_soon(
                        self._confirm_delivery, delivery_tag, frame,
                    )
        else:
            self._confirm_delivery(frame.delivery_tag, frame)

    async def _reader(self) -> None:
        while True:
            try:
                frame = await self._get_frame()

                if isinstance(frame, spec.Basic.Deliver):
                    with suppress(Exception):
                        await self._on_deliver(frame)
                    continue
                elif isinstance(
                    frame, (spec.Basic.GetOk, spec.Basic.GetEmpty),
                ):
                    with suppress(Exception):
                        await self._on_get(frame)
                elif isinstance(frame, spec.Basic.Return):
                    with suppress(Exception):
                        await self._on_return(frame)
                    continue
                elif isinstance(frame, spec.Basic.Cancel):
                    if frame.consumer_tag is None:
                        continue
                    self.consumers.pop(frame.consumer_tag, None)
                    continue
                elif isinstance(frame, spec.Basic.CancelOk):
                    if frame.consumer_tag is not None:
                        self.consumers.pop(frame.consumer_tag, None)
                elif isinstance(frame, (spec.Basic.Ack, spec.Basic.Nack)):
                    with suppress(Exception):
                        await self._on_confirm(frame)
                    continue
                elif isinstance(frame, spec.Channel.Close):
                    exc: BaseException = exception_by_code(frame)
                    self.write_queue.put_nowait(
                        ChannelFrame(
                            channel_number=self.number,
                            frames=[spec.Channel.CloseOk()],
                        ),
                    )

                    self.connection.channels.pop(self.number, None)
                    await self._cancel_tasks(exc)
                    return
                elif isinstance(frame, spec.Channel.CloseOk):
                    if self.__close_event.is_set():
                        await self._cancel_tasks(asyncio.TimeoutError())
                        self.connection.channels.pop(self.number, None)
                        return

                await self.rpc_frames.put(frame)
            except asyncio.CancelledError:
                return
            except Exception as e:  # pragma: nocover
                log.debug("Channel reader exception %r", exc_info=e)
                await self._cancel_tasks(e)
                raise

    @task
    async def _on_close(self, exc: Optional[ExceptionType] = None) -> None:
        if self.connection.is_opened and not self.__close_event.is_set():
            await self.rpc(
                spec.Channel.Close(
                    reply_code=self.__close_reply_code,
                    class_id=self.__close_class_id,
                    method_id=self.__close_method_id,
                ),
                timeout=self.connection.connection_tune.heartbeat or None,
            )
        self.connection.channels.pop(self.number, None)

    async def basic_get(
        self, queue: str = "", no_ack: bool = False,
        timeout: TimeoutType = None,
    ) -> DeliveredMessage:

        countdown = Countdown(timeout)
        async with countdown.enter_context(self.getter_lock):
            self.getter = self.create_future()

            await self.rpc(
                spec.Basic.Get(queue=queue, no_ack=no_ack),
                timeout=countdown.get_timeout(),
            )

            frame: Union[spec.Basic.GetEmpty, spec.Basic.GetOk]
            message: DeliveredMessage

            frame, message = await countdown(self.getter)
            del self.getter

        return message

    async def basic_cancel(
        self, consumer_tag: str, *, nowait: bool = False,
        timeout: TimeoutType = None
    ) -> spec.Basic.CancelOk:
        return await self.rpc(
            spec.Basic.Cancel(consumer_tag=consumer_tag, nowait=nowait),
            timeout=timeout,
        )

    async def basic_consume(
        self,
        queue: str,
        consumer_callback: ConsumerCallback,
        *,
        no_ack: bool = False,
        exclusive: bool = False,
        arguments: ArgumentsType = None,
        consumer_tag: str = None,
        timeout: TimeoutType = None
    ) -> spec.Basic.ConsumeOk:

        consumer_tag = consumer_tag or "ctag%i.%s" % (
            self.number,
            UUID(int=getrandbits(128), version=4).hex,
        )

        if consumer_tag in self.consumers:
            raise DuplicateConsumerTag(self.number)

        self.consumers[consumer_tag] = awaitable(consumer_callback)

        return await self.rpc(
            spec.Basic.Consume(
                queue=queue,
                no_ack=no_ack,
                exclusive=exclusive,
                consumer_tag=consumer_tag,
                arguments=arguments,
            ),
            timeout=timeout,
        )

    async def basic_ack(
        self, delivery_tag: int, multiple: bool = False, wait: bool = True,
    ) -> None:
        drain_future = self.create_future() if wait else None

        await self.write_queue.put(
            ChannelFrame(
                frames=[
                    spec.Basic.Ack(
                        delivery_tag=delivery_tag,
                        multiple=multiple,
                    ),
                ],
                channel_number=self.number,
                drain_future=drain_future,
            ),
        )

        if drain_future is not None:
            await drain_future

    async def basic_nack(
        self,
        delivery_tag: int,
        multiple: bool = False,
        requeue: bool = True,
        wait: bool = True,
    ) -> None:
        if not self.connection.basic_nack:
            raise MethodNotImplemented

        drain_future = self.create_future() if wait else None

        await self.write_queue.put(
            ChannelFrame(
                frames=[
                    spec.Basic.Nack(
                        delivery_tag=delivery_tag,
                        multiple=multiple,
                        requeue=requeue,
                    ),
                ],
                channel_number=self.number,
                drain_future=drain_future,
            ),
        )

        if drain_future is not None:
            await drain_future

    async def basic_reject(
        self, delivery_tag: int, *, requeue: bool = True, wait: bool = True
    ) -> None:
        drain_future = self.create_future()
        await self.write_queue.put(
            ChannelFrame(
                channel_number=self.number,
                frames=[
                    spec.Basic.Reject(
                        delivery_tag=delivery_tag,
                        requeue=requeue,
                    ),
                ],
                drain_future=drain_future,
            ),
        )

        if drain_future is not None:
            await drain_future

    def _split_body(self, body: bytes) -> List[ContentBody]:
        if not body:
            return []

        if len(body) < self.max_content_size:
            return [ContentBody(body)]

        with io.BytesIO(body) as fp:
            reader = partial(fp.read, self.max_content_size)
            return list(map(ContentBody, iter(reader, b"")))

    async def basic_publish(
        self,
        body: bytes,
        *,
        exchange: str = "",
        routing_key: str = "",
        properties: spec.Basic.Properties = None,
        mandatory: bool = False,
        immediate: bool = False,
        timeout: TimeoutType = None,
        wait: bool = True,
    ) -> Optional[ConfirmationFrameType]:
        drain_future = self.create_future() if wait else None
        countdown = Countdown(timeout=timeout)

        publish_frame = spec.Basic.Publish(
            exchange=exchange,
            routing_key=routing_key,
            mandatory=mandatory,
            immediate=immediate,
        )

        content_header = ContentHeader(
            properties=properties or spec.Basic.Properties(delivery_mode=1),
            body_size=len(body),
        )

        if not content_header.properties.message_id:
            # UUID compatible random bytes
            rnd_uuid = UUID(int=getrandbits(128), version=4)
            content_header.properties.message_id = rnd_uuid.hex

        confirmation: Optional[ConfirmationType] = None

        async with countdown.enter_context(self.lock):
            self.delivery_tag += 1

            if self.publisher_confirms:
                message_id = content_header.properties.message_id

                if self.delivery_tag not in self.confirmations:
                    self.confirmations[
                        self.delivery_tag
                    ] = self.create_future()

                confirmation = self.confirmations[self.delivery_tag]
                self.message_id_delivery_tag[message_id] = self.delivery_tag

                if confirmation is None:
                    return

                confirmation.add_done_callback(
                    lambda _: self.message_id_delivery_tag.pop(
                        message_id, None,
                    ),
                )

            body_frames: List[Union[FrameType, ContentBody]]
            body_frames = [publish_frame, content_header]
            body_frames += self._split_body(body)

            await countdown(
                self.write_queue.put(
                    ChannelFrame(
                        frames=body_frames,
                        channel_number=self.number,
                        drain_future=drain_future,
                    ),
                ),
            )

            if drain_future:
                await drain_future

            if not self.publisher_confirms:
                return None

            if confirmation is None:
                return None

            return await countdown(confirmation)

    async def basic_qos(
        self,
        *,
        prefetch_size: int = None,
        prefetch_count: int = None,
        global_: bool = False,
        timeout: TimeoutType = None
    ) -> spec.Basic.QosOk:
        return await self.rpc(
            spec.Basic.Qos(
                prefetch_size=prefetch_size or 0,
                prefetch_count=prefetch_count or 0,
                global_=global_,
            ),
            timeout=timeout,
        )

    async def basic_recover(
        self, *, nowait: bool = False, requeue: bool = False,
        timeout: TimeoutType = None
    ) -> spec.Basic.RecoverOk:
        frame: Union[spec.Basic.RecoverAsync, spec.Basic.Recover]
        if nowait:
            frame = spec.Basic.RecoverAsync(requeue=requeue)
        else:
            frame = spec.Basic.Recover(requeue=requeue)

        return await self.rpc(frame, timeout=timeout)

    async def exchange_declare(
        self,
        exchange: str = "",
        *,
        exchange_type: str = "direct",
        passive: bool = False,
        durable: bool = False,
        auto_delete: bool = False,
        internal: bool = False,
        nowait: bool = False,
        arguments: Optional[Dict[str, Any]] = None,
        timeout: TimeoutType = None
    ) -> spec.Exchange.DeclareOk:
        return await self.rpc(
            spec.Exchange.Declare(
                exchange=str(exchange),
                exchange_type=str(exchange_type),
                passive=bool(passive),
                durable=bool(durable),
                auto_delete=bool(auto_delete),
                internal=bool(internal),
                nowait=bool(nowait),
                arguments=arguments,
            ),
            timeout=timeout,
        )

    async def exchange_delete(
        self,
        exchange: str = "",
        *,
        if_unused: bool = False,
        nowait: bool = False,
        timeout: TimeoutType = None
    ) -> spec.Exchange.DeleteOk:
        return await self.rpc(
            spec.Exchange.Delete(
                exchange=exchange, nowait=nowait, if_unused=if_unused,
            ),
            timeout=timeout,
        )

    async def exchange_bind(
        self,
        destination: str = "",
        source: str = "",
        routing_key: str = "",
        *,
        nowait: bool = False,
        arguments: dict = None,
        timeout: TimeoutType = None
    ) -> spec.Exchange.BindOk:
        return await self.rpc(
            spec.Exchange.Bind(
                destination=destination,
                source=source,
                routing_key=routing_key,
                nowait=nowait,
                arguments=arguments,
            ),
            timeout=timeout,
        )

    async def exchange_unbind(
        self,
        destination: str = "",
        source: str = "",
        routing_key: str = "",
        *,
        nowait: bool = False,
        arguments: dict = None,
        timeout: TimeoutType = None
    ) -> spec.Exchange.UnbindOk:
        return await self.rpc(
            spec.Exchange.Unbind(
                destination=destination,
                source=source,
                routing_key=routing_key,
                nowait=nowait,
                arguments=arguments,
            ),
            timeout=timeout,
        )

    async def flow(
        self, active: bool,
        timeout: TimeoutType = None,
    ) -> spec.Channel.FlowOk:
        return await self.rpc(
            spec.Channel.Flow(active=active),
            timeout=timeout,
        )

    async def queue_bind(
        self,
        queue: str,
        exchange: str,
        routing_key: str = "",
        nowait: bool = False,
        arguments: dict = None,
        timeout: TimeoutType = None,
    ) -> spec.Queue.BindOk:
        return await self.rpc(
            spec.Queue.Bind(
                queue=queue,
                exchange=exchange,
                routing_key=routing_key,
                nowait=nowait,
                arguments=arguments,
            ),
            timeout=timeout,
        )

    async def queue_declare(
        self,
        queue: str = "",
        *,
        passive: bool = False,
        durable: bool = False,
        exclusive: bool = False,
        auto_delete: bool = False,
        nowait: bool = False,
        arguments: dict = None,
        timeout: TimeoutType = None
    ) -> spec.Queue.DeclareOk:
        return await self.rpc(
            spec.Queue.Declare(
                queue=queue,
                passive=bool(passive),
                durable=bool(durable),
                exclusive=bool(exclusive),
                auto_delete=bool(auto_delete),
                nowait=bool(nowait),
                arguments=arguments,
            ),
            timeout=timeout,
        )

    async def queue_delete(
        self,
        queue: str = "",
        if_unused: bool = False,
        if_empty: bool = False,
        nowait: bool = False,
        timeout: TimeoutType = None,
    ) -> spec.Queue.DeleteOk:
        return await self.rpc(
            spec.Queue.Delete(
                queue=queue,
                if_unused=if_unused,
                if_empty=if_empty,
                nowait=nowait,
            ),
            timeout=timeout,
        )

    async def queue_purge(
        self, queue: str = "", nowait: bool = False,
        timeout: TimeoutType = None,
    ) -> spec.Queue.PurgeOk:
        return await self.rpc(
            spec.Queue.Purge(queue=queue, nowait=nowait),
            timeout=timeout,
        )

    async def queue_unbind(
        self,
        queue: str = "",
        exchange: str = "",
        routing_key: str = "",
        arguments: dict = None,
        timeout: TimeoutType = None,
    ) -> spec.Queue.UnbindOk:
        return await self.rpc(
            spec.Queue.Unbind(
                routing_key=routing_key,
                arguments=arguments,
                queue=queue,
                exchange=exchange,
            ),
            timeout=timeout,
        )

    async def tx_commit(
        self, timeout: TimeoutType = None,
    ) -> spec.Tx.CommitOk:
        return await self.rpc(spec.Tx.Commit(), timeout=timeout)

    async def tx_rollback(
        self, timeout: TimeoutType = None,
    ) -> spec.Tx.RollbackOk:
        return await self.rpc(spec.Tx.Rollback(), timeout=timeout)

    async def tx_select(self, timeout: TimeoutType = None) -> spec.Tx.SelectOk:
        return await self.rpc(spec.Tx.Select(), timeout=timeout)

    async def confirm_delivery(
        self, nowait: bool = False,
        timeout: TimeoutType = None,
    ) -> spec.Confirm.SelectOk:
        return await self.rpc(
            spec.Confirm.Select(nowait=nowait),
            timeout=timeout,
        )
