import logging from datetime import date, timedelta from decimal import Decimal, ROUND_HALF_UP from django.db import transaction from django.utils import timezone from django.db.models import Q from django_cronjob_utils import CronTask, register_task from django_cronjob_utils.models import CronExecution from apps.investments.models import Investment, Reward from apps.referrals.models import Commission from apps.transactions.models import Transaction from apps.users.models import User logger = logging.getLogger(__name__) DECIMAL_QUANTIZE = Decimal("0.00000001") COMMISSION_RATES = { 1: Decimal("50.00"), 2: Decimal("30.00"), 3: Decimal("20.00"), } CALCULATE_TASK_CODE = "R001" DISTRIBUTE_TASK_CODE = "D001" @register_task("calculate-rewards", CALCULATE_TASK_CODE) class CalculateRewardsTask(CronTask): """ Phase 1: Calculate daily rewards. This task should be run for a specific date (e.g., Yesterday). It identifies active investments for that date and creates PENDING Reward records. """ def execute(self, date: date) -> dict: target_date = date # target_date is the date for which we calculate rewards reward_date = target_date # Calculate new rewards (pending distribution) rewards_calculated = self._calculate_pending_rewards(reward_date) return { "error": False, "message": f"Calculated {rewards_calculated} new rewards for {reward_date}." } def _calculate_pending_rewards(self, reward_date: date) -> int: """ Identify eligible investments and create pending Reward records. reward_date: The date for which we're calculating rewards (typically yesterday). """ # Eligibility: # 1. Status is Pending or Active # 2. Covers the reward_date # 3. First reward starts the day AFTER start_date (not on start_date itself) eligible_investments = Investment.objects.filter( status__in=[Investment.STATUS_PENDING, Investment.STATUS_ACTIVE], start_date__date__lt=reward_date, # Must have started BEFORE reward_date (not on it) end_date__date__gte=reward_date, # And not yet expired on reward_date ).select_for_update() count = 0 with transaction.atomic(): for investment in eligible_investments: # Calculate reward amount reward_amount = self._quantize( investment.amount * investment.daily_reward_rate / Decimal("100") ) # Create as pending (distributed_at=None) # We use get_or_create to make it idempotent reward, created = Reward.objects.get_or_create( investment=investment, reward_date=reward_date, defaults={ "amount": reward_amount, "distributed_at": None, }, ) # Update investment status # If PENDING, move to ACTIVE # If reward_date reached end_date, move to COMPLETED new_status = investment.status if investment.status == Investment.STATUS_PENDING: new_status = Investment.STATUS_ACTIVE if reward_date >= investment.end_date.date(): new_status = Investment.STATUS_COMPLETED if investment.status != new_status: investment.status = new_status investment.save(update_fields=["status", "updated_at"]) if created: count += 1 return count def _quantize(self, value: Decimal | float | int) -> Decimal: return Decimal(value).quantize(DECIMAL_QUANTIZE, rounding=ROUND_HALF_UP) @register_task("distribute-rewards", DISTRIBUTE_TASK_CODE) class DistributeRewardsTask(CronTask): """ Phase 2: Distribute rewards. DEPENDENCY: Depends on 'calculate-rewards' completing successfully for the same date. This task distributes the rewards calculated for the given date. """ def should_run(self) -> bool: """ Custom check: Ensure CalculateRewardsTask (INV001) completed successfully for this execution date. """ # Check standard duplicate prevention first if not super().should_run(): return False # Check dependency dependency_completed = CronExecution.objects.filter( task_code=CALCULATE_TASK_CODE, execution_date=self.execution_date, success=True ).exists() if not dependency_completed: logger.warning( f"Dependency {CALCULATE_TASK_CODE} not completed for {self.execution_date}. " f"Skipping {self.task_code}." ) return False return True def execute(self, date: date) -> dict: target_date = date now = timezone.now() # Determine items to distribute: # Rewards with reward_date = target_date AND distributed_at is NULL rewards_distributed, commissions_created = self._distribute_rewards(target_date, now) return { "error": False, "message": f"Distributed {rewards_distributed} rewards and {commissions_created} commissions for {target_date}." } def _distribute_rewards(self, target_date: date, now) -> tuple[int, int]: # pending: distributed_at is NULL and reward_date == target_date # (We only process the specific date we are asked to run for, # to align with the dependency check for THAT date) pending_rewards = Reward.objects.filter( distributed_at__isnull=True, reward_date=target_date ).select_related("investment", "investment__user").select_for_update() reward_count = 0 commission_count = 0 with transaction.atomic(): for reward in pending_rewards: # Credit the user self._credit_user( user=reward.investment.user, amount=reward.amount, tx_type=Transaction.TYPE_REWARD, metadata={"investment_id": reward.investment.id, "reward_id": reward.id}, ) # Update investment ROI totals investment = Investment.objects.select_for_update().get(pk=reward.investment.id) investment.roi_total_amount = self._quantize(investment.roi_total_amount + reward.amount) if investment.amount > 0: investment.roi_total_percent = self._quantize( (investment.roi_total_amount / investment.amount) * Decimal("100") ) investment.save(update_fields=["roi_total_amount", "roi_total_percent", "updated_at"]) # Mark as distributed reward.distributed_at = now reward.save(update_fields=["distributed_at"]) reward_count += 1 # Calculate and distribute commissions immediately c_count = self._create_commissions(reward, reward.investment.user, now) commission_count += c_count return reward_count, commission_count def _create_commissions(self, reward: Reward, downline_user: User, now) -> int: commission_count = 0 # Level 1 = Direct Upline uplines: list[User] = [] current_user = downline_user.referred_by while current_user and len(uplines) < 3: uplines.append(current_user) current_user = current_user.referred_by for i, upline_user in enumerate(uplines): level = i + 1 rate = COMMISSION_RATES.get(level) if not rate: break amount = self._quantize(reward.amount * rate / Decimal("100")) commission, created = Commission.objects.get_or_create( reward=reward, level=level, upline_user=upline_user, defaults={ "downline_user": downline_user, "commission_rate": rate, "amount": amount, "status": Commission.STATUS_DISTRIBUTED, "distributed_at": now, }, ) if created: commission_count += 1 self._credit_user( user=upline_user, amount=amount, tx_type=Transaction.TYPE_COMMISSION, metadata={ "reward_id": reward.id, "downline_user_id": downline_user.id, "level": level, }, ) return commission_count def _credit_user( self, user: User, amount: Decimal, tx_type: str, metadata: dict ) -> None: """Credit a user's balance and log the transaction atomically.""" user_locked = User.objects.select_for_update().get(pk=user.pk) balance_before = user_locked.credit_balance user_locked.credit_balance = self._quantize(balance_before + amount) user_locked.save(update_fields=["credit_balance"]) Transaction.objects.create( user=user_locked, transaction_type=tx_type, amount=amount, status=Transaction.STATUS_COMPLETED, balance_before=balance_before, balance_after=user_locked.credit_balance, data=metadata, ) def _quantize(self, value: Decimal | float | int) -> Decimal: return Decimal(value).quantize(DECIMAL_QUANTIZE, rounding=ROUND_HALF_UP)