from base64 import b64decode
from collections import Counter
import json
import logging
import os

from django.conf import settings
from django.http import JsonResponse

from staff.monitorings.celery_queue import get_celery_broker_messages_collection


logger = logging.getLogger('monitoring.mongodb')
MAX_CELERY_QUEUE_SIZE = 500
MAX_CELERY_TRIPS_QUEUE_SIZE = 1500


def celery_queue_size_exceeds_limit():
    def get_message_body(message):
        return json.loads(b64decode(json.loads(message['payload'])['body']))

    messages = get_celery_broker_messages_collection()

    default_q = {'queue': os.getenv('DEPLOY_STAGE_ID', 'celery')}
    trips_q = {'queue': settings.CELERY_TRIPS_QUEUE}

    default_queue_size = messages.find(default_q).count()
    trips_queue_size = messages.find(trips_q).count()

    if default_queue_size or trips_queue_size:
        logger.info('Celery default queue size: %s. trips queue size: %s.', default_queue_size, trips_queue_size)

    result = {}
    if default_queue_size > MAX_CELERY_QUEUE_SIZE:
        result['default_queue_size'] = default_queue_size
        counter = Counter((get_message_body(msg).get('task') for msg in messages.find(default_q).limit(5000)))
        result['tasks_in_queue'] = sorted(
            [[k, v] for k, v in counter.items()],
            key=lambda x: x[1],
            reverse=True
        )
    if trips_queue_size > MAX_CELERY_TRIPS_QUEUE_SIZE:
        result['celery_trips_queue_size'] = trips_queue_size

    return result


def check_celery_queue_size(request):
    return JsonResponse(celery_queue_size_exceeds_limit())
