"""Model managers."""

import warnings
from functools import wraps
from itertools import count

from celery.utils.time import maybe_timedelta
from django.conf import settings
from django.db import connections, models, router, transaction

from .utils import now, raw_delete

W_ISOLATION_REP = """
Polling results with transaction isolation level 'repeatable-read'
within the same transaction may give outdated results.

Be sure to commit the transaction for each poll iteration.
"""


class TxIsolationWarning(UserWarning):
    """Warning emitted if the transaction isolation level is suboptimal."""


def transaction_retry(max_retries=1):
    """Decorate a function to retry database operations.

    For functions doing database operations, adding
    retrying if the operation fails.

    Keyword Arguments:
    -----------------
        max_retries (int): Maximum number of retries.  Default one retry.

    """
    def _outer(fun):

        @wraps(fun)
        def _inner(*args, **kwargs):
            _max_retries = kwargs.pop('exception_retry_count', max_retries)
            for retries in count(0):
                try:
                    return fun(*args, **kwargs)
                except Exception:   # pragma: no cover
                    # Depending on the database backend used we can experience
                    # various exceptions. E.g. psycopg2 raises an exception
                    # if some operation breaks the transaction, so saving
                    # the task result won't be possible until we rollback
                    # the transaction.
                    if retries >= _max_retries:
                        raise
        return _inner

    return _outer


class ResultManager(models.Manager):
    """Generic manager for celery results."""

    def warn_if_repeatable_read(self):
        if 'mysql' in self.current_engine().lower():
            cursor = self.connection_for_read().cursor()
            # MariaDB and MySQL since 8.0 have different transaction isolation
            # variables: the former has tx_isolation, while the latter has
            # transaction_isolation
            if cursor.execute("SHOW VARIABLES WHERE variable_name IN "
                              "('tx_isolation', 'transaction_isolation');"):
                isolation = cursor.fetchone()[1]
                if isolation == 'REPEATABLE-READ':
                    warnings.warn(TxIsolationWarning(W_ISOLATION_REP.strip()))

    def connection_for_write(self):
        return connections[router.db_for_write(self.model)]

    def connection_for_read(self):
        return connections[self.db]

    def current_engine(self):
        try:
            return settings.DATABASES[self.db]['ENGINE']
        except AttributeError:
            return settings.DATABASE_ENGINE

    def get_all_expired(self, expires):
        """Get all expired results."""
        return self.filter(date_done__lt=now() - maybe_timedelta(expires))

    def delete_expired(self, expires):
        """Delete all expired results."""
        with transaction.atomic():
            raw_delete(queryset=self.get_all_expired(expires))


class TaskResultManager(ResultManager):
    """Manager for :class:`celery.models.TaskResult` models."""

    _last_id = None

    def get_task(self, task_id):
        """Get result for task by ``task_id``.

        Keyword Arguments:
        -----------------
            exception_retry_count (int): How many times to retry by
                transaction rollback on exception.  This could
                happen in a race condition if another worker is trying to
                create the same task.  The default is to retry once.

        """
        try:
            return self.get(task_id=task_id)
        except self.model.DoesNotExist:
            if self._last_id == task_id:
                self.warn_if_repeatable_read()
            self._last_id = task_id
            return self.model(task_id=task_id)

    @transaction_retry(max_retries=2)
    def store_result(self, content_type, content_encoding,
                     task_id, result, status,
                     traceback=None, meta=None,
                     periodic_task_name=None,
                     task_name=None, task_args=None, task_kwargs=None,
                     worker=None, using=None):
        """Store the result and status of a task.

        Arguments:
        ---------
            content_type (str): Mime-type of result and meta content.
            content_encoding (str): Type of encoding (e.g. binary/utf-8).
            task_id (str): Id of task.
            periodic_task_name (str): Celery Periodic task name.
            task_name (str): Celery task name.
            task_args (str): Task arguments.
            task_kwargs (str): Task kwargs.
            result (str): The serialized return value of the task,
                or an exception instance raised by the task.
            status (str): Task status.  See :mod:`celery.states` for a list of
                possible status values.
            worker (str): Worker that executes the task.
            using (str): Django database connection to use.
            traceback (str): The traceback string taken at the point of
                exception (only passed if the task failed).
            meta (str): Serialized result meta data (this contains e.g.
                children).

        Keyword Arguments:
        -----------------
            exception_retry_count (int): How many times to retry by
                transaction rollback on exception.  This could
                happen in a race condition if another worker is trying to
                create the same task.  The default is to retry twice.

        """
        fields = {
            'status': status,
            'result': result,
            'traceback': traceback,
            'meta': meta,
            'content_encoding': content_encoding,
            'content_type': content_type,
            'periodic_task_name': periodic_task_name,
            'task_name': task_name,
            'task_args': task_args,
            'task_kwargs': task_kwargs,
            'worker': worker
        }
        obj, created = self.using(using).get_or_create(task_id=task_id,
                                                       defaults=fields)
        if not created:
            for k, v in fields.items():
                setattr(obj, k, v)
            obj.save(using=using)
        return obj


class GroupResultManager(ResultManager):
    """Manager for :class:`celery.models.GroupResult` models."""

    _last_id = None

    def get_group(self, group_id):
        """Get result for group by ``group_id``.

        Keyword Arguments:
        -----------------
            exception_retry_count (int): How many times to retry by
                transaction rollback on exception.  This could
                happen in a race condition if another worker is trying to
                create the same task.  The default is to retry once.

        """
        try:
            return self.get(group_id=group_id)
        except self.model.DoesNotExist:
            if self._last_id == group_id:
                self.warn_if_repeatable_read()
            self._last_id = group_id
            return self.model(group_id=group_id)

    @transaction_retry(max_retries=2)
    def store_group_result(self, content_type, content_encoding,
                           group_id, result, using=None):
        fields = {
            'result': result,
            'content_encoding': content_encoding,
            'content_type': content_type,
        }

        if not using:
            using = self.db

        obj, created = self.using(using).get_or_create(group_id=group_id,
                                                       defaults=fields)
        if not created:
            for k, v in fields.items():
                setattr(obj, k, v)
            obj.save(using=self.db)
        return obj
