promotions/apps/promotions/models.py

from django.conf import settings
from django.db import models, transaction
from django.db.models import F
from django.utils import timezone

from oauth2_provider.settings import oauth2_settings

from apps.promotions.handlers import PromotionTypeChoices, BasePromotionHandler
from utils.models import BaseModel

from django.utils.translation import gettext_lazy as _


class EventQuerySet(models.QuerySet):
    pass

class Event(BaseModel):
    label = models.CharField(max_length=255, db_index=True)
    data = models.JSONField()  # from to amount application user etc

    objects = EventQuerySet.as_manager()

    def __str__(self):
        return self.label

    def analyze(self):
        for plan in Plan.objects.get_related_to_event(self):
            plan.process_event(self)



class PlanQuerySet(models.QuerySet):

    # this is for list of discount code for example
    def calculate(self, user, application=None):
        for plan in self.filter(application=application)[:10]:
            if not Promotion.objects.used(user):
                plan.calculate(user)

    def get_related_to_event(self, event):
        return self.filter(event_list__icontains=event.label)


class Plan(BaseModel):
    # user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.PROTECT)
    application = models.ForeignKey(oauth2_settings.APPLICATION_MODEL, on_delete=models.PROTECT,
                                    related_name='+', null=True, blank=True)
    event_list = models.JSONField(default=list)
    title = models.CharField(max_length=255, verbose_name=_('title'), db_index=True)
    description = models.TextField(verbose_name=_('description'), blank=True)
    promotion_type = models.CharField(max_length=64, verbose_name=_('promotion type'), db_index=True, choices=PromotionTypeChoices.choices)
    policy = models.JSONField(default=dict, verbose_name=_('policy'), null=True, blank=True)
    balance = models.IntegerField(verbose_name=_('balance'), default=0)


    objects = PlanQuerySet.as_manager()

    def __str__(self):
        return self.title

    @property
    def handler(self) -> BasePromotionHandler:
        return PromotionTypeChoices(self.promotion_type).handler(
            application=self.application,
            policy=self.policy,
            balance=self.balance,
        )

    def calculate_promotion(self, **kwargs):
        handler = self.handler
        self.base_amount = kwargs.get('base_amount', 0)
        return handler.calculate(**kwargs)

    def reserve_promotion_amount(self, amount):
        updated = False
        if amount and self.balance >= amount:
            updated = Plan.objects.filter(pk=self.pk, balance__gte=amount).update(
                balance=F('balance') - amount
            )
        return updated

    def process_event(self, event=None):
        referral_code = event.data.get('referral_code')
        user_uuid = event.data.get('user_uuid')
        promotion_amount = self.calculate_promotion()
        if referral_code and promotion_amount and not self.promotions.filter(user_id=referral_code, data__user_uuid=user_uuid).exists():
            with transaction.atomic():
                if self.reserve_promotion_amount(promotion_amount):
                    promotion, created = Promotion.objects.get_or_create()
                    if created:
                        raise Exception('Promotion already reserved')
                else:
                    raise Exception('Promotion already reserved')
            promotion.promote()


class PromotionQuerySet(models.QuerySet):
    def used(self, user):
        return self.filter(user=user, state__in=[Promotion.StateChoices.CREATED])


class Promotion(BaseModel):
    class StateChoices(models.IntegerChoices):
        CREATED = 0, _('Created')
        ASSIGNED = 1, _('Assigned')
        PROMOTED = 2, _('Promoted')
        EXPIRED = 10, _('Expired')

    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.PROTECT)
    application = models.ForeignKey(oauth2_settings.APPLICATION_MODEL, on_delete=models.PROTECT,
                                    related_name='+', null=True, blank=True)
    plan = models.ForeignKey(Plan, on_delete=models.PROTECT, related_name='promotions', null=True, blank=True)
    state = models.IntegerField(choices=StateChoices.choices, default=StateChoices.CREATED)
    base_amount = models.IntegerField(null=True, blank=True)
    promotion_amount = models.IntegerField(null=True, blank=True)
    data = models.JSONField(null=True, blank=True)

    objects = PromotionQuerySet.as_manager()

    def change_state(self, from_states: list | tuple, to_state, same_ok=True, raise_exception=False,
                     error_message=None,
                     data: None | dict = None):
        data = data or {}
        data['updated_at'] = timezone.now()

        if same_ok and self.state == to_state:
            return self

        if self.state in from_states:
            updated = Promotion.objects.filter(pk=self.pk, state__in=from_states).update(
                state=to_state,
                **data
            )
            if updated:  # double check in db level
                self.state = to_state
                return self

        if raise_exception:
            raise Exception(error_message)

    def promote(self, **kwargs):
        with transaction.atomic():
            self.plan.reserve_promotion_amount(self)
            self.change_state(from_states=[Promotion.StateChoices.CREATED], to_state=Promotion.StateChoices.ASSIGNED,
                              same_ok=False, raise_exception=True)

        handler = self.plan.handler
        if handler.promote(promotion=self, **kwargs):
            self.change_state(from_states=[Promotion.StateChoices.ASSIGNED], to_state=Promotion.StateChoices.PROMOTED)
            return True
        return False