igny8/backend/igny8_core/tasks/plan_limits.py

"""
Plan Limits Reset Task
Scheduled task to reset monthly plan limits at billing period end
"""
from celery import shared_task
from django.utils import timezone
from datetime import timedelta
import logging

from igny8_core.auth.models import Account
from igny8_core.business.billing.services.limit_service import LimitService

logger = logging.getLogger(__name__)


@shared_task(name='reset_monthly_plan_limits')
def reset_monthly_plan_limits():
    """
    Reset monthly plan limits for accounts whose billing period has ended.
    
    This task should run daily (recommended: midnight UTC).
    It finds all accounts where the billing period has ended and resets their monthly usage.
    
    Monthly limits that get reset:
    - content_ideas
    - content_words
    - images_basic
    - images_premium  
    - image_prompts
    
    Hard limits (sites, users, keywords, clusters) are NOT reset.
    """
    logger.info("Starting monthly plan limits reset task")
    
    today = timezone.now().date()
    reset_count = 0
    error_count = 0
    
    # Find all active accounts with subscriptions
    accounts = Account.objects.filter(
        status='active',
        subscription__isnull=False
    ).select_related('subscription', 'plan')
    
    logger.info(f"Found {accounts.count()} active accounts with subscriptions")
    
    for account in accounts:
        try:
            subscription = account.subscription
            
            # Check if billing period has ended
            if subscription.current_period_end and subscription.current_period_end.date() <= today:
                logger.info(f"Resetting limits for account {account.id} ({account.name}) - "
                           f"period ended {subscription.current_period_end.date()}")
                
                # Reset monthly limits
                result = LimitService.reset_monthly_limits(account)
                
                # Update subscription period
                from dateutil.relativedelta import relativedelta
                
                # Calculate new period based on billing cycle
                plan = account.plan
                if plan.billing_cycle == 'monthly':
                    new_period_start = subscription.current_period_end + timedelta(days=1)
                    new_period_end = new_period_start + relativedelta(months=1) - timedelta(days=1)
                elif plan.billing_cycle == 'annual':
                    new_period_start = subscription.current_period_end + timedelta(days=1)
                    new_period_end = new_period_start + relativedelta(years=1) - timedelta(days=1)
                else:
                    # Default to monthly
                    new_period_start = subscription.current_period_end + timedelta(days=1)
                    new_period_end = new_period_start + relativedelta(months=1) - timedelta(days=1)
                
                # Update subscription
                subscription.current_period_start = new_period_start
                subscription.current_period_end = new_period_end
                subscription.save(update_fields=['current_period_start', 'current_period_end'])
                
                reset_count += 1
                
                logger.info(f"Reset complete for account {account.id}: "
                           f"New period {new_period_start} to {new_period_end}")
                
        except Exception as e:
            error_count += 1
            logger.error(f"Error resetting limits for account {account.id}: {str(e)}", exc_info=True)
    
    logger.info(f"Monthly plan limits reset task complete: "
               f"{reset_count} accounts reset, {error_count} errors")
    
    return {
        'reset_count': reset_count,
        'error_count': error_count,
        'total_accounts': accounts.count(),
    }


@shared_task(name='check_approaching_limits')
def check_approaching_limits(threshold_percentage=80):
    """
    Check for accounts approaching their plan limits and send notifications.
    
    Args:
        threshold_percentage: Percentage at which to trigger warning (default 80%)
    
    This task should run daily.
    It checks both hard and monthly limits and sends warnings when usage exceeds threshold.
    """
    logger.info(f"Starting limit warning check task (threshold: {threshold_percentage}%)")
    
    warning_count = 0
    
    # Find all active accounts
    accounts = Account.objects.filter(status='active').select_related('plan')
    
    for account in accounts:
        try:
            if not account.plan:
                continue
            
            # Get usage summary
            summary = LimitService.get_usage_summary(account)
            
            warnings = []
            
            # Check hard limits
            for limit_type, data in summary.get('hard_limits', {}).items():
                if data['percentage_used'] >= threshold_percentage:
                    warnings.append({
                        'type': 'hard',
                        'limit_type': limit_type,
                        'display_name': data['display_name'],
                        'current': data['current'],
                        'limit': data['limit'],
                        'percentage': data['percentage_used'],
                    })
            
            # Check monthly limits
            for limit_type, data in summary.get('monthly_limits', {}).items():
                if data['percentage_used'] >= threshold_percentage:
                    warnings.append({
                        'type': 'monthly',
                        'limit_type': limit_type,
                        'display_name': data['display_name'],
                        'current': data['current'],
                        'limit': data['limit'],
                        'percentage': data['percentage_used'],
                        'resets_in_days': summary.get('days_until_reset', 0),
                    })
            
            if warnings:
                warning_count += 1
                logger.info(f"Account {account.id} ({account.name}) has {len(warnings)} limit warnings")
                
                # TODO: Send email notification
                # from igny8_core.business.billing.services.email_service import send_limit_warning_email
                # send_limit_warning_email(account, warnings)
                
        except Exception as e:
            logger.error(f"Error checking limits for account {account.id}: {str(e)}", exc_info=True)
    
    logger.info(f"Limit warning check complete: {warning_count} accounts with warnings")
    
    return {
        'warning_count': warning_count,
        'total_accounts': accounts.count(),
    }