# coding=utf-8
from __future__ import absolute_import, print_function

import gevent
import traceback

import redis
from collections import OrderedDict
from gevent.event import Event
from suanpan import error, utils, g
from suanpan.log import logger

MQ_DELIVER_KEY = "_suanpan_mq_deliver"


class MQ(object):
    def __init__(
            self,  # pylint: disable=unused-argument
            redisHost="localhost",
            redisPort=6379,
            redisRealtime=False,
            redisKeepalive=True,
            # redisKeepaliveIDLE=120,
            # redisKeepaliveCNT=2,
            # redisKeepaliveINTVL=30,
            redisSocketConnectTimeout=1,
            redisUnixSocketPath=None,
            options=None,
            client=None,
            __p2p_stream=False,
            **kwargs,
    ):
        self.options = options or {}
        self._p2p = __p2p_stream or kwargs.get('__p2p_stream', False)
        self.options.update(
            host=redisHost,
            port=redisPort,
            realtime=redisRealtime,
            keepalive=redisKeepalive,
            # keepidle=redisKeepaliveIDLE,
            # keepcnt=redisKeepaliveCNT,
            # keepintvl=redisKeepaliveINTVL,
            socketConnectTimeout=redisSocketConnectTimeout,
            unixSocketPath=redisUnixSocketPath,
        )
        if client and not isinstance(client, redis.Redis):
            raise error.MQError(f"Invalid redis client: {client}")

        if client:
            self.client = client
        elif self.options["unixSocketPath"]:
            self.client = redis.Redis(
                decode_responses=True, unix_socket_path=self.options["unixSocketPath"]
            )
        else:
            self.client = redis.Redis(
                host=self.options["host"],
                port=self.options["port"],
                decode_responses=True,
                socket_keepalive=self.options["keepalive"],
                # socket_keepalive_options={
                #     socket.TCP_KEEPIDLE: self.options["keepidle"],
                #     socket.TCP_KEEPCNT: self.options["keepcnt"],
                #     socket.TCP_KEEPINTVL: self.options["keepintvl"],
                # },
                socket_connect_timeout=self.options["socketConnectTimeout"],
            )

        self._stop_event = Event()
        self._stop_event.clear()

        if self._p2p:
            self._getP2PConnections()

    def close(self):
        self._stop_event.set()
        self._stop_event.is_set()
        try:
            self.client.close()
            self.client.connection_pool.disconnect()
        finally:
            pass

    @property
    def connected(self):
        return bool(self.client.connection)

    def _getQueue(self, name):
        try:
            return self.client.xinfo_stream(name)
        except redis.RedisError:
            return {}

    def _getQueueGroups(self, name):
        try:
            return self.client.xinfo_groups(name)
        except redis.RedisError:
            return []

    def _lenQueue(self, name):
        try:
            return self.client.xlen(name)
        except redis.RedisError:
            return 0

    def listQueueNames(self):
        return self.client.keys("mq-*")

    def getQueueInfo(self, name):
        # queue = self._getQueue(name)
        length = self._lenQueue(name)
        groups = self._getQueueGroups(name)
        pending = sum(group.get("pending", 0) for group in groups)
        return {
            "name": name,
            "groups": groups,
            # "length": queue.get("length", 0),
            "length": length,
            "pending": pending,
        }

    def createQueue(
            self, name, group="default", consumeID="0", force=False, existOk=False
    ):
        logger.debug(f"Create Queue {name}-{group}")
        if force:
            self.deleteQueue(name)
        return self._createQueue(
            name, group=group, consumeID=consumeID, existOk=existOk
        )

    def _createQueue(self, name, group="default", consumeID="0", existOk=False):
        try:
            return self.client.xgroup_create(name, group, id=consumeID, mkstream=True)
        except redis.RedisError as e:
            # tracebackInfo = traceback.format_exc()
            logger.warn(f"Redis create queue error: {e}")
            if not existOk:
                raise error.MQError(f"Redis queue {name} existed") from e
            return None

    def deleteQueue(self, *names):
        return self.client.delete(*names)

    # def hasQueue(self, name, group="default"):
    #     try:
    #         queue = self.client.xinfo_stream(name)
    #         groups = self.client.xinfo_groups(name)
    #         return any(g["name"].decode() == group for g in groups)
    #     except Exception:
    #         return False

    def sendMessage(
            self, queue, data, messageID="*", maxlen=1000, trimImmediately=False
    ):
        if self._p2p:
            success = data.get('success')
            if success != 'true':
                return

            for target, data4in in self._sendP2PMessages(queue, data):
                logger.debug(f'send p2p message to {target}: {data4in}')
                self.client.xadd(
                    target, data4in, id=messageID, maxlen=maxlen, approximate=(not trimImmediately)
                )
        else:
            return self.client.xadd(
                queue, data, id=messageID, maxlen=maxlen, approximate=(not trimImmediately)
            )

    def sendEvent(self, queue, data, messageID="*", maxlen=1000, trimImmediately=False):
        return self.client.xadd(
            queue, data, id=messageID, maxlen=maxlen, approximate=(not trimImmediately)
        )

    def _getP2PConnections(self):
        from suanpan.graph import graph

        self._outputConnections = OrderedDict()
        node = graph.getNodeById(g.nodeId)
        for port in node.metadata['def'].ports:
            if port.ioType == 'out':
                port_id = port.uuid
                to_dashboard = False
                connections = graph.getOutputConnections(g.nodeId, port_id)
                for tgt in connections:
                    tgt_node = graph.getNodeById(tgt.process)
                    if tgt_node.metadata['def'].type in [22, 10002]:
                        to_dashboard = True
                        break

                self._outputConnections[port_id] = (to_dashboard, connections)

    def _sendP2PMessages(self, queue, data):
        send_master = False
        for port_id, conns in self._outputConnections.items():
            to_dashboard, connections = conns
            send_master = send_master or to_dashboard
            if not to_dashboard:
                for tgt in connections:
                    d = data.get(port_id)
                    if d:
                        data4in = {"id": data['request_id'], tgt.port: d}
                        extra = data.get('extra')
                        if extra:
                            data4in["extra"] = extra
                        target_queue = f'{queue}-{tgt.process}'
                        yield target_queue, data4in

                data.pop(port_id, None)

        if send_master:
            # connection to dashboard, data must send to master
            yield queue, data

    def recvMessages(
            self,
            queue,
            group="default",
            consumer="unknown",
            noack=False,
            block=False,
            count=1,
            consumeID=">",
    ):
        block = None if not block else 0 if block is True else block
        messages = self.client.xreadgroup(
            group, consumer, {queue: consumeID}, count=count, block=block, noack=noack
        )
        messages = list(self._parseMessagesGenerator(messages, group))

        lostMessageIDs = [
            message["id"]
            for message in messages
            if message["id"] and not message["data"]
        ]
        if lostMessageIDs:
            self.client.xack(queue, group, *lostMessageIDs)
            logger.warning(f"Messages have lost: {lostMessageIDs}")

        return [message for message in messages if message["data"]]

    def subscribeQueue(
            self,
            queue,
            group="default",
            consumer="unknown",
            noack=False,
            block=True,
            count=1,
            consumeID=">",
            delay=0,
            errDelay=1,
            errCallback=logger.error,
    ):
        self.createQueue(queue, group=group, existOk=True)
        logger.debug("Subscribing Messages")
        while True:
            try:
                messages = self.recvMessages(
                    queue,
                    group=group,
                    consumer=consumer,
                    noack=noack,
                    block=block,
                    count=count,
                    consumeID=consumeID,
                )
            except Exception as e:  # pylint: disable=broad-except
                if self._stop_event.is_set():
                    logger.debug(f"stop event is set, break.")
                    break

                errCallback(e)
                logger.debug(f"Error in receiving messages. Wait {errDelay}s")
                gevent.sleep(errDelay)
                self.createQueue(queue, group=group, existOk=True)
                continue

            if not messages:
                logger.debug(f"Received no messages. Wait {delay}s")
                gevent.sleep(delay)
                continue

            for message in messages:
                yield message
                self.client.xack(queue, group, message["id"])

    def recvEvents(
            self,
            queue,
            block=False,
            count=1,
            consumeID="$",
    ):
        block = None if not block else 0 if block is True else block
        events = self.client.xread({queue: consumeID}, count=count, block=block)
        events = list(self._parseMessagesGenerator(events, None))

        return [event for event in events if event["data"]]

    def subscribeEvent(
            self,
            queue,
            block=True,
            count=1,
            consumeID="$",
            delay=0,
            errDelay=1,
            errCallback=logger.error,
    ):
        logger.debug("Subscribing Event")
        while True:
            try:
                events = self.recvEvents(queue, block=block, count=count, consumeID=consumeID)
            except Exception as e:  # pylint: disable=broad-except
                if self._stop_event.is_set():
                    logger.debug(f"stop event is set, break.")
                    break

                errCallback(e)
                logger.debug(f"Error in receiving events. Wait {errDelay}s")
                gevent.sleep(errDelay)
                continue

            if not events:
                # logger.debug(f"Received no events. Wait {delay}s")
                gevent.sleep(delay)
                continue

            for event in events:
                yield event

    def recvPendingMessagesInfo(
            self, queue, group="default", consumer="unknown", start="-", end="+", count=None
    ):
        return self.client.xpending_range(
            queue, group, start, end, count, consumername=consumer
        )

    def retryPendingMessages(
            self,
            queue,
            group="default",
            consumer="unknown",
            count=100,
            maxTimes=3,
            timeout=1,
            errCallback=logger.error,
            maxlen=1000,
            trimImmediately=False,
    ):
        logger.debug("Retrying Pending Messages")

        def _getPendingMessages():
            try:
                return self.recvMessages(
                    queue,
                    group=group,
                    consumer=consumer,
                    block=False,
                    count=count,
                    consumeID="0",
                )
            except Exception:  # pylint: disable=broad-except
                logger.warning("Error in getting pending messages:")
                logger.warning(traceback.format_exc())
                return []

        pendingMessages = {msg["id"]: msg for msg in _getPendingMessages()}
        if not pendingMessages:
            logger.debug("Nothing to retry!")
            return

        pendingInfos = {
            msg["message_id"]: msg
            for msg in self.recvPendingMessagesInfo(
                queue, group=group, consumer=consumer, count=count
            )
        }

        for mid in pendingMessages.keys():
            message = pendingMessages[mid]
            info = pendingInfos.get(mid, {})
            message = utils.merge(message, info)
            data = message["data"]
            deliveredTimes = int(data.get(MQ_DELIVER_KEY, 1))
            if deliveredTimes >= maxTimes:
                logger.error(
                    f"Message {mid} retry failed {deliveredTimes} times. Drop!"
                )
                errCallback(message)
                self.client.xack(queue, group, message["id"])
                continue
            timeSinceDelivered = message.get("time_since_delivered", 0)
            if timeSinceDelivered < timeout:
                logger.warning(
                    f"Message {mid} maybe lost: {timeSinceDelivered} < {timeout}"
                )
                continue
            success = self.client.xack(queue, group, message["id"])
            if success:
                data.update({MQ_DELIVER_KEY: deliveredTimes + 1})
                newMID = self.sendMessage(
                    queue, data, maxlen=maxlen, trimImmediately=trimImmediately
                )
                logger.warning(f"Message send back to queue: {mid} -> {newMID}")

    def _parseMessagesGenerator(self, messages, group):
        for message in messages:
            queue, items = message
            for item in items:
                mid, data = item
                yield {"id": mid, "data": data, "queue": queue, "group": group}

    def groupsInfo(self, key):
        return self.client.xinfo_groups(key)

    def streamInfo(self, key):
        return self.client.xinfo_stream(key)
