import uuid
from copy import deepcopy
from typing import Dict

import pytest
from asgiref.sync import sync_to_async
from django.conf import settings
from django.test.client import AsyncClient  # type: ignore
from django.urls import reverse

from promptgame.gameui.llm_wrappers import LLMChoice
from promptgame.gameui.models import LLM_CHOICES, Account, Defense
from promptgame.gameui.user_account_tools import create_user_and_account
from promptgame.tests.helpers import async_login


# fixture returning some correct inputs
@pytest.fixture()
def correct_post_args() -> Dict[str, str]:
    return {
        "access_code": "foo",
        "defense_pre_prompt": "bar",
        "defense_post_prompt": "baz",
        "llm_choice": LLMChoice.GPT_35_TURBO_0613.name,
    }


@pytest.fixture()
def success_message() -> str:
    return "Your defense was saved"


@pytest.fixture()
def failure_message() -> str:
    return "Your defense wasn&#x27;t saved"


@pytest.fixture()
def failure_message_pre_prompt_tokens() -> str:
    return f"Your defense wasn&#x27;t saved: your opening defense is too long (&gt;{settings.MAX_DEFENSE_PRE_PROMPT_LENGTH_TOKENS} tokens)"


@pytest.fixture()
def failure_message_post_prompt_tokens() -> str:
    return f"Your defense wasn&#x27;t saved: your closing defense is too long (&gt;{settings.MAX_DEFENSE_POST_PROMPT_LENGTH_TOKENS} tokens)"


@pytest.fixture()
def failure_message_access_code_tokens() -> str:
    return f"Your defense wasn&#x27;t saved: your access code is too long (&gt;{settings.MAX_ACCESS_CODE_LENGTH_TOKENS} tokens)"


@pytest.mark.asyncio
async def test_not_logged_in_fail(
    async_client: AsyncClient,
    correct_post_args: Dict[str, str],
) -> None:
    """Test that client must be logged in to create a defense."""
    resp = await async_client.post(reverse("security_submit"), correct_post_args)
    assert resp.status_code == 403


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_create_defense(
    async_client: AsyncClient,
    mock_content_moderation_noflag: None,
    mock_streaming_llm_output_granted: None,
    correct_post_args: Dict[str, str],
    success_message: str,
    failure_message: str,
) -> None:
    """Test that basic defense creation works."""
    user, account = await sync_to_async(create_user_and_account)()
    old_defense = await sync_to_async(lambda account: account.active_defense)(account)
    await async_login(user, async_client)
    resp = await async_client.post(reverse("security_submit"), correct_post_args)
    assert resp.status_code == 200
    resp_content = resp.content.decode("utf8")
    assert success_message in resp_content
    assert failure_message not in resp_content
    await account.arefresh_from_db()
    new_defense = await sync_to_async(lambda account: account.active_defense)(account)
    assert old_defense != new_defense
    assert new_defense.access_code == correct_post_args["access_code"]


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_create_defense_fail(
    async_client: AsyncClient,
    mock_content_moderation_noflag: None,
    mock_streaming_llm_output_denied: None,
    correct_post_args: Dict[str, str],
    success_message: str,
    failure_message: str,
) -> None:
    """Check that defense creation fails if LLM says 'access denied'."""
    user, account = await sync_to_async(create_user_and_account)()
    old_defense = await sync_to_async(lambda account: account.active_defense)(account)
    await async_login(user, async_client)
    resp = await async_client.post(reverse("security_submit"), correct_post_args)
    assert resp.status_code == 200
    resp_content = resp.content.decode("utf8")
    assert success_message not in resp_content
    assert failure_message in resp_content
    await account.arefresh_from_db()
    new_defense = await sync_to_async(lambda account: account.active_defense)(account)
    assert old_defense == new_defense


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_create_defense_fail_token_limit(
    async_client: AsyncClient,
    mock_content_moderation_noflag: None,
    correct_post_args: Dict[str, str],
    success_message: str,
    failure_message_pre_prompt_tokens: str,
    failure_message_post_prompt_tokens: str,
    failure_message_access_code_tokens: str,
) -> None:
    """Check that defense creation fails if token limits exceeded"""
    user, account = await sync_to_async(create_user_and_account)()
    old_defense = await sync_to_async(lambda account: account.active_defense)(account)
    await async_login(user, async_client)

    # too many tokens in defense pre-prompt
    token_test_post_args = deepcopy(correct_post_args)
    long_pre_prompt = " a" * (1 + settings.MAX_DEFENSE_PRE_PROMPT_LENGTH_TOKENS)
    token_test_post_args["defense_pre_prompt"] = long_pre_prompt
    resp = await async_client.post(
        reverse("security_submit"),
        token_test_post_args,
    )
    resp_content = resp.content.decode("utf8")
    assert success_message not in resp_content
    assert failure_message_pre_prompt_tokens in resp_content
    await account.arefresh_from_db()
    new_defense = await sync_to_async(lambda account: account.active_defense)(account)
    assert old_defense == new_defense

    # too many tokens in defense post-prompt
    token_test_post_args = deepcopy(correct_post_args)
    long_post_prompt = " a" * (1 + settings.MAX_DEFENSE_POST_PROMPT_LENGTH_TOKENS)
    token_test_post_args["defense_post_prompt"] = long_post_prompt
    resp = await async_client.post(
        reverse("security_submit"),
        token_test_post_args,
    )
    resp_content = resp.content.decode("utf8")
    assert success_message not in resp_content
    assert failure_message_post_prompt_tokens in resp_content
    await account.arefresh_from_db()
    new_defense = await sync_to_async(lambda account: account.active_defense)(account)
    assert old_defense == new_defense

    # too many tokens in access code
    token_test_post_args = deepcopy(correct_post_args)
    long_access_code = " a" * (1 + settings.MAX_ACCESS_CODE_LENGTH_TOKENS)
    token_test_post_args["access_code"] = long_access_code
    resp = await async_client.post(
        reverse("security_submit"),
        token_test_post_args,
    )
    resp_content = resp.content.decode("utf8")
    assert success_message not in resp_content
    assert failure_message_access_code_tokens in resp_content
    await account.arefresh_from_db()
    new_defense = await sync_to_async(lambda account: account.active_defense)(account)
    assert old_defense == new_defense


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_create_defense_moderation(
    async_client: AsyncClient,
    mock_content_moderation_flag_on_fail_or_error: None,
    mock_streaming_llm_output_granted: None,
    success_message: str,
):
    # Check that defense creation fails if moderation says the input is naughty.
    user, account = await sync_to_async(create_user_and_account)()
    await async_login(user, async_client)
    resp = await async_client.post(
        reverse("security_submit"),
        {
            "access_code": "foo",
            "defense_pre_prompt": "bar",
            "defense_post_prompt": "fail",  # this will make it fail
            "llm_choice": LLMChoice.GPT_35_TURBO_0613.name,
        },
    )
    naughty_string = "Your input contained the word &#x27;fail&#x27;."
    assert naughty_string in resp.content.decode("utf8")

    # Check that defense creation succeeds if moderation gives an api error
    # but that the defense's auto_moderated field is set to False.
    resp = await async_client.post(
        reverse("security_submit"),
        {
            "access_code": "foo",
            "defense_pre_prompt": "bar",
            "defense_post_prompt": "error",  # this will make it error
            "llm_choice": LLMChoice.GPT_35_TURBO_0613.name,
        },
    )
    assert success_message in resp.content.decode("utf8")
    account = await Account.objects.select_related("active_defense").aget(pk=account.pk)
    defense = account.active_defense
    assert not defense.auto_moderated
    assert defense.post_prompt == "error"

    # If moderation passes, defense creation should succeed.
    resp = await async_client.post(
        reverse("security_submit"),
        {
            "access_code": "foo",
            "defense_pre_prompt": "bar",
            "defense_post_prompt": "success",
            "llm_choice": LLMChoice.GPT_35_TURBO_0613.name,
        },
    )
    assert success_message in resp.content.decode("utf8")
    account = await Account.objects.select_related("active_defense").aget(pk=account.pk)
    defense = account.active_defense
    assert defense.auto_moderated
    assert defense.post_prompt == "success"


@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,
    mock_content_moderation_flag_on_fail_or_error: None,
    async_client: AsyncClient,
) -> None:
    # submit an attack
    user, account = await sync_to_async(create_user_and_account)()
    await async_login(user, async_client)
    unique_ident = str(uuid.uuid4())
    resp = await async_client.post(
        reverse("security_submit"),
        {
            "access_code": unique_ident,
            "defense_pre_prompt": unique_ident,
            "defense_post_prompt": unique_ident,
            "llm_choice": llm_choice_str,
        },
    )

    # it's getting routed to the right model
    # (however, it won't get saved because the thing is not outputting access
    # granted)
    resp_text = resp.content.decode("utf8")
    llm_choice = LLMChoice[llm_choice_str]
    for other_choice in 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=}"


@pytest.mark.asyncio
@pytest.mark.django_db
@pytest.mark.parametrize("llm_choice_str", [c[0] for c in LLM_CHOICES])
async def test_llm_saved(
    llm_choice_str: str,
    async_client: AsyncClient,
    mock_streaming_llm_output_granted: None,
    mock_content_moderation_flag_on_fail_or_error: None,
) -> None:
    # submit an attack
    user, account = await sync_to_async(create_user_and_account)()
    await async_login(user, async_client)
    unique_ident = str(uuid.uuid4())
    resp = await async_client.post(
        reverse("security_submit"),
        {
            "access_code": unique_ident,
            "defense_pre_prompt": unique_ident,
            "defense_post_prompt": unique_ident,
            "llm_choice": llm_choice_str,
        },
    )

    # now we test that it gets saved
    active_defense = await Defense.objects.aget(account=account)
    assert active_defense.llm_choice == llm_choice_str
