import os
import pymongo

from django.conf import settings
from django.utils import timezone
from kombu.five import Empty
from kombu.transport.mongodb import Channel, Transport
from kombu.utils.encoding import bytes_to_str
from kombu.utils.json import loads, dumps

from intranet.femida.src.utils.datetime import shifted_now


VISIBILITY_TIMEOUT = getattr(settings, 'CELERY_VISIBILITY_TIMEOUT', None)
if os.getenv('ENABLE_VISIBILITY_TIMEOUT', '1') != '1':
    VISIBILITY_TIMEOUT = None


class FemidaMongoChannel(Channel):
    """
    Канал для mongodb в качестве celery-брокера, переписанный под нужды Фемиды.
    От стандартного брокера его отличает возможность использовать настройку,
    аналогичную visibility_timeout для redis или sqs
    https://docs.celeryproject.org/en/stable/getting-started/brokers/redis.html#visibility-timeout

    В стандартной реализации сообщение берётся из очереди и сразу удаляется.
    Ack(nowledge) глобально не делает ничего, а лишь хранит эту информацию на уровне воркера.
    Есть возможность вернуть обратно в очередь сообщения, которые не отправили ack
    (например, это происходит, когда отключается celery воркер).

    В данной реализации сообщение берётся из очереди и скрывается до заданного времени.
    Ack удаляет сообщение из очереди. Соот-но, если ack не произойдёт до заданного времени,
    какой-нибудь воркер снова получит сообщение из очереди.
    """
    def basic_ack(self, delivery_tag, multiple=False):
        """
        Удаляет сообщение из очереди, используя delivery_tag
        """
        message = self.qos.get(delivery_tag)
        delivery_info = message.delivery_info
        for queue in self._lookup(delivery_info['exchange'], delivery_info['routing_key']):
            self.messages.delete_one({
                'delivery_tag': delivery_tag,
                'queue': queue,
            })
        return super().basic_ack(delivery_tag, multiple)

    def _get_message(self, queue):
        """
        Получает из очереди сообщение.

        Если используется настройка VISIBILITY_TIMEOUT,
        то берётся первое видимое сообщение из очереди,
        и ему задаётся атрибут `hidden_till` – время, до которого сообщение нужно скрыть.
        Под видимым сообщением подразумевается
        сообщение, с незаданным либо уже прошедшим `hidden_till`.

        Если настройка не используется, сообщение сразу удаляется из очереди,
        как в стандартной реализации.
        """
        if VISIBILITY_TIMEOUT is None:
            return self.messages.find_one_and_delete(
                filter={'queue': queue},
                sort=[('priority', pymongo.ASCENDING)],
            )
        return self.messages.find_one_and_update(
            filter={
                'queue': queue,
                '$or': [
                    {'hidden_till': None},
                    {'hidden_till': {'$lt': timezone.now()}},
                ],
            },
            update={
                '$set': {
                    'hidden_till': shifted_now(seconds=VISIBILITY_TIMEOUT),
                },
            },
            sort=[('priority', pymongo.ASCENDING)],
        )

    def _get(self, queue):
        """
        ! Не отличается от метода супер-класса,
        здесь только подменяется фактическое получение сообщения из очереди.
        """
        if queue in self._fanout_queues:
            msg = next(self._get_broadcast_cursor(queue), None)
        else:
            msg = self._get_message(queue)

        if self.ttl:
            self._update_queues_expire(queue)

        if msg is None:
            raise Empty()

        return loads(bytes_to_str(msg['payload']))

    def _put(self, queue, message, **kwargs):
        """
        Кладёт сообщение в очередь.
        От стандартной реализации отличают 2 вещи.
        1. Добавлен delivery_tag, по которому потом
           можно отправить ack (удалить сообщ. из очереди).
        2. Используется upsert вместо insert. Т.к. в данной реализации сообщение
           не удаляется из очереди, его восстановление (например, при отключении воркера)
           должно происходить по-другому.
           В стандартной реализации сообщение просто повторно кладётся в очередь,
           в данной реализации – делается upsert по delivery_tag.
        """
        delivery_tag = message['properties']['delivery_tag']
        data = {
            'payload': dumps(message),
            'queue': queue,
            'priority': self._get_message_priority(message, reverse=True),
            'delivery_tag': delivery_tag,
            'hidden_till': None,
        }

        if self.ttl:
            data['expire_at'] = self._get_expire(queue, 'x-message-ttl')

        self.messages.update_one(
            filter={
                'delivery_tag': delivery_tag,
                'queue': queue,
            },
            update={'$set': data},
            upsert=True,
        )


class FemidaMongoTransport(Transport):
    Channel = FemidaMongoChannel
