from __future__ import annotations

import abc
import json
import time
from contextlib import contextmanager
from typing import Any, Callable, Iterator, NamedTuple, Optional

import pika
from pika.connection import Parameters

from rabbitmq_alphamoon.logger import get_logger

LOGGER = get_logger('rabbitmq_alphamoon')


class IQueueConnector(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def consume_forever(self, process_message_callback: Callable[[Any], None]):
        pass

    @abc.abstractmethod
    def publish(self, message: Any) -> None:
        pass

    @property
    @abc.abstractmethod
    def message_count(self) -> int:
        pass


_RabbitMQConnection = NamedTuple(
    '_RabbitMQConnection',
    [
        ('connection', pika.BlockingConnection),
        ('channel', pika.adapters.blocking_connection.BlockingChannel),
        ('queue_declaration', pika.frame.Method),
    ],
)


class RabbitMQConnector(IQueueConnector):
    def __init__(
        self,
        connection_parameters: Parameters,
        queue: str,
        durable: bool = True,
        exclusive: bool = False,
        auto_delete: bool = False,
    ):
        self._connection_parameters = connection_parameters
        self._queue = queue
        self._durable = durable
        self._exclusive = exclusive
        self._auto_delete = auto_delete

        self._connection: Optional[_RabbitMQConnection] = None
        self._queue_declaration = None

    def _connect(self):
        LOGGER.debug("Connecting to broker, queue %s", self._queue)
        connection = pika.BlockingConnection(self._connection_parameters)
        channel = connection.channel()
        channel.confirm_delivery()
        queue_declaration = channel.queue_declare(
            queue=self._queue,
            durable=self._durable,
            exclusive=self._exclusive,
            auto_delete=self._auto_delete,
        )
        self._connection = _RabbitMQConnection(
            connection=connection,
            channel=channel,
            queue_declaration=queue_declaration,
        )
        LOGGER.info("Connection with broker established, queue %s", self._queue)

    def _assure_connection(self):
        while True:
            try:
                self._connect()
            except pika.exceptions.AMQPConnectionError as error:
                error_type = type(error).__name__  # e.g. 'AMQPConnectionError'
                LOGGER.warning("%s, queue %s: %s, retrying...", error_type, self._queue, error)
                time.sleep(1)
            else:
                break

    def _disconnect(self):
        LOGGER.debug("Disconnecting from broker, queue %s", self._queue)
        self._connection.connection.close()
        self._connection = None
        LOGGER.info("Disconnected from broker, queue %s", self._queue)

    @contextmanager
    def open_connection(self) -> Iterator[RabbitMQConnector]:
        self._assure_connection()
        try:
            yield self  # simplify one-time connector usage, especially when `consume_forever` used
        finally:
            self._disconnect()

    def _get_opened_connection(self) -> _RabbitMQConnection:
        if self._connection is None:
            raise RuntimeError("Connection not opened. Forgot to run in `open_connection` context?")
        return self._connection

    def _reconnect_and_get_opened_connection(self) -> _RabbitMQConnection:
        LOGGER.info("Reconnecting to queue, queue %s", self._queue)
        try:
            self._disconnect()
        except pika.exceptions.AMQPConnectionError as error:
            LOGGER.debug("Ignored disconnection error, queue %s, error %s", self._queue, error)
        else:
            LOGGER.warning("Connection was established, reconnecting anyway, queue %s", self._queue)

        self._assure_connection()
        return self._get_opened_connection()

    def consume_forever(self, process_message_callback: Callable[[Any], None]):
        connection = self._get_opened_connection()

        def callback_fun(
            channel: pika.channel.Channel,
            method_frame: pika.spec.Basic.Deliver,
            header_frame: pika.BasicProperties,
            body: bytes,
        ):
            LOGGER.info("Received message, deserializing, queue %s", self._queue)
            message = json.loads(body)

            LOGGER.info("Processing message, queue %s", self._queue)
            process_message_callback(message)

            LOGGER.debug("Sending ack, queue %s", self._queue)
            channel.basic_ack(delivery_tag=method_frame.delivery_tag)
            LOGGER.info("Sent ack, queue %s", self._queue)

        try:
            while True:
                try:
                    LOGGER.debug("Starting consumption, queue %s", self._queue)
                    connection.channel.basic_qos(prefetch_count=1)
                    connection.channel.basic_consume(
                        queue=self._queue,
                        on_message_callback=callback_fun,
                        auto_ack=False,
                    )
                    connection.channel.start_consuming()
                except KeyboardInterrupt:
                    break
                except pika.exceptions.AMQPConnectionError as error:
                    error_type = type(error).__name__  # e.g. 'ConnectionClosedByBroker'
                    LOGGER.warning("%s, queue %s: %s, retrying...", error_type, self._queue, error)
                    time.sleep(1)
                    connection = self._reconnect_and_get_opened_connection()
                except pika.exceptions.AMQPChannelError as error:
                    error_type = type(error).__name__  # e.g. 'ChannelWrongStateError'
                    LOGGER.error("%s, queue %s: %s, stopping...", error_type, self._queue, error)
                    break
        finally:
            LOGGER.debug("Stopping consumption, queue %s", self._queue)
            connection.channel.stop_consuming()

    def publish(self, message: Any) -> None:
        connection = self._get_opened_connection()
        LOGGER.debug("Serializing message, queue %s", self._queue)
        body = json.dumps(
            message,
            ensure_ascii=False,
            separators=(',', ':'),  # compact encoding
        )

        try:
            LOGGER.debug("Publishing message, queue %s", self._queue)
            connection.channel.basic_publish(
                exchange='',
                routing_key=self._queue,
                body=body.encode(),
                properties=pika.BasicProperties(
                    delivery_mode=pika.spec.PERSISTENT_DELIVERY_MODE,  # persists during restarts
                ),
                mandatory=True,  # ack's when message is received
            )
            LOGGER.info("Published message, queue %s", self._queue)
        except pika.exceptions.UnroutableError as error:
            LOGGER.error("Publishing failed, queue %s, %s", self._queue, error)

    @property
    def message_count(self) -> int:
        connection = self._get_opened_connection()
        return connection.queue_declaration.method.message_count
