quantum-ai/authentication/models.py

from django.contrib.auth.models import AbstractUser
from django.db import models
from decimal import Decimal


class User(AbstractUser):
    email = models.EmailField(unique=True)
    wallet_balance = models.DecimalField(max_digits=10, decimal_places=2, default=Decimal('0.00'))
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    
    USERNAME_FIELD = 'email'
    REQUIRED_FIELDS = ['username']
    
    def __str__(self):
        return self.email
    
    def has_sufficient_balance(self, amount):
        return self.wallet_balance >= Decimal(str(amount))
    
    def deduct_balance(self, amount, description="", agent_slug=""):
        if self.has_sufficient_balance(amount):
            self.wallet_balance -= Decimal(str(amount))
            self.save()
            
            # Create transaction record
            from wallet.models import WalletTransaction
            WalletTransaction.objects.create(
                user=self,
                amount=-Decimal(str(amount)),
                type='agent_usage',
                description=description,
                agent_slug=agent_slug
            )
            return True
        return False
    
    def add_balance(self, amount, description="", stripe_session_id=""):
        self.wallet_balance += Decimal(str(amount))
        self.save()
        
        # Create transaction record
        from wallet.models import WalletTransaction
        WalletTransaction.objects.create(
            user=self,
            amount=Decimal(str(amount)),
            type='top_up',
            description=description,
            stripe_session_id=stripe_session_id
        )