from django.db.models import DateTimeField, ExpressionWrapper, F, Manager, Max, Q
from django.utils import timezone

from .choices import TaskStatus


class TaskManager(Manager):
    def retryable(self):
        now = timezone.now()

        return self.annotate(
            latest_log_timestamp=Max("logs__timestamp"),
            retryable_time=ExpressionWrapper(
                F("latest_log_timestamp") + F("retry_after"),
                output_field=DateTimeField(),
            ),
        ).filter(
            Q(status=TaskStatus.PENDING)
            | Q(status=TaskStatus.FAILED, retryable_time__lte=now)
        )
