import random
from datetime import datetime

from django.core.management.base import BaseCommand

from promptgame.gameui import llm_wrappers
from promptgame.gameui.models import Account, AuthAttempt

MIN_BALANCE = 10000
MAX_BALANCE = 100000


class Command(BaseCommand):
    help = "Populate the auth attempt table"

    def add_arguments(self, parser):
        parser.add_argument(
            "num_attempts", type=int, help="The number of auth attempts to be created"
        )

    def handle(self, *args, **kwargs):
        num_attempts = kwargs["num_attempts"]

        accounts = Account.objects.all()

        if accounts.count() < 2:
            self.stdout.write(
                self.style.ERROR(
                    "At least 2 accounts are required to create auth attempts."
                )
            )
            return

        for _ in range(num_attempts):
            attacking_account = random.choice(accounts)
            defending_account = random.choice(
                [acc for acc in accounts if acc != attacking_account]
            )
            was_successful = random.choice([True, False])
            attacking_transfer_amount = random.randint(MIN_BALANCE, MAX_BALANCE)
            defending_transfer_amount = random.randint(MIN_BALANCE, MAX_BALANCE)

            auth_attempt = AuthAttempt(
                defending_account=defending_account,
                attacking_account=attacking_account,
                timestamp=datetime.now(),
                was_successful=was_successful,
                attacking_transfer_amount=attacking_transfer_amount,
                defending_transfer_amount=defending_transfer_amount,
                attacking_balance=attacking_account.balance,
                defending_balance=defending_account.balance,
                llm_choice=random.choice(list(llm_wrappers.LLMChoice)).name,
            )
            auth_attempt.save()

        self.stdout.write(self.style.SUCCESS(f"{num_attempts} auth attempts created"))
