import logging
import pika
from pika.exchange_type import ExchangeType

from amqp.model.model import AMQRoute
from amqp.rabbitmq.rabbit_mq_producer import RabbitMQProducer
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE, DLQ_EXCHANGE, AMQ_RPC_EXCHANGE
from amqp.router.router_consumer import RouterConsumer
from amqp.router.utils import sanitize_as_rabbitmq_name


class RabbitMQConsumer(RabbitMQProducer):
    PREFETCH_COUNT = 1
    DLQ_ARGS = {"x-dead-letter-exchange": DLQ_EXCHANGE}

    def __init__(self, configuration):
        super().__init__(configuration, RouterConsumer(configuration))

    def cleanup(self):
        logging.info(" [*] RabbitMQConsumer.cleanup()")
        self.router.cleanup()  # de-register route(s)
        if self.router.is_open(self.channel):
            try:
                self.channel.close()
            except (pika.connection.exceptions.AMQPError, pika.connection.exceptions.ChannelClosedByBroker) as e:
                logging.error(" [E] RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)

    def register_inbound_routes(self, callback):
        if self.router.is_open(self.channel):
            queue_args = self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
            route_mapping = self.amq_configuration.amq_adapter.route_mapping
            logging.info(" [*] RabbitMQConsumer register routes, cfg mapping: [%s]", route_mapping)
            requested_routes = self.router.unwrap_route_list(route_mapping)
            for route in requested_routes:
                try:
                    queue_name = route.queue if route.queue else self.router.get_consume_queue_name()
                    logging.info(" [*] RabbitMQConsumer register route: [%s]", route)
                    self.channel.queue_declare(queue_name, durable=True, exclusive=False,
                                               auto_delete=False, arguments=queue_args)
                    exchange = route.exchange if route.exchange else AMQ_RPC_EXCHANGE
                    self.channel.exchange_declare(exchange, exchange_type=pika.exchange_type.ExchangeType.topic)
                    self.channel.queue_bind(queue_name, exchange, route.key)
                    self.channel.basic_consume(
                        queue=queue_name,
                        on_message_callback=callback,
                        auto_ack=False
                    )
                    self.router.register_route(
                        AMQRoute(
                            route.component_name, exchange, queue_name,
                            sanitize_as_rabbitmq_name(route.key),
                            route.timeout, route.valid_until
                        )
                    )
                    logging.info(" [*] AMQ connection initialized for route: %s, listens on: %s/%s",
                          route, exchange, queue_name)
                except pika.connection.exceptions.AMQPError as err:
                    logging.error(" [E] Callback registration INCOMPLETE, %s", err)
        else:
            logging.warning(" [W] Channel is null, cannot register routes.")

    def register_rpc_handler(self, rpc_consumer):
        cm_reply_to_queue = self.router.get_reply_to_queue_name()
        self.channel.queue_declare(cm_reply_to_queue, durable=True, exclusive=False, auto_delete=False)
        self.channel.exchange_declare(
            AMQ_REPLY_TO_EXCHANGE, exchange_type=pika.exchange_type.ExchangeType.direct
        )
        self.channel.queue_bind(cm_reply_to_queue, AMQ_REPLY_TO_EXCHANGE, cm_reply_to_queue)
        self.channel.basic_consume(
            queue=cm_reply_to_queue,
            on_message_callback=rpc_consumer,
            auto_ack=False
        )

    """
    def limit_rate_inside_time_window(self, message: AMQMessage) -> int:
        next_batch_size = BackpressureSubscriber.NO_OP
        now = time.time_ns() // 1000000
        queue_size = self.getQueueSize(message)

        # 1. Backpressure? if unprocessed count exceeds threshold, wait for some time
        if queue_size > THRESHOLD:
            # TODO - signal queue backpressure explicitly to upstream/metrics
            # wait: estimated-per-message-processing-time * SQRT(over-the-limit)
            self.wait_for_milliseconds(
                (WINDOW / self.request_for_window_duration) * (queue_size - THRESHOLD) ** 0.5)
            queue_size = self.getQueueSize(message)  # refresh queue size after wait

        # 2. batch complete? determine next batch size and wait till window end. reset window.
        if self.fulfilled >= self.request_for_window_duration:
            next_batch_size = self.calculateNextBatchSize(queue_size)
            self.request_for_window_duration = next_batch_size
            remaining_time = (self.window_started_at + WINDOW) - now
            if remaining_time > 0:
                self.wait_for_milliseconds(remaining_time)
                queue_size = self.getQueueSize(message)  # refresh queue size after wait
            self.window_started_at = time.time_ns() // 1000000
            self.fulfilled = 0

        # 3. Window elapsed? Reset window & recalculate batch size
        if now >= self.window_started_at + WINDOW:
            if next_batch_size == BackpressureSubscriber.NO_OP:
                next_batch_size = self.calculateNextBatchSize(queue_size)
                self.request_for_window_duration = next_batch_size
            self.window_started_at = time.time_ns() // 1000000
            self.fulfilled = 0

        # 4. send the message from the current batch
        self.send_message(message)
        self.fulfilled += 1

        # 5. report and setup state for the next call
        end = time.time_ns() // 1000000
        self.on_next_timestamp = end
        return next_batch_size

    def calculateNextBatchSize(self, queue_size: int) -> int:
        next_batch_size = self.request_for_window_duration
        if queue_size <= THRESHOLD:
            # processing proceeds well, can increase the rate
            next_batch_size += self.subscriber.increaseRate(self.request_for_window_duration)
            logging.debug(f"INCREASE rate to {next_batch_size}. qSize={queue_size}")
        else:
            decrease = self.subscriber.backoff(queue_size - THRESHOLD)
            if self.request_for_window_duration > decrease:
                next_batch_size -= decrease
            logging.debug(f"DECREASE rate to {next_batch_size}. qSize={queue_size}")
        return next_batch_size
    """
    """
    def getQueueSize(self, message: AMQMessage) -> int:
        try:
            route_queue_name = self.router.findRoute(message)
            queue_name = route_queue_name.queue if route_queue_name else self.TASK_DISPATCH_QUEUE_NAME
            logger.info(f"Get Queue size for {queue_name}")
            return max(len(self.awaitingACK), self.channel.message_count(queue_name))
        except Exception:
            return len(self.awaitingACK)
    """