igny8/backend/igny8_core/modules/billing/services.py

"""
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)