import functools
import logging

from typing import Optional, List
from uuid import uuid4

from celery import Celery, Task
from celery import signals
from celery.exceptions import Ignore
from celery.result import AsyncResult
from celery.utils.nodenames import default_nodename

from django import db
from django.db import transaction
from django.conf import settings

from django_celery_monitoring.mixins import CachedChainTaskMixin
from django_tools_log_context.celery import CtxAwareMixin


logger = logging.getLogger(__name__)


class NoRetry(Exception):
    """
    Исключение, которое можно бросить, когда ретраить таск нет смысла
    """


def get_retry_countdown(retry):
    retry_periods = settings.CELERY_RETRY_PERIODS
    return retry_periods[min(len(retry_periods) - 1, retry)]


class FemidaTask(CtxAwareMixin, CachedChainTaskMixin, Task):
    """
    Если on_transaction_commit == True, то таска будет выполнена только после коммита транзакции
    """
    def apply_async(self, args=None, kwargs=None, task_id=None, **options):
        on_transaction_commit = options.pop('on_transaction_commit', True)
        task_id = task_id or str(uuid4())
        _apply_async = lambda: super(FemidaTask, self).apply_async(
            args=args,
            kwargs=kwargs,
            task_id=task_id,
            **options
        )
        if on_transaction_commit:
            transaction.on_commit(_apply_async)
        else:
            _apply_async()
        return AsyncResult(task_id)

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        if hasattr(self, 'failure_callback'):
            self.failure_callback(*args, **kwargs)
        return super().on_failure(exc, task_id, args, kwargs, einfo)

    def __call__(self, *args, **kwargs):
        try:
            return super().__call__(*args, **kwargs)
        except Ignore:
            # Т.к. мы логируем старт выполнения таска,
            # нужно явно его удалить в случае игнора
            result = AsyncResult(self.request.id)
            result.forget()
            raise


class FemidaCelery(Celery):

    task_cls = 'intranet.femida.src.celery_app:FemidaTask'

    def task(self, *args, **opts):
        # Note: меняем дефолтное поведение
        # – таски не должны быть общими между разными celery-app.
        # Например, через YMQ не должно быть возможно выполнить наши внутр.таски
        opts.setdefault('shared', False)
        return super().task(*args, **opts)

    def autoretry_task(self, *args_task, **opts_task):
        def decorator(func):
            @self.task(*args_task, **opts_task)
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                try:
                    return func(*args, **kwargs)
                except (Ignore, NoRetry):
                    raise
                except Exception as exc:
                    wrapper.retry(
                        countdown=get_retry_countdown(wrapper.request.retries),
                        exc=exc,
                        args=args,
                        kwargs=kwargs,
                    )
            return wrapper
        return decorator

    def start_consuming(self, queue: Optional[str] = None, destination: Optional[List[str]] = None):
        """
        Включает получение тасков из очереди `queue` воркерами `destination`.
        Шорткат для .control.add_consumer с дефолтными значениями.
        """
        queue = queue or self.conf['CELERY_DEFAULT_QUEUE']
        destination = destination or [default_nodename(None)]
        result = self.control.add_consumer(queue, destination=destination, reply=True)
        logger.info('Turned celery consumer ON %s: %s', destination, result)
        return result

    def stop_consuming(self, queue: Optional[str] = None, destination: Optional[List[str]] = None):
        """
        Отключает получение тасков из очереди `queue` воркерами `destination`.
        Шорткат для .control.cancel_consumer с дефолтными значениями.
        """
        queue = queue or self.conf['CELERY_DEFAULT_QUEUE']
        destination = destination or [default_nodename(None)]
        result = self.control.cancel_consumer(queue, destination=destination, reply=True)
        logger.warning('Turned celery consumer OFF %s: %s', destination, result)
        return result


app = FemidaCelery('femida')
app.config_from_object('django.conf:settings', namespace='CELERY')
app.autodiscover_tasks(lambda: settings.INSTALLED_APPS)

beamery_app = FemidaCelery('beamery-to-femida', set_as_current=False)
beamery_app.config_from_object('django.conf:settings', namespace='BEAMERY_CELERY')
beamery_app.autodiscover_tasks(lambda: settings.INSTALLED_APPS)

# Подчищаем протухшие соединения с базой перед каждой таской
signals.task_prerun.connect(db.close_old_connections)
