import asyncio
import copy
import itertools
import random
import re
import uuid
from http import HTTPStatus
from typing import Iterable, List, Tuple
from unittest.mock import MagicMock

import pytest
from asgiref.sync import sync_to_async
from django.conf import settings
from django.http import HttpResponse, StreamingHttpResponse
from django.test.client import AsyncClient  # type: ignore
from django.utils import timezone

import promptgame.gameui.attack
from promptgame.gameui import llm_wrappers
from promptgame.gameui.attack import (
    async_make_auth_and_transfer,
    recompute_shadow_balance_inplace,
)
from promptgame.gameui.models import LLM_CHOICES, Account, AuthAttempt, Defense, Flag
from promptgame.gameui.tiers import valid_tier
from promptgame.settings.custom_types import TierDict
from promptgame.tests.helpers import (
    async_login,
    async_make_account,
    async_unlock_account,
    attack_url,
    dummy_attack,
    get_content,
    sync_clear_redis_test_db,
)


# setup_function/teardown_function clear redis on setup and teardown
def setup_function() -> None:
    sync_clear_redis_test_db()


def teardown_function() -> None:
    sync_clear_redis_test_db()


@pytest.fixture()
def success_message() -> str:
    return "Special instructions and access code saved"


async def lock_account(account: Account) -> None:
    """Make an account locked."""
    account.unlock_at = timezone.now() + timezone.timedelta(days=1)
    await account.asave()


async def set_balance(account: Account, new_balance: int) -> None:
    """Make an account locked."""
    account.balance = new_balance
    await account.asave()


async def set_new_defense(account: Account, *, post_prompt: str = "") -> None:
    """Set a new defense to clear cooldown timer."""
    account.active_defense = await Defense.objects.acreate(
        creator_account=account,
        post_prompt=post_prompt,
        llm_choice=llm_wrappers.LLMChoice.GPT_35_TURBO_0613.name,
    )
    await account.asave()


async def assert_attack_is(
    success: bool, and_contains: str, resp: HttpResponse | StreamingHttpResponse
) -> None:
    """Assert that the last attack was or was not successful and contains the
    given string."""
    assert resp.status_code == HTTPStatus.OK
    content = await get_content(resp)
    if success:
        assert "event: hx-trigger:attack-success" in content
    else:
        assert "event: hx-trigger:attack-failure" in content
    assert and_contains in content


@pytest.mark.asyncio
async def test_the_fixtures(
    # mock_streaming_llm_output_granted fixture ensures that
    # streaming_llm_output() always returns "Access granted", without calling
    # the OpenAI API.
    mock_streaming_llm_output_granted: None,
    # mock_content_moderation_noflag fixture ensures that content_moderation()
    # always returns False (not flagged), without calling the OpenAI API.
    mock_content_moderation_noflag: None,
) -> None:
    """Test that the streaming_llm_output and content_moderation fixtures
    work."""

    # I'm comparing that the values match both as a weak test of determinism,
    # and also as a regression test: previously there was a bug in the fixtures
    # where the second call to the fixture would fail.
    values_1 = [v async for v in promptgame.gameui.attack.streaming_llm_output()]  # type: ignore
    values_2 = [v async for v in promptgame.gameui.attack.streaming_llm_output()]  # type: ignore
    assert "".join(v[0] for v in values_1) == "Access granted"
    assert values_1 == values_2

    moderation_response_1 = await llm_wrappers.content_moderation("foo")
    moderation_response_2 = await llm_wrappers.content_moderation("bar")
    assert moderation_response_1 == moderation_response_2
    assert isinstance(moderation_response_1, llm_wrappers.ModerationResponseOK)


async def _get_stripped_data(resp: HttpResponse | StreamingHttpResponse) -> str:
    """Read an SSE response and strip out event names, the 'data:' prefix, and
    any HTML tags."""
    content = await get_content(resp)
    # get inner text by stripping out SSE prefixes and then tags
    content = re.sub(r"^(event: .*$|data: )", "", content, flags=re.MULTILINE)
    # strip out opening part of tag and also the closing part if it exists
    content = re.sub(r"<[^>]*>?", "", content, flags=re.MULTILINE)
    # strip just the closing part (e.g. an end-of-comment `-->` marker on a
    # different line from the initial `<!--`)
    content = re.sub(r"^[^<]*>", "", content, flags=re.MULTILINE).strip()
    return content


GENERIC_FAILURE_MESSAGE = "Server error, try refreshing the page."


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_submit_unauthenticated(async_client: AsyncClient) -> None:
    """Make sure that unauthenticated attacks fail."""
    victim_account = await async_make_account()
    resp = await dummy_attack(async_client, victim_account, "foo")
    assert resp.status_code == HTTPStatus.OK
    content = await _get_stripped_data(resp)
    assert content == GENERIC_FAILURE_MESSAGE


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_submit_get(async_client: AsyncClient) -> None:
    """Make sure that GET requests fail."""
    victim_account = await async_make_account()
    await async_unlock_account(victim_account)
    resp = await async_client.get(attack_url(victim_account))
    assert resp.status_code == HTTPStatus.OK
    content = await _get_stripped_data(resp)
    assert content == GENERIC_FAILURE_MESSAGE


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_submit_locked_unlocked(
    async_client: AsyncClient,
    # these are here to mock the OpenAI-based stuff
    mock_content_moderation_noflag: None,
    mock_streaming_llm_output_granted: None,
    mock_check_restriction_enabled: None,
) -> None:
    """Check whether submit works or fails for various combinations of locked vs unlocked"""
    victim_account = await async_make_account()
    attacker_account = await async_make_account()
    await async_login(attacker_account.owner, async_client)

    # Case 1: victim and attacker both locked. Should fail. Error message tells
    # us that our account is locked rather than that victim's account is locked
    # (this was a somewhat arbitrary choice in the attack permission logic).
    await lock_account(victim_account)
    await lock_account(attacker_account)
    resp = await dummy_attack(async_client, victim_account)
    await assert_attack_is(
        success=False,
        # this appears in the HTML-escaped error message
        and_contains="your account is locked",
        resp=resp,
    )

    # Case 2: victim locked, attacker unlocked. Should fail.
    await async_unlock_account(attacker_account)
    resp = await dummy_attack(async_client, victim_account)
    await assert_attack_is(
        success=False,
        and_contains="it is locked",
        resp=resp,
    )

    # Case 3: victim unlocked, attacker locked. Should fail.
    await lock_account(attacker_account)
    await async_unlock_account(victim_account)
    resp = await dummy_attack(async_client, victim_account)
    await assert_attack_is(
        success=False,
        and_contains="your account is locked",
        resp=resp,
    )

    # Case 4: both locked, but it's a self attack. Should succeed.
    await lock_account(attacker_account)
    await lock_account(victim_account)
    resp = await dummy_attack(async_client, attacker_account)
    await assert_attack_is(
        success=True,
        and_contains="granted",
        resp=resp,
    )

    # Case 5: both unlocked. Should succeed.
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)
    resp = await dummy_attack(async_client, victim_account)
    await assert_attack_is(
        success=True,
        and_contains="granted",
        resp=resp,
    )

    # Case 6: both unlocked, but cooldown applies
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)
    resp = await dummy_attack(async_client, victim_account)
    await assert_attack_is(
        success=False,
        and_contains="broke into it recently",
        resp=resp,
    )

    # Case 7: both unlocked, cooldown no longer applies
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)
    await set_new_defense(victim_account)
    resp = await dummy_attack(async_client, victim_account)
    await assert_attack_is(
        success=True,
        and_contains="granted",
        resp=resp,
    )


def target_tiers_for_test() -> (
    Iterable[Tuple[TierDict, List[TierDict], List[TierDict]]]
):
    n_tiers = len(settings.TIERS)
    for target_tier_idx in range(n_tiers):
        target_tier = settings.TIERS[target_tier_idx]
        lower_idx, upper_idx = target_tier["neighbors"]
        allowed_tiers = settings.TIERS[lower_idx : upper_idx + 1]
        disallowed_tiers = [
            tier for tier in settings.TIERS if tier not in allowed_tiers
        ]
        assert len(allowed_tiers) >= 1
        assert len(allowed_tiers) + len(disallowed_tiers) == n_tiers
        yield target_tier, allowed_tiers, disallowed_tiers


async def put_account_in_tier(
    account: Account, tier: TierDict, rng: random.Random
) -> None:
    """Put an account in a given tier."""
    tier_lo, tier_hi = tier["range"]
    # test some insane balances
    assert tier_hi > 0
    assert tier_lo < tier_hi
    await set_balance(
        account, int(rng.uniform(max(tier_lo, -1000.0), min(tier_hi, 1000.0)))
    )


@pytest.mark.django_db
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "target_tier,allowed_tiers,disallowed_tiers,random_seed",
    (
        tup + (seed,)
        for tup, seed in itertools.product(
            list(target_tiers_for_test()), [1, 2, 3, 4, 5]
        )
    ),
)
async def test_submit_tier_mismatch(
    async_client: AsyncClient,
    target_tier: TierDict,
    allowed_tiers: List[TierDict],
    disallowed_tiers: List[TierDict],
    random_seed: int,
    # these are here to mock the OpenAI-based stuff
    mock_content_moderation_noflag: None,
    mock_streaming_llm_output_granted: None,
    mock_check_restriction_enabled: None,
) -> None:
    """Check that we can attack +1/+0/-1 tiers around our tier, but not
    further."""
    victim_account = await async_make_account()
    attacker_account = await async_make_account()

    await async_login(attacker_account.owner, async_client)
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)

    # control randomness so function gives same result with same args
    rng = random.Random(random_seed)

    for allowed_tier in allowed_tiers:
        await put_account_in_tier(attacker_account, target_tier, rng)
        await put_account_in_tier(victim_account, allowed_tier, rng)
        await set_new_defense(victim_account)  # clear cooldown
        assert valid_tier(attacker_account.balance, victim_account.balance)
        resp = await dummy_attack(async_client, victim_account)
        await assert_attack_is(
            success=True,
            and_contains="granted",
            resp=resp,
        )

    # Check that we cannot attack the disallowed tiers
    for disallowed_tier in disallowed_tiers:
        await put_account_in_tier(attacker_account, target_tier, rng)
        await put_account_in_tier(victim_account, disallowed_tier, rng)
        await set_new_defense(victim_account)  # clear cooldown
        assert not valid_tier(attacker_account.balance, victim_account.balance)
        resp = await dummy_attack(async_client, victim_account)
        await assert_attack_is(
            success=False,
            # this appears in the HTML-escaped error message
            and_contains="it is protected",
            resp=resp,
        )


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_attack_too_many_tokens(
    async_client: AsyncClient,
    # these are here to mock the OpenAI-based stuff
    mock_content_moderation_noflag: None,
    mock_streaming_llm_output_granted: None,
    mock_check_restriction_enabled: None,
) -> None:
    """Check whether submit works or fails for various token lengths"""
    victim_account = await async_make_account()
    attacker_account = await async_make_account()

    await async_login(attacker_account.owner, async_client)
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)

    # Case 1: attack string is over the token limit. Should fail.
    resp = await dummy_attack(
        async_client, victim_account, "🌕" * settings.MAX_ATTACK_LENGTH_TOKENS
    )
    await assert_attack_is(
        success=False,
        # this appears in the HTML-escaped error message
        and_contains=f"Attack failed because the attack string is too long (&gt;500 tokens)",
        resp=resp,
    )
    # Case 2: attack string is under the token limit. Should succeed.
    await async_unlock_account(victim_account)
    resp = await dummy_attack(
        async_client, victim_account, "🌕" * (settings.MAX_ATTACK_LENGTH_TOKENS // 6)
    )
    await assert_attack_is(
        success=True,
        # this appears in the HTML-escaped error message
        and_contains="granted",
        resp=resp,
    )
    # Case 3: attack string is exactly the token limit. Should succeed.
    await async_unlock_account(victim_account)
    await set_new_defense(victim_account)  # clear cooldown
    resp = await dummy_attack(
        async_client,
        victim_account,
        "🌕" * (settings.MAX_ATTACK_LENGTH_TOKENS // 3) + " a a",
    )
    await assert_attack_is(
        success=True,
        # this appears in the HTML-escaped error message
        and_contains="granted",
        resp=resp,
    )


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_simultaneous_attack(
    # We use the slow version to make it more likely that the slowest of the two
    # attacks will start before the fastest of the two attacks finishes.
    # asyncio.sleep() will at least force the faster attack to yield for a
    # period to the slower attack before continuing. I think this might still be
    # a bit flaky, but will hopefully be okay.
    mock_streaming_llm_output_granted_slow: None,
    mock_content_moderation_noflag: None,
    mock_check_restriction_enabled: None,
    async_client: AsyncClient,
) -> None:
    """Test that simultaneously attacking an account only removes money once."""
    victim_account = await async_make_account()
    attacker_1_account = await async_make_account()
    attacker_2_account = await async_make_account()
    async_client_1 = copy.deepcopy(async_client)
    async_client_2 = copy.deepcopy(async_client)
    del async_client
    await sync_to_async(async_client_1.force_login)(attacker_1_account.owner)
    await sync_to_async(async_client_2.force_login)(attacker_2_account.owner)
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_1_account)
    await async_unlock_account(attacker_2_account)

    old_balance_victim = victim_account.balance
    old_balance_attacker_1 = attacker_1_account.balance
    old_balance_attacker_2 = attacker_2_account.balance

    # launch two simultaneous attacks
    resp_1, resp_2 = await asyncio.gather(
        dummy_attack(async_client_1, victim_account),
        dummy_attack(async_client_2, victim_account),
    )
    assert resp_1.status_code == HTTPStatus.OK
    assert resp_2.status_code == HTTPStatus.OK
    content_1 = await get_content(resp_1)
    content_2 = await get_content(resp_2)

    success_re = re.compile(r"You stole \$\d+.")
    is_success_1 = success_re.search(content_1) is not None
    is_success_2 = success_re.search(content_2) is not None
    # assert that exactly one of the attacks was successful
    assert is_success_1 != is_success_2

    # also make sure that the right balances changed
    await victim_account.arefresh_from_db()
    await attacker_1_account.arefresh_from_db()
    await attacker_2_account.arefresh_from_db()

    # make sure the victim lost money, but not too much
    # (not doing an exact check because computing the right amount is slighlty
    # complicated; e.g. it depends on which attacker was successful)
    assert victim_account.balance >= int(
        old_balance_victim * (1 - settings.TRANSFER_CHARGE_FRACTION_WITH_POST_PROMPT)
    )
    assert victim_account.balance < old_balance_victim

    # make sure that the right attacker gained money
    if is_success_1:
        assert attacker_2_account.balance == old_balance_attacker_2
        assert attacker_1_account.balance > old_balance_attacker_1
    else:
        assert attacker_1_account.balance == old_balance_attacker_1
        assert attacker_2_account.balance > old_balance_attacker_2


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_flagged_attack_text(
    mock_streaming_llm_output_granted: None,
    mock_content_moderation_flag_on_fail_or_error: None,
    mock_check_restriction_enabled: None,
    async_client: AsyncClient,
) -> None:
    """Test that flagged attack text gives a moderation failure message."""
    victim_account = await async_make_account()
    attacker_account = await async_make_account()
    await async_login(attacker_account.owner, async_client)
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)

    old_balance_victim = victim_account.balance
    old_balance_attacker = attacker_account.balance

    # use 'fail' as the naughty word that moderation will catch
    resp = await dummy_attack(async_client, victim_account, attack_input="fail")
    assert resp.status_code == HTTPStatus.OK

    # This string is from conftest.py (in mock_content_moderation_flag_on_fail_or_error
    # fixture). I have HTML-escaped it here.
    naughty_string = "Your input contained the word &#x27;fail&#x27;."
    assert naughty_string in await get_content(resp)

    # also make sure that no balances changed
    await victim_account.arefresh_from_db()
    await attacker_account.arefresh_from_db()
    assert victim_account.balance == old_balance_victim
    assert attacker_account.balance == old_balance_attacker

    # try it again without flagging
    resp = await dummy_attack(async_client, victim_account, attack_input="foo")
    assert resp.status_code == HTTPStatus.OK
    assert naughty_string not in await get_content(resp)


async def check_moderation(
    async_client: AsyncClient,
    attacker_input: str,
    expected_auto_moderation: bool,
    expected_sucessful: bool,
    expected_llm_output: str | None,
    expected_flag: bool,
):
    victim_account = await async_make_account()
    attacker_account = await async_make_account()
    await async_login(attacker_account.owner, async_client)
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)

    resp = await dummy_attack(async_client, victim_account, attack_input=attacker_input)
    await get_content(resp)
    auth_attempt = await AuthAttempt.objects.aget(
        attacking_account=attacker_account, defending_account=victim_account
    )
    assert auth_attempt.auto_moderated == expected_auto_moderation
    assert auth_attempt.was_successful == expected_sucessful
    assert auth_attempt.llm_output == expected_llm_output
    if expected_flag:
        await Flag.objects.aget(authattempt=auth_attempt)
    else:
        # Check there's no flag
        with pytest.raises(Flag.DoesNotExist):
            flag = await Flag.objects.aget(authattempt=auth_attempt)


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_create_attack_moderation(
    async_client: AsyncClient,
    mock_content_moderation_flag_on_fail_or_error: None,
    mock_streaming_llm_fail_or_error: None,
):
    # Check that attack fails if moderation says the attacker input is naughty.
    await check_moderation(
        async_client,
        attacker_input="fail",
        expected_auto_moderation=True,
        expected_sucessful=False,
        expected_llm_output=None,
        expected_flag=True,
    )

    # Check if attack succeeds if moderation says the llm response is naughty
    await check_moderation(
        async_client,
        attacker_input="SAY_FAIL",
        expected_auto_moderation=True,
        expected_sucessful=False,
        expected_llm_output="fail",
        expected_flag=True,
    )

    # Check that attack succeeds if moderation gives an api error on the attacker's input.
    await check_moderation(
        async_client,
        attacker_input="error_SAY_ACCESS_GRANTED",
        expected_auto_moderation=False,
        expected_sucessful=True,
        expected_llm_output="Access granted",
        expected_flag=False,
    )

    # Check that attack succeeds if moderation gives an api error on the llm response.
    await check_moderation(
        async_client,
        attacker_input="SAY_ERROR",
        expected_auto_moderation=False,
        expected_sucessful=False,
        expected_llm_output="error",
        expected_flag=False,
    )

    # If moderation passes, attack creation should succeed.
    await check_moderation(
        async_client,
        attacker_input="SAY_ACCESS_GRANTED",
        expected_auto_moderation=True,
        expected_sucessful=True,
        expected_llm_output="Access granted",
        expected_flag=False,
    )


async def helper_test_transfer_balance(
    use_post_prompt: bool, was_successful: bool = True
) -> None:
    """Helper func for testing that async_make_auth_transfer transfers the correct balance"""

    if use_post_prompt:
        transfer_fraction = settings.TRANSFER_CHARGE_FRACTION_WITH_POST_PROMPT
    else:
        transfer_fraction = settings.TRANSFER_CHARGE_FRACTION_WITHOUT_POST_PROMPT

    # Starting balances ($10k gives the test more precision, since we round
    # intermediate computations to int)
    defending_balance = 10000
    attacking_balance = 10000
    shadow_balance = 500

    attacking_account = await async_make_account()
    attacking_account.balance = attacking_balance
    attacking_account.shadow_balance = shadow_balance
    await attacking_account.asave()
    await async_unlock_account(attacking_account)
    defending_account = await async_make_account()
    defending_account.balance = defending_balance
    await defending_account.asave()
    await async_unlock_account(defending_account)
    await set_new_defense(
        defending_account, post_prompt="something" if use_post_prompt else ""
    )
    # assign crazy amounts to make sure we're refreshing from DB properly
    attacking_account.balance = 12387891
    attacking_account.shadow_balance = 99999
    defending_account.balance = 1829192

    moderation_mock = MagicMock()

    # Run async_make_auth_and_transfer
    transfer_result = await async_make_auth_and_transfer(
        attacking_account=attacking_account,
        defending_account=defending_account,
        was_successful=was_successful,
        attacker_input="input",
        llm_output="output",
        moderation_response=moderation_mock,
        auto_moderated=True,
    )

    # freshen up accounts (to make sure we saved)
    await attacking_account.arefresh_from_db()
    await defending_account.arefresh_from_db()

    # Step by step what should be happening with balance, including rounding to
    # int
    if was_successful:
        shadow_transfer_amount = 0
        defending_transfer_amount = int(defending_balance * transfer_fraction)
    else:
        shadow_transfer_amount = int(
            shadow_balance * settings.SHADOW_BALANCE_TRANSFER_FRACTION
        )
        defending_transfer_amount = 0
    defending_balance -= defending_transfer_amount
    attacking_balance += defending_transfer_amount
    shadow_balance -= shadow_transfer_amount
    defending_balance += shadow_transfer_amount

    assert transfer_result.defending_transfer_amount == defending_transfer_amount
    assert transfer_result.shadow_transfer_amount == shadow_transfer_amount
    assert attacking_account.balance == attacking_balance
    assert defending_account.balance == defending_balance
    await sync_to_async(recompute_shadow_balance_inplace)(attacking_account)
    assert abs(attacking_account.shadow_balance - shadow_balance) < 1


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_transfer_with_post_prompt_success(
    mock_auth_attempt_create: None, mock_atomic: None
) -> None:
    await helper_test_transfer_balance(use_post_prompt=True, was_successful=True)


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_transfer_without_post_prompt_success(
    mock_auth_attempt_create: None, mock_atomic: None
) -> None:
    await helper_test_transfer_balance(use_post_prompt=False, was_successful=True)


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_transfer_with_post_prompt_failure(
    mock_auth_attempt_create: None, mock_atomic: None
) -> None:
    await helper_test_transfer_balance(use_post_prompt=True, was_successful=False)


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_transfer_without_post_prompt_failure(
    mock_auth_attempt_create: None, mock_atomic: None
) -> None:
    await helper_test_transfer_balance(use_post_prompt=False, was_successful=False)


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_permission_cooldown_race_condition() -> None:
    # Test a race condition:
    #
    # - The attacker submits a successful attack but gets their account locked
    #   while the attack is in progress.
    # - In this case, the resulting attack should not transfer any money, should
    #   be marked as unsuccessful, and should not set a cooldown timer.
    #
    # We simulate the race condition by directly calling
    # async_make_auth_and_transfer() with was_successful=True, but with a locked
    # attacker account. We then pull out the resulting AuthAttempt and verify
    # that it's marked with was_successful=False and transferred 0 money in both
    # directions.
    defending_account = await async_make_account()
    await async_unlock_account(defending_account)

    attacking_account = await async_make_account()
    # Lock the attacking account to simulate the race condition
    await lock_account(attacking_account)

    moderation_mock = MagicMock()

    # Run async_make_auth_and_transfer with a successful attack
    transfer_result = await async_make_auth_and_transfer(
        attacking_account=attacking_account,
        defending_account=defending_account,
        was_successful=True,
        attacker_input="input",
        llm_output="output",
        moderation_response=moderation_mock,
        auto_moderated=True,
    )

    # Refreshing accounts to make sure we saved
    await attacking_account.arefresh_from_db()
    await defending_account.arefresh_from_db()

    # Asserting that no balance was transferred due to the account being locked
    assert transfer_result.defending_transfer_amount == 0
    assert transfer_result.shadow_transfer_amount == 0
    assert transfer_result.transfer_blocked

    # Pull out the corresponding AuthAttempt and make sure it was not successful
    # as well
    auth_attempt = await AuthAttempt.objects.aget(
        attacking_account=attacking_account, defending_account=defending_account
    )
    was_successful = auth_attempt.was_successful
    assert not was_successful
    assert auth_attempt.attacking_transfer_amount == 0
    assert auth_attempt.defending_transfer_amount == 0


@pytest.mark.asyncio
@pytest.mark.django_db
@pytest.mark.parametrize("llm_choice_str", [c[0] for c in LLM_CHOICES])
async def test_llm_routing(
    llm_choice_str: str,
    mock_streaming_llm_append_model_name: None,
    async_client: AsyncClient,
    mock_content_moderation_flag_on_fail_or_error: None,
) -> None:
    victim_account = await async_make_account()
    attacker_account = await async_make_account()
    await async_login(attacker_account.owner, async_client)
    await async_unlock_account(victim_account)
    await async_unlock_account(attacker_account)
    active_defense = await Defense.objects.aget(account=victim_account)
    active_defense.llm_choice = llm_choice_str
    await active_defense.asave()

    unique_identifier = str(uuid.uuid4())
    post_result = await async_client.post(
        attack_url(victim_account),
        {
            "attacker_input": unique_identifier,
        },
    )
    resp_text = await get_content(post_result)

    # make sure that it outputs the right model
    # (which signals that it was routed to the right place)
    llm_choice = llm_wrappers.LLMChoice[llm_choice_str]
    for other_choice in llm_wrappers.LLMChoice:
        if other_choice == llm_choice:
            assert other_choice.name in resp_text, f"{other_choice=}, {llm_choice=}"
        else:
            assert other_choice.name not in resp_text, f"{other_choice=}, {llm_choice=}"

    # now get the attack by looking for the UUID
    chosen_attack: AuthAttempt = await AuthAttempt.objects.aget(
        attacking_account=attacker_account,
        defending_account=victim_account,
        attacker_input=unique_identifier,
    )
    assert llm_choice_str in str(chosen_attack.llm_output)
    assert chosen_attack.llm_choice == llm_choice_str
