import json
import time
import logging
import functools
from concurrent.futures import ThreadPoolExecutor

import pika

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s - %(levelname)s - %(message)s")


class MQConsumer:

    def __init__(self, config):
        self.config = config
        self.executor = ThreadPoolExecutor(config.workers)

        self._connection = None
        self._channel = None
        self._closing = False

    # ==============================
    # 连接 RabbitMQ
    # ==============================

    def _connect(self):

        params = pika.ConnectionParameters(
            host=self.config.host,
            port=self.config.port,
            virtual_host=self.config.vhost,
            credentials=pika.PlainCredentials(
                self.config.username,
                self.config.password
            ),
            heartbeat=60,
            blocked_connection_timeout=300
        )

        connection = pika.BlockingConnection(params)
        channel = connection.channel()

        channel.basic_qos(prefetch_count=self.config.prefetch)

        return connection, channel

    # ==============================
    # 启动消费者
    # ==============================

    def listen(self, queue, handler):

        while not self._closing:

            try:

                self._connection, self._channel = self._connect()

                self._channel.queue_declare(
                    queue=queue,
                    durable=True,
                    arguments={
                        "x-dead-letter-exchange": self.config.dlx_exchange
                    }
                )
                # 声明延迟交换机
                self._channel.exchange_declare(
                    exchange=self.config.delay_exchange,
                    exchange_type="x-delayed-message",
                    durable=True,
                    arguments={"x-delayed-type": "direct"}
                )

                # 绑定 queue
                self._channel.queue_bind(
                    queue=queue,
                    exchange=self.config.delay_exchange,
                    routing_key=queue
                )

                self._channel.basic_consume(
                    queue=queue,
                    on_message_callback=lambda ch, method, properties, body:
                    self._on_message(ch, method, properties, body, queue, handler),
                    auto_ack=False
                )

                logging.info(f"consumer started -> {queue}")

                self._channel.start_consuming()

            except Exception as e:

                logging.error(f"consumer error: {e}")

                try:
                    if self._connection and self._connection.is_open:
                        self._connection.close()
                except:
                    pass

                time.sleep(5)

    # ==============================
    # 收到消息
    # ==============================

    def _on_message(self, ch, method, properties, body, queue, handler):

        try:
            data = json.loads(body)
        except:
            data = body

        future = self.executor.submit(handler, data)

        future.add_done_callback(
            lambda f: self._on_task_done(
                f, ch, method, properties, body, queue
            )
        )

    # ==============================
    # 任务执行完成
    # ==============================

    def _on_task_done(self, future, ch, method, properties, body, queue):

        def ack():
            if ch.is_open:
                ch.basic_ack(method.delivery_tag)

        headers = properties.headers or {}
        retry = headers.get("x-retry", 0)

        try:

            result = future.result()

            logging.info(f"task success: {result}")

            cb = functools.partial(ack)
            ch.connection.add_callback_threadsafe(cb)

        except Exception as e:

            delay = self._get_retry_delay(retry)

            if delay is None:

                logging.error(f"task failed max retry: {e}")

                cb = functools.partial(ack)
                ch.connection.add_callback_threadsafe(cb)

                return

            def retry_publish():

                if ch.is_open:

                    ch.basic_publish(
                        exchange=self.config.delay_exchange,
                        routing_key=queue,
                        body=body,
                        properties=pika.BasicProperties(
                            headers={
                                "x-delay": delay,
                                "x-retry": retry + 1
                            },
                            message_id=properties.message_id,
                            delivery_mode=2
                        )
                    )

                    ch.basic_ack(method.delivery_tag)

            cb = functools.partial(retry_publish)

            ch.connection.add_callback_threadsafe(cb)

            logging.warning(
                f"retry later queue={queue} retry={retry} delay={delay}"
            )

    # ==============================
    # 重试策略
    # ==============================

    def _get_retry_delay(self, retry):

        retry_policy = [

            5 * 1000,
            30 * 1000,
            60 * 1000,
            5 * 60 * 1000,
            10 * 60 * 1000
        ]

        if retry >= len(retry_policy):
            return None

        return retry_policy[retry]

    # ==============================
    # 优雅关闭
    # ==============================

    def shutdown(self):

        self._closing = True

        logging.info("consumer shutting down...")

        try:

            if self._channel and self._channel.is_open:
                self._channel.stop_consuming()

            if self._connection and self._connection.is_open:
                self._connection.close()

        except:
            pass

        self.executor.shutdown(wait=True)

        logging.info("consumer stopped")