import logging
import functools

from arq import Retry
from sentry_sdk import capture_exception

from intranet.trip.lib.asgi_log_middleware import arq_log_context
from intranet.trip.lib.arq_utils.settings import arq_settings


logger = logging.getLogger(__name__)


async def _execute_job(job_coro, ctx, *args, **kwargs):
    try:
        return await job_coro(ctx, *args, **kwargs)
    except Retry as e:
        capture_exception(e.__cause__)
        raise
    except Exception as e:
        capture_exception(e)
        raise


def job(func):
    @functools.wraps(func)
    async def wrapper(ctx, *args, **kwargs):
        ctx['is_last_try'] = ctx['job_try'] == arq_settings.ARQ_MAX_JOB_TRIES
        with arq_log_context(func, ctx):
            result = await _execute_job(func, ctx, *args, **kwargs)
        return result

    wrapper.arq_config = {
        'is_cron': False
    }
    return wrapper


def periodic_job(predicate=True, **cron_kwargs):
    def decorator(func):
        @functools.wraps(func)
        async def wrapper(ctx, *args, **kwargs):
            ctx['is_last_try'] = ctx['job_try'] == arq_settings.ARQ_MAX_JOB_TRIES
            with arq_log_context(func, ctx):
                result = await _execute_job(func, ctx, *args, **kwargs)
            return result

        wrapper.arq_config = {
            'is_cron': True,
            'predicate': predicate,
            'cron_kwargs': cron_kwargs,
        }
        return wrapper

    return decorator


def default_retry_defer_callable(job_try):
    return 2 ** job_try  # 2, 4, 8, 16... seconds


def retry_on_exception(retry_defer_callable=default_retry_defer_callable):
    def decorator(coro):
        @functools.wraps(coro)
        async def wrapped(ctx, *args, **kwargs):
            try:
                return await coro(ctx, *args, **kwargs)
            except Exception as err:
                logger.exception(
                    'Exception, msg: %s, in task %s:%s on %s try',
                    err,
                    ctx['job_id'],
                    coro.__qualname__,
                    ctx['job_try'],
                )
                raise Retry(defer=retry_defer_callable(ctx['job_try'])) from err
        return wrapped
    return decorator
