import uuid
from http import HTTPStatus
from typing import List, Optional

import pytest
from bs4 import BeautifulSoup
from django.test.client import AsyncClient, Client  # type: ignore
from django.urls import reverse

from promptgame.gameui.llm_wrappers import LLMChoice
from promptgame.gameui.models import LLM_CHOICES, Account, AuthAttempt
from promptgame.gameui.user_account_tools import (
    async_create_user_and_account,
    create_user_and_account,
)
from promptgame.tests.helpers import (
    SSEMessage,
    async_login,
    get_content,
    login,
    parse_sse_messages,
    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()


async def make_account_and_login(async_client: AsyncClient) -> Account:
    user, account = await async_create_user_and_account()
    await async_login(user, async_client)
    return account


@pytest.mark.django_db
def test_view_page(client: Client) -> None:
    """Check that we can actually view the page with the form."""
    user, account = create_user_and_account()
    login(user, client)

    resp = client.get(reverse("sandbox"))
    assert resp.status_code == HTTPStatus.OK
    content = resp.content.decode("utf-8")
    assert "Attack sandbox" in content

    # parse as HTML so we can check that there's a link to Our own self-attack
    # page. Link text should be "Go Back".
    target_link = reverse("subaccount_login", args=[account.account_id])
    soup = BeautifulSoup(content, "html.parser")
    link = soup.find("a", href=target_link)
    assert link is not None
    assert link.text == "Go Back"


async def do_test_attack(
    async_client: AsyncClient,
    pre_prompt: Optional[str] = None,
    attack: Optional[str] = None,
    post_prompt: Optional[str] = None,
    llm_choice: LLMChoice = LLMChoice.GPT_35_TURBO_0613,
) -> List[SSEMessage]:
    form_values = {
        "defense_pre_prompt": pre_prompt,
        "attacker_input": attack,
        "defense_post_prompt": post_prompt,
        "llm_choice": llm_choice.name,
    }
    form_values_filtered = {k: v for k, v in form_values.items() if v is not None}
    resp = await async_client.post(reverse("sandbox_submit"), form_values_filtered)
    content = await get_content(resp)
    messages = parse_sse_messages(content)
    return messages


def extract_llm_response(stream: List[SSEMessage]) -> str:
    # Get back the LLM response from delta messages.
    # The delta messages have `event: hx-swap:none`, then a data body that looks
    # like this:
    #
    # <span id="llm-response-bubble" hx-swap-oob="beforeend">extra token from
    # LLM</span>
    contents = []
    for message in stream:
        if (
            message.data is not None
            and 'id="llm-response-bubble" hx-swap-oob="beforeend"' in message.data
        ):
            # find the thing inside the <span></span>
            inner_content = message.data.split(">", 1)[1].split("<", 1)[0]
            contents.append(inner_content)
    return "".join(contents)


def extract_user_messages(stream: List[SSEMessage]) -> List[str]:
    # messages to the user look like this:
    # <div class="flex flex-col justify-center">
    #   <div class="chat-info-message {{ type }}">
    #     {{ message }}
    #   </div>
    # </div>
    # We want to just extract the inside
    user_messages = []
    for message in stream:
        if (
            message.data is not None
            and '<div class="chat-info-message ' in message.data
        ):
            # parse with BS4 and return
            soup = BeautifulSoup(message.data, "html.parser")
            # get inner text
            user_messages.append(soup.text)
    return user_messages


def check_successful(stream: List[SSEMessage]) -> None:
    assert "access granted" in extract_llm_response(stream).lower()
    assert any("Your attack was successful" in m for m in extract_user_messages(stream))


def check_denied_access(stream: List[SSEMessage]) -> None:
    assert "access granted" not in extract_llm_response(stream).lower()
    assert any("Your attack failed" in m for m in extract_user_messages(stream))


def check_llm_api_failed(stream: List[SSEMessage]) -> None:
    # this failure happens when the OpenAI API fails
    assert any(
        "Could not generate AI response; please refresh and try again." in m
        for m in extract_user_messages(stream)
    )


def check_generic_failure(stream: List[SSEMessage]) -> None:
    # this is the failure that is given when we can't actually serve the request
    assert len(extract_llm_response(stream)) == 0
    assert any(
        "Server error, try refreshing the page" in m
        for m in extract_user_messages(stream)
    )


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_no_log_in(
    mock_streaming_llm_output_granted: None, async_client: AsyncClient
) -> None:
    # if we don't log in then we should get an error
    stream = await do_test_attack(
        async_client=async_client, pre_prompt="foo", attack="bar", post_prompt="baz"
    )
    check_generic_failure(stream)


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_access_granted_basic(
    mock_streaming_llm_output_granted: None, async_client: AsyncClient
) -> None:
    # make sure we get success and our thing is logged
    account = await make_account_and_login(async_client)
    pre_prompt = "foo42"
    attack = "bar82"
    post_prompt = "baz37"
    stream = await do_test_attack(
        async_client=async_client,
        pre_prompt=pre_prompt,
        attack=attack,
        post_prompt=post_prompt,
    )
    check_successful(stream)

    # get most recent self-attack (authattempt where attacker and defender are this account)
    auth_attempt = (
        await AuthAttempt.objects.filter(
            attacking_account=account, defending_account=account
        )
        .order_by("-timestamp")
        .afirst()
    )
    assert auth_attempt is not None
    assert auth_attempt.was_successful
    assert auth_attempt.defense_pre_prompt == pre_prompt
    assert auth_attempt.attacker_input == attack
    assert auth_attempt.defense_post_prompt == post_prompt


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_zero_length_attack(
    mock_streaming_llm_output_granted: None, async_client: AsyncClient
) -> None:
    # make sure that attacks also succeed with the empty string
    # (we don't check that entries actually end up in the DB; I assume the last
    # test was sufficient)
    await make_account_and_login(async_client)
    stream = await do_test_attack(
        async_client=async_client, pre_prompt="", attack="", post_prompt=""
    )
    check_successful(stream)


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_missing_fields(
    mock_streaming_llm_output_granted: None, async_client: AsyncClient
) -> None:
    # but don't allow attack if fields are missing entirely
    await make_account_and_login(async_client)
    stream = await do_test_attack(async_client=async_client)
    check_generic_failure(stream)


@pytest.mark.asyncio
@pytest.mark.django_db
async def test_denied_attacks(
    mock_streaming_llm_output_denied: None,
    async_client: AsyncClient,
) -> None:
    await make_account_and_login(async_client)
    stream = await do_test_attack(
        async_client=async_client, pre_prompt="", attack="", post_prompt=""
    )
    check_denied_access(stream)


@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:
    """Test that requests are routed to the right LLM."""
    await make_account_and_login(async_client)
    # identify the submitted attack with a uuid
    uuid_string = str(uuid.uuid4())
    llm_choice = LLMChoice[llm_choice_str]
    stream = await do_test_attack(
        async_client=async_client,
        pre_prompt=uuid_string,
        attack=uuid_string,
        post_prompt=uuid_string,
        llm_choice=llm_choice,
    )
    resp_text = extract_llm_response(stream)

    # iterate through all the model choices and check that the model choice is
    # in the output IFF it's the model we chose
    for other_choice in LLMChoice:
        if other_choice == llm_choice:
            assert other_choice.name in resp_text
        else:
            assert other_choice.name not in resp_text, llm_choice

    submitted_attack: AuthAttempt = await AuthAttempt.objects.aget(
        attacker_input=uuid_string
    )
    assert llm_choice_str in str(submitted_attack.llm_output)
    assert submitted_attack.llm_choice == llm_choice_str
