from django.db import models


class AutoAnnotationField:
    annotations_cls_attr_name = '_annotations'

    def __init__(self, expression):
        self.expression = expression

    def contribute_to_class(self, cls, name, **kwargs):
        if hasattr(cls, self.annotations_cls_attr_name):
            getattr(cls, self.annotations_cls_attr_name)[name] = self.expression
        else:
            setattr(cls, self.annotations_cls_attr_name, {name: self.expression})


class AutoAnnotationsManager(models.Manager):
    """
    Usage: objects = AnnotationManager.from_queryset(CustomQuerySet)()
    """
    annotations_cls_attr_name = AutoAnnotationField.annotations_cls_attr_name

    def get_queryset(self):
        qs = super().get_queryset()
        if hasattr(qs, 'prefetch_related_annotations'):
            qs = qs.prefetch_related_annotations()
        annotations = getattr(self.model, self.annotations_cls_attr_name, None)
        if not annotations:
            return qs
        return qs.annotate(**annotations)


class AutoAnnotationsQuerySetMixin:
    def prefetch_related_annotations(self):
        raise NotImplementedError


__all__ = [
    'AutoAnnotationField',
    'AutoAnnotationsManager',
]