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'), db_index=True) created_at = models.DateTimeField(auto_now_add=True, db_index=True) updated_at = models.DateTimeField(auto_now=True) USERNAME_FIELD = 'email' REQUIRED_FIELDS = ['username'] class Meta: indexes = [ models.Index(fields=['email', 'wallet_balance']), models.Index(fields=['created_at', 'wallet_balance']), models.Index(fields=['-created_at']), ] 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 )