from django.core.management.base import BaseCommand
from django.db import transaction
from apps.accounts.models import User
from apps.accounts.models import UserKYC, UserShippingAddress
from apps.business.wallet.models import Wallet, WalletLedger
from apps.business.withdrawals.models import WithdrawalRequest, BankAccount
from apps.business.orders.models import Order
from apps.payments.models import PaymentTransaction
from apps.business.distributor.models import DistributorID, BrandAmbassadorEligibility

EXEMPT_MOBILES = ["8668192080"]


class Command(BaseCommand):
    help = "Delete all users and related data except for exempted mobile numbers."

    def handle(self, *args, **options):
        BATCH_SIZE = 500
        with transaction.atomic():
            users_to_delete = User.objects.exclude(mobile__in=EXEMPT_MOBILES)
            user_ids = list(users_to_delete.values_list("id", flat=True))

            self.stdout.write(f"Deleting related data for {len(user_ids)} users in batches...")

            def batched_delete(qs, id_field="user_id"):
                for i in range(0, len(user_ids), BATCH_SIZE):
                    batch_ids = user_ids[i:i+BATCH_SIZE]
                    filter_kwargs = {f"{id_field}__in": batch_ids}
                    qs.filter(**filter_kwargs).delete()

            batched_delete(UserKYC.objects)
            batched_delete(UserShippingAddress.objects)
            batched_delete(WalletLedger.objects)
            batched_delete(Wallet.objects, id_field="user_id")
            batched_delete(WithdrawalRequest.objects)
            batched_delete(BankAccount.objects)
            batched_delete(Order.objects)
            batched_delete(PaymentTransaction.objects)
            batched_delete(DistributorID.objects)
            batched_delete(BrandAmbassadorEligibility.objects)

            # Finally, delete the users in batches
            for i in range(0, len(user_ids), BATCH_SIZE):
                batch_ids = user_ids[i:i+BATCH_SIZE]
                User.objects.filter(id__in=batch_ids).delete()

            self.stdout.write(self.style.SUCCESS(f"Deleted all users and related data except exempted mobiles."))