""" Credit Service for managing credit transactions and deductions """ from django.db import transaction from django.utils import timezone from .models import CreditTransaction, CreditUsageLog from .constants import CREDIT_COSTS from .exceptions import InsufficientCreditsError, CreditCalculationError from igny8_core.auth.models import Account class CreditService: """Service for managing credits""" @staticmethod def get_credit_cost(operation_type, amount=None): """ Get credit cost for operation. Args: operation_type: Type of operation (from CREDIT_COSTS) amount: Optional amount (word count, image count, etc.) Returns: int: Number of credits required Raises: CreditCalculationError: If operation type is unknown """ base_cost = CREDIT_COSTS.get(operation_type, 0) if base_cost == 0: raise CreditCalculationError(f"Unknown operation type: {operation_type}") # Variable cost operations if operation_type == 'content_generation' and amount: # Per 100 words return max(1, int(base_cost * (amount / 100))) elif operation_type == 'optimization' and amount: # Per 200 words return max(1, int(base_cost * (amount / 200))) elif operation_type == 'image_generation' and amount: # Per image return base_cost * amount elif operation_type == 'idea_generation' and amount: # Per idea return base_cost * amount # Fixed cost operations return base_cost @staticmethod def check_credits(account, operation_type, amount=None): """ Check if account has sufficient credits for an operation. Args: account: Account instance operation_type: Type of operation amount: Optional amount (word count, image count, etc.) Raises: InsufficientCreditsError: If account doesn't have enough credits """ required = CreditService.get_credit_cost(operation_type, amount) if account.credits < required: raise InsufficientCreditsError( f"Insufficient credits. Required: {required}, Available: {account.credits}" ) return True @staticmethod def check_credits_legacy(account, required_credits): """ Legacy method: Check if account has enough credits (for backward compatibility). Args: account: Account instance required_credits: Number of credits required Raises: InsufficientCreditsError: If account doesn't have enough credits """ if account.credits < required_credits: raise InsufficientCreditsError( f"Insufficient credits. Required: {required_credits}, Available: {account.credits}" ) @staticmethod @transaction.atomic def deduct_credits(account, amount, operation_type, description, metadata=None, cost_usd=None, model_used=None, tokens_input=None, tokens_output=None, related_object_type=None, related_object_id=None): """ Deduct credits and log transaction. Args: account: Account instance amount: Number of credits to deduct operation_type: Type of operation (from CreditUsageLog.OPERATION_TYPE_CHOICES) description: Description of the transaction metadata: Optional metadata dict cost_usd: Optional cost in USD model_used: Optional AI model used tokens_input: Optional input tokens tokens_output: Optional output tokens related_object_type: Optional related object type related_object_id: Optional related object ID Returns: int: New credit balance """ # Check sufficient credits (legacy: amount is already calculated) CreditService.check_credits_legacy(account, amount) # Deduct from account.credits account.credits -= amount account.save(update_fields=['credits']) # Create CreditTransaction CreditTransaction.objects.create( account=account, transaction_type='deduction', amount=-amount, # Negative for deduction balance_after=account.credits, description=description, metadata=metadata or {} ) # Create CreditUsageLog CreditUsageLog.objects.create( account=account, operation_type=operation_type, credits_used=amount, cost_usd=cost_usd, model_used=model_used or '', tokens_input=tokens_input, tokens_output=tokens_output, related_object_type=related_object_type or '', related_object_id=related_object_id, metadata=metadata or {} ) return account.credits @staticmethod @transaction.atomic def deduct_credits_for_operation(account, operation_type, amount=None, description=None, metadata=None, cost_usd=None, model_used=None, tokens_input=None, tokens_output=None, related_object_type=None, related_object_id=None): """ Deduct credits for an operation (convenience method that calculates cost automatically). Args: account: Account instance operation_type: Type of operation amount: Optional amount (word count, image count, etc.) description: Optional description (auto-generated if not provided) metadata: Optional metadata dict cost_usd: Optional cost in USD model_used: Optional AI model used tokens_input: Optional input tokens tokens_output: Optional output tokens related_object_type: Optional related object type related_object_id: Optional related object ID Returns: int: New credit balance """ # Calculate credit cost credits_required = CreditService.get_credit_cost(operation_type, amount) # Check sufficient credits CreditService.check_credits(account, operation_type, amount) # Auto-generate description if not provided if not description: if operation_type == 'clustering': description = f"Clustering operation" elif operation_type == 'idea_generation': description = f"Generated {amount or 1} idea(s)" elif operation_type == 'content_generation': description = f"Generated content ({amount or 0} words)" elif operation_type == 'image_generation': description = f"Generated {amount or 1} image(s)" else: description = f"{operation_type} operation" return CreditService.deduct_credits( account=account, amount=credits_required, operation_type=operation_type, description=description, metadata=metadata, cost_usd=cost_usd, model_used=model_used, tokens_input=tokens_input, tokens_output=tokens_output, related_object_type=related_object_type, related_object_id=related_object_id ) @staticmethod @transaction.atomic def add_credits(account, amount, transaction_type, description, metadata=None): """ Add credits (purchase, subscription, etc.). Args: account: Account instance amount: Number of credits to add transaction_type: Type of transaction (from CreditTransaction.TRANSACTION_TYPE_CHOICES) description: Description of the transaction metadata: Optional metadata dict Returns: int: New credit balance """ # Add to account.credits account.credits += amount account.save(update_fields=['credits']) # Create CreditTransaction CreditTransaction.objects.create( account=account, transaction_type=transaction_type, amount=amount, # Positive for addition balance_after=account.credits, description=description, metadata=metadata or {} ) return account.credits @staticmethod def calculate_credits_for_operation(operation_type, **kwargs): """ Calculate credits needed for an operation. Legacy method - use get_credit_cost() instead. Args: operation_type: Type of operation **kwargs: Operation-specific parameters Returns: int: Number of credits required Raises: CreditCalculationError: If calculation fails """ # Map legacy operation types if operation_type == 'ideas': operation_type = 'idea_generation' elif operation_type == 'content': operation_type = 'content_generation' elif operation_type == 'images': operation_type = 'image_generation' # Extract amount from kwargs amount = None if 'word_count' in kwargs: amount = kwargs.get('word_count') elif 'image_count' in kwargs: amount = kwargs.get('image_count') elif 'idea_count' in kwargs: amount = kwargs.get('idea_count') return CreditService.get_credit_cost(operation_type, amount)