from datetime import timedelta
from time import sleep

from bson import ObjectId
from django.utils import timezone
from pymongo import MongoClient

from pymongo.errors import ConfigurationError

from django_q.brokers import Broker
from django_q.conf import Conf


def _timeout():
    return timezone.now() - timedelta(seconds=Conf.RETRY)


class Mongo(Broker):
    def __init__(self, list_key=Conf.PREFIX):
        super(Mongo, self).__init__(list_key)
        self.collection = self.get_collection()

    @staticmethod
    def get_connection(list_key=Conf.PREFIX):
        return MongoClient(**Conf.MONGO)

    def get_collection(self):
        if not Conf.MONGO_DB:
            try:
                Conf.MONGO_DB = self.connection.get_default_database().name
            except ConfigurationError:
                Conf.MONGO_DB = 'django-q'
        return self.connection[Conf.MONGO_DB][self.list_key]

    def queue_size(self):
        return self.collection.count({'lock': {'$lte': _timeout()}})

    def lock_size(self):
        return self.collection.count({'lock': {'$gt': _timeout()}})

    def purge_queue(self):
        return self.delete_queue()

    def ping(self):
        return self.info is not None

    def info(self):
        if not self._info:
            self._info = 'MongoDB {}'.format(self.connection.server_info()['version'])
        return self._info

    def fail(self, task_id):
        self.delete(task_id)

    def enqueue(self, task):
        inserted_id = self.collection.insert_one({'payload': task, 'lock': _timeout()}).inserted_id
        return str(inserted_id)

    def dequeue(self):
        task = self.collection.find_one_and_update({'lock': {'$lte': _timeout()}}, {'$set': {'lock': timezone.now()}})
        if task:
            return [(str(task['_id']), task['payload'])]
        # empty queue, spare the cpu
        sleep(Conf.POLL)

    def delete_queue(self):
        return self.collection.drop()

    def delete(self, task_id):
        self.collection.delete_one({'_id': ObjectId(task_id)})

    def acknowledge(self, task_id):
        return self.delete(task_id)
