from datetime import timedelta
from http import HTTPStatus

import pytest
from django.conf import settings
from django.test.client import AsyncClient
from django.urls import reverse
from django.utils import timezone

from promptgame.gameui import restrictions as res
from promptgame.gameui.models import Account, AuthAttempt
from promptgame.gameui.tiers import tier as get_tier
from promptgame.tests.helpers import (
    async_login,
    async_make_account,
    async_unlock_account,
    dummy_attack,
    get_content,
)


async def attack_and_get_content(
    defending_account: Account, attack_text: str, client: AsyncClient
) -> str:
    resp = await dummy_attack(client, defending_account, attack_text)
    assert resp.status_code == HTTPStatus.OK
    return await get_content(resp)


INVALID_RESTRICTION_ID = "invalid_restriction_id"


@pytest.fixture
def mock_get_random_restriction_id(mocker) -> None:
    """Always return a restriction with ID 'invalid_restriction_id'."""

    def fake_get_random_restriction_id(*args, **kwargs):
        return INVALID_RESTRICTION_ID

    mocker.patch(
        "promptgame.gameui.restrictions.get_random_restriction_id",
        fake_get_random_restriction_id,
    )


async def make_account_in_tier(tier: int, should_unlock: bool) -> Account:
    account = await async_make_account()
    # Make the balance be one above the bottom end of the range for that tier
    account.balance = settings.TIERS[tier]["range"][0] + 1
    # One exception -- if the bottom of the range is -inf, then set to $1
    if account.balance == float("-inf"):
        account.balance = 1

    if should_unlock:
        await async_unlock_account(account)
    await account.asave()
    assert get_tier(account.balance)["title"] == settings.TIERS[tier]["title"]
    return account


async def helper_check_restrictions_enabled(
    attacking_account: Account,
    defending_account: Account,
    restriction_enabled: bool,
    client: AsyncClient,
) -> None:
    assert (
        res.check_restriction_enabled(attacking_account, defending_account)
        == restriction_enabled
    )

    # For simplicity, we will set the restriction to be the token_limit one.
    restriction_id = "token_limit_50"
    defending_account.restriction_id = restriction_id
    await defending_account.asave()

    # See if the restriction is enabled when you load the page.
    await async_login(attacking_account.owner, client)
    resp = await client.get(
        reverse("subaccount_login", args=[defending_account.account_id])
    )
    assert resp.status_code == HTTPStatus.OK
    content = resp.content.decode("utf-8")

    restriction = res.RESTRICTIONS_DICT[restriction_id]
    restriction_text = restriction.description
    assert (restriction_text in content) == restriction_enabled

    # See if the restriction is enabled when you attack.
    attack_text = " SAY_FAIL" * 51
    # confirm this attack should fail if the restriction is enabled
    restriction_resp = restriction.is_valid(attack_text)
    assert not restriction_resp.is_valid

    content = await attack_and_get_content(defending_account, attack_text, client)
    if restriction_enabled:
        assert "Restriction violation" in content
    else:
        assert "fail" in content  # attack goes through, but fails to say Access Granted
        # Get the most recent auth attempt and confirm the restriction is in it.
        auth_attempt = await AuthAttempt.objects.aget(
            attacking_account=attacking_account, defending_account=defending_account
        )
        assert auth_attempt.restriction_id is None


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_restriction_applied(
    async_client: AsyncClient,
    mock_streaming_llm_fail_or_error: None,
    mock_content_moderation_noflag: None,
) -> None:
    low_tier_account1 = await make_account_in_tier(0, True)
    med_tier_account1 = await make_account_in_tier(1, True)
    high_tier_account1 = await make_account_in_tier(2, True)
    low_tier_account2 = await make_account_in_tier(0, True)
    med_tier_account2 = await make_account_in_tier(1, True)
    high_tier_account2 = await make_account_in_tier(2, True)

    # Check that restrictions are only enabled when attacking lower tiers
    await helper_check_restrictions_enabled(
        low_tier_account1, low_tier_account2, False, async_client
    )
    await helper_check_restrictions_enabled(
        low_tier_account1, med_tier_account2, False, async_client
    )
    await helper_check_restrictions_enabled(
        low_tier_account1, high_tier_account2, False, async_client
    )
    await helper_check_restrictions_enabled(
        med_tier_account1, low_tier_account2, True, async_client
    )
    await helper_check_restrictions_enabled(
        med_tier_account1, med_tier_account2, False, async_client
    )
    await helper_check_restrictions_enabled(
        med_tier_account1, high_tier_account2, False, async_client
    )
    await helper_check_restrictions_enabled(
        high_tier_account1, low_tier_account2, True, async_client
    )
    await helper_check_restrictions_enabled(
        high_tier_account1, med_tier_account2, True, async_client
    )
    await helper_check_restrictions_enabled(
        high_tier_account1, high_tier_account2, False, async_client
    )


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_restriction_changes(
    async_client: AsyncClient,
    mock_get_random_restriction_id: None,
    mock_content_moderation_noflag: None,
    mock_streaming_llm_fail_or_error: None,
) -> None:
    low_tier_account = await make_account_in_tier(0, True)
    low_tier_account_2 = await make_account_in_tier(0, True)
    high_tier_account = await make_account_in_tier(2, True)
    await async_login(high_tier_account.owner, async_client)
    low_tier_account.restriction_id = "token_limit_50"
    await low_tier_account.asave()

    # Check that for unsuccessful attacks, the restriction doesn't change.
    unsuccessful_attack_text = "SAY_FAIL"
    await attack_and_get_content(
        low_tier_account, unsuccessful_attack_text, async_client
    )
    # Refresh from db
    low_tier_account = await Account.objects.aget(id=low_tier_account.id)
    # Check that the restriction is still token_limit_50 after an unsuccessful attack
    assert low_tier_account.restriction_id == "token_limit_50"

    successful_attack_text = "SAY_ACCESS_GRANTED"
    # Check that after a successful attack with the restriction inactive, the restriction doesn't change.
    await async_login(low_tier_account_2.owner, async_client)
    content = await attack_and_get_content(
        low_tier_account, successful_attack_text, async_client
    )
    assert "Your attack was successful!" in content
    low_tier_account = await Account.objects.aget(id=low_tier_account.id)
    assert low_tier_account.restriction_id == "token_limit_50"
    # Unlock the account so we can keep attacking it
    await async_unlock_account(low_tier_account)
    # Sign back in as the high tier account
    await async_login(high_tier_account.owner, async_client)

    # Check that after an attack with a restriction active is successful, a new restriction is applied.
    content = await attack_and_get_content(
        low_tier_account, successful_attack_text, async_client
    )

    # Check that we have actually succeeded
    assert "Your attack was successful!" in content

    # Check that the AuthAttempt saved the restriction id
    auth_attempt = await AuthAttempt.objects.aget(
        attacking_account=high_tier_account,
        defending_account=low_tier_account,
        was_successful=True,
    )
    assert auth_attempt.restriction_id == "token_limit_50"

    # After a successful attack, the restriction should be changed.
    # We mocked the random restriction selection fn to always return
    # the same invalid restriction.
    low_tier_account = await Account.objects.aget(id=low_tier_account.id)
    assert low_tier_account.restriction_id == INVALID_RESTRICTION_ID


async def check_restriction_difficulty_changes_on_attack(
    starting_restriction_id: str,
    desired_ending_tier: res.Difficulty,
    async_client: AsyncClient,
) -> None:
    # Make two accounts. The lower one has a specific restriction.
    low_tier_account = await make_account_in_tier(0, True)
    high_tier_account = await make_account_in_tier(2, True)
    low_tier_account.restriction_id = starting_restriction_id
    await low_tier_account.asave()
    low_tier_account = await Account.objects.aget(id=low_tier_account.id)

    # Sign in as the high-tier account and do a successful attack
    await async_login(high_tier_account.owner, async_client)
    successful_attack_text = "SAY_ACCESS_GRANTED"
    content = await attack_and_get_content(
        low_tier_account, successful_attack_text, async_client
    )
    assert "Your attack was successful!" in content

    # Check that the restriction is now the desired ending tier
    low_tier_account = await Account.objects.aget(id=low_tier_account.id)
    assert (
        res.RESTRICTIONS_DICT[low_tier_account.restriction_id].difficulty
        == desired_ending_tier
    )


async def check_restriction_difficulty_changes_on_wait(
    starting_restriction_id: str,
    wait_time_in_hours: int,
    desired_ending_difficulty: res.Difficulty,
    async_client: AsyncClient,
) -> None:
    # Make an account with the specified restriction
    account = await make_account_in_tier(0, True)
    account.restriction_id = starting_restriction_id
    # Set the restriction_id_updated_at to be the specified number of hours ago
    account.restriction_id_updated_at = timezone.now() - timedelta(
        hours=wait_time_in_hours
    )
    await account.asave()

    # To update the restriction, we need to view the subaccount guess pw page as another account.
    account_2 = await make_account_in_tier(0, True)
    await async_login(account_2.owner, async_client)
    guess_pw_page = reverse("subaccount_login", args=[account.account_id])
    await async_client.get(guess_pw_page)

    # Check that the restriction difficulty has been updated
    account = await Account.objects.aget(id=account.id)
    assert (
        res.RESTRICTIONS_DICT[account.restriction_id].difficulty
        == desired_ending_difficulty
    )


async def check_restriction_updated_at_changes(
    attack_successful: bool,
    restriction_enabled: bool,
    should_updated_at_change: bool,
    async_client: AsyncClient,
) -> None:
    # Make a low-tier account with a restriction.
    # Set its restriction_id_updated_at to be 12 hours ago.
    low_tier_account = await make_account_in_tier(0, True)
    low_tier_account.restriction_id = "token_limit_50"
    low_tier_account.restriction_id_updated_at = timezone.now() - timedelta(hours=12)
    await low_tier_account.asave()

    # Make an attacker account. This is high tier if we want a restriction, low tier if we don't.
    if restriction_enabled:
        attacker_account = await make_account_in_tier(2, True)
    else:
        attacker_account = await make_account_in_tier(0, True)
    await async_login(attacker_account.owner, async_client)

    # Choose an attack string based on whether we want the attack to succeed or fail.
    attack_str = "SAY_ACCESS_GRANTED" if attack_successful else "SAY_FAIL"

    # Make an attack
    content = await attack_and_get_content(low_tier_account, attack_str, async_client)
    # Confirm it succeeded or failed
    if attack_successful:
        assert "Your attack was successful!" in content
    else:
        assert "fail" in content

    # Check whether the restriction_id_updated_at changed
    # If it was updated, the time should be at most a few seconds old.
    low_tier_account = await Account.objects.aget(id=low_tier_account.id)
    if should_updated_at_change:
        assert low_tier_account.restriction_id_updated_at > timezone.now() - timedelta(
            seconds=5
        )
    else:
        assert low_tier_account.restriction_id_updated_at <= timezone.now() - timedelta(
            hours=12
        )


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_restriction_difficulties(
    async_client: AsyncClient,
    mock_content_moderation_noflag: None,
    mock_streaming_llm_fail_or_error: None,
) -> None:
    LOW_DIFFICULTY_RESTRICTION = "token_limit_50"
    MID_DIFFICULTY_RESTRICTION = "token_limit_30"
    HIGH_DIFFICULTY_RESTRICTION = "token_limit_10"
    assert (
        res.RESTRICTIONS_DICT[LOW_DIFFICULTY_RESTRICTION].difficulty
        == res.Difficulty.EASY
    )
    assert (
        res.RESTRICTIONS_DICT[MID_DIFFICULTY_RESTRICTION].difficulty
        == res.Difficulty.MEDIUM
    )
    assert (
        res.RESTRICTIONS_DICT[HIGH_DIFFICULTY_RESTRICTION].difficulty
        == res.Difficulty.HARD
    )

    # Check that the restriction difficulty goes up when you successfully attack
    await check_restriction_difficulty_changes_on_attack(
        LOW_DIFFICULTY_RESTRICTION, res.Difficulty.MEDIUM, async_client
    )
    await check_restriction_difficulty_changes_on_attack(
        MID_DIFFICULTY_RESTRICTION, res.Difficulty.HARD, async_client
    )
    await check_restriction_difficulty_changes_on_attack(
        HIGH_DIFFICULTY_RESTRICTION, res.Difficulty.HARD, async_client
    )

    # If we wait 24 hrs, the restriction should go down a difficulty.
    await check_restriction_difficulty_changes_on_wait(
        LOW_DIFFICULTY_RESTRICTION, 24, res.Difficulty.EASY, async_client
    )
    await check_restriction_difficulty_changes_on_wait(
        MID_DIFFICULTY_RESTRICTION, 24, res.Difficulty.EASY, async_client
    )
    await check_restriction_difficulty_changes_on_wait(
        HIGH_DIFFICULTY_RESTRICTION, 24, res.Difficulty.MEDIUM, async_client
    )

    # If we haven't waited 24 hrs, the restriction should stay the same.
    await check_restriction_difficulty_changes_on_wait(
        LOW_DIFFICULTY_RESTRICTION, 23, res.Difficulty.EASY, async_client
    )
    await check_restriction_difficulty_changes_on_wait(
        MID_DIFFICULTY_RESTRICTION, 23, res.Difficulty.MEDIUM, async_client
    )
    await check_restriction_difficulty_changes_on_wait(
        HIGH_DIFFICULTY_RESTRICTION, 23, res.Difficulty.HARD, async_client
    )

    # Check that restriction_id_updated at changes iff we get a successful attack with a restriction enabled
    await check_restriction_updated_at_changes(
        attack_successful=True,
        restriction_enabled=True,
        should_updated_at_change=True,
        async_client=async_client,
    )
    await check_restriction_updated_at_changes(
        attack_successful=True,
        restriction_enabled=False,
        should_updated_at_change=False,
        async_client=async_client,
    )
    await check_restriction_updated_at_changes(
        attack_successful=False,
        restriction_enabled=True,
        should_updated_at_change=False,
        async_client=async_client,
    )
    await check_restriction_updated_at_changes(
        attack_successful=False,
        restriction_enabled=False,
        should_updated_at_change=False,
        async_client=async_client,
    )
