from decimal import Decimal

from django.core.management.base import BaseCommand
from django.db import transaction
from django.db.models import Case, DecimalField, Sum, Value, When

from apps.business.wallet.models import Wallet, WalletLedger


class Command(BaseCommand):
    help = "Recalculate wallet main/repurchase balances from WalletLedger and update mismatches"

    def add_arguments(self, parser):
        parser.add_argument("--batch-size", type=int, default=2000)
        parser.add_argument("--dry-run", action="store_true", help="Compute mismatches without updating balances")

    def handle(self, *args, **options):
        batch_size = max(1, int(options["batch_size"]))
        dry_run = bool(options["dry_run"])

        self.stdout.write(self.style.NOTICE("Aggregating wallet ledger totals..."))

        main_agg = WalletLedger.objects.filter(wallet_type=WalletLedger.WalletType.EARNINGS).values("user_id").annotate(
            credits=Sum(
                Case(
                    When(transaction_type=WalletLedger.TransactionType.CREDIT, then="amount"),
                    default=Value(0),
                    output_field=DecimalField(max_digits=18, decimal_places=2),
                )
            ),
            debits=Sum(
                Case(
                    When(transaction_type=WalletLedger.TransactionType.DEBIT, then="amount"),
                    default=Value(0),
                    output_field=DecimalField(max_digits=18, decimal_places=2),
                )
            ),
        )
        repurchase_agg = WalletLedger.objects.filter(wallet_type=WalletLedger.WalletType.REPURCHASE).values("user_id").annotate(
            credits=Sum(
                Case(
                    When(transaction_type=WalletLedger.TransactionType.CREDIT, then="amount"),
                    default=Value(0),
                    output_field=DecimalField(max_digits=18, decimal_places=2),
                )
            ),
            debits=Sum(
                Case(
                    When(transaction_type=WalletLedger.TransactionType.DEBIT, then="amount"),
                    default=Value(0),
                    output_field=DecimalField(max_digits=18, decimal_places=2),
                )
            ),
        )

        zero = Decimal("0.00")
        main_map = {row["user_id"]: (row["credits"] or zero) - (row["debits"] or zero) for row in main_agg}
        repurchase_map = {row["user_id"]: (row["credits"] or zero) - (row["debits"] or zero) for row in repurchase_agg}
        user_ids = set(main_map.keys()) | set(repurchase_map.keys())

        existing_wallets = {w.user_id: w for w in Wallet.objects.filter(user_id__in=user_ids)}

        to_create = []
        for user_id in user_ids:
            if user_id not in existing_wallets:
                to_create.append(
                    Wallet(
                        user_id=user_id,
                        main_balance=main_map.get(user_id, zero),
                        repurchase_balance=repurchase_map.get(user_id, zero),
                    )
                )

        updated_count = 0
        created_count = len(to_create)
        mismatched = 0
        update_rows = []

        for user_id, wallet in existing_wallets.items():
            expected_main = main_map.get(user_id, zero)
            expected_repurchase = repurchase_map.get(user_id, zero)
            if wallet.main_balance != expected_main or wallet.repurchase_balance != expected_repurchase:
                mismatched += 1
                wallet.main_balance = expected_main
                wallet.repurchase_balance = expected_repurchase
                update_rows.append(wallet)

        if not dry_run:
            with transaction.atomic():
                if to_create:
                    Wallet.objects.bulk_create(to_create, batch_size=batch_size)
                if update_rows:
                    Wallet.objects.bulk_update(update_rows, ["main_balance", "repurchase_balance"], batch_size=batch_size)
            updated_count = len(update_rows)

        self.stdout.write(
            self.style.SUCCESS(
                f"Wallet reconciliation complete. users_with_ledger={len(user_ids)} mismatched={mismatched} created={created_count} updated={updated_count} dry_run={dry_run}"
            )
        )