import pytest
from asgiref.sync import sync_to_async
from django.conf import settings

from promptgame.gameui.attack import get_transfer_charges
from promptgame.tests.helpers import async_make_account


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_get_transfer_amounts() -> None:
    victim_account = await async_make_account()
    attacker_account = await async_make_account()

    # Make the balance a prime number so that we can test rounding
    victim_balance = 9601
    shadow_balance = int(settings.SHADOW_BALANCE_CAP / 2) + 1
    victim_account.balance = victim_balance
    victim_account.active_defense.post_prompt = ""
    attacker_account.shadow_balance = shadow_balance
    xfer_charges = await sync_to_async(get_transfer_charges)(
        attacking_account=attacker_account, defending_account=victim_account
    )

    # Check that they're both ints
    assert isinstance(xfer_charges.steal_amount, int)
    assert isinstance(xfer_charges.shadow_transfer, int)

    # Check that they're both nonnegative
    assert xfer_charges.steal_amount >= 0
    assert xfer_charges.shadow_transfer >= 0

    # Check that you make more money if you the victim has a post_prompt
    # The shadow_transfer charge shouldn't change
    victim_account.active_defense.post_prompt = "foo"
    attacker_account.shadow_balance = shadow_balance
    charges_with_post_prompt = await sync_to_async(get_transfer_charges)(
        attacking_account=attacker_account, defending_account=victim_account
    )
    assert charges_with_post_prompt.steal_amount > xfer_charges.steal_amount
    assert charges_with_post_prompt.shadow_transfer == xfer_charges.shadow_transfer

    # This part is super specific to the exact algorithm; we might need to change this if we modify how we calculate transfers.
    assert xfer_charges.shadow_transfer == int(
        shadow_balance * settings.SHADOW_BALANCE_TRANSFER_FRACTION
    )
    assert xfer_charges.steal_amount == int(
        victim_balance * settings.TRANSFER_CHARGE_FRACTION_WITHOUT_POST_PROMPT
    )
    assert charges_with_post_prompt.steal_amount == int(
        victim_balance * settings.TRANSFER_CHARGE_FRACTION_WITH_POST_PROMPT
    )
