""" 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, etc.) for variable costs Returns: int: Number of credits required """ base_cost = CREDIT_COSTS.get(operation_type, 0) # Variable costs based on amount 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))) return base_cost @staticmethod def check_credits(account, required_credits=None, operation_type=None, amount=None): """ Check if account has enough credits. Args: account: Account instance required_credits: Number of credits required (legacy parameter) operation_type: Type of operation (new parameter) amount: Optional amount for variable costs (new parameter) Raises: InsufficientCreditsError: If account doesn't have enough credits """ # Support both old and new API if operation_type: required_credits = CreditService.get_credit_cost(operation_type, amount) elif required_credits is None: raise ValueError("Either required_credits or operation_type must be provided") 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 CreditService.check_credits(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 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. DEPRECATED: Use get_credit_cost() instead. Kept for backward compatibility. Args: operation_type: Type of operation **kwargs: Operation-specific parameters Returns: int: Number of credits required Raises: CreditCalculationError: If calculation fails """ # Map old operation types to new ones operation_mapping = { 'ideas': 'idea_generation', 'content': 'content_generation', 'images': 'image_generation', 'reparse': 'image_prompt_extraction', } mapped_type = operation_mapping.get(operation_type, operation_type) # Handle variable costs if mapped_type == 'content_generation': word_count = kwargs.get('word_count') or kwargs.get('content_count', 1000) * 100 return CreditService.get_credit_cost(mapped_type, word_count) elif mapped_type == 'clustering': keyword_count = kwargs.get('keyword_count', 0) # Clustering is fixed cost per request return CreditService.get_credit_cost(mapped_type) elif mapped_type == 'idea_generation': idea_count = kwargs.get('idea_count', 1) # Fixed cost per request return CreditService.get_credit_cost(mapped_type) elif mapped_type == 'image_generation': image_count = kwargs.get('image_count', 1) return CreditService.get_credit_cost(mapped_type) * image_count return CreditService.get_credit_cost(mapped_type)