import asyncio
from typing import Tuple

import openai
import pytest
from django.contrib.auth.models import User
from django.test.client import AsyncClient  # type: ignore

from promptgame.gameui.attack import streaming_llm_output
from promptgame.gameui.llm_wrappers import (
    LLMChoice,
    ModerationResponseAPIError,
    content_moderation,
)
from promptgame.gameui.models import Account
from promptgame.gameui.user_account_tools import async_create_user_and_account
from promptgame.tests.helpers import async_login

SHORT_TEST_WAIT_TIME = 0.2
LONGER_TEST_WAIT_TIME = 1
SUPER_LONG_TEST_WAIT_TIME = 10000


@pytest.fixture
def mock_streaming_llm_output_granted_super_long(mocker) -> None:
    """Waits SUPER_LONG_TEST_WAIT_TIME seconds between each yield."""

    async def fake_async_fn(*args, **kwargs):
        await asyncio.sleep(SUPER_LONG_TEST_WAIT_TIME)
        yield ("Access", True)
        await asyncio.sleep(SUPER_LONG_TEST_WAIT_TIME)
        yield (" granted", True)

    mocker.patch(
        "promptgame.gameui.attack.streaming_llm_output",
        fake_async_fn,
    )


@pytest.fixture
def mock_chat_completion_super_long(mocker) -> None:
    """Takes SUPER_LONG_TEST_WAIT_TIME seconds before responding."""

    async def fake_async_fn(*args, **kwargs):
        await asyncio.sleep(SUPER_LONG_TEST_WAIT_TIME)

    mocker.patch(
        "openai.ChatCompletion.acreate",
        fake_async_fn,
    )


@pytest.fixture
def mock_chat_completion_streaming_super_long(mocker) -> None:
    """Takes SUPER_LONG_TEST_WAIT_TIME seconds before responding."""

    async def fake_async_fn(*args, **kwargs):
        return async_generator()

    async def async_generator():
        await asyncio.sleep(SUPER_LONG_TEST_WAIT_TIME)
        yield {"choices": [{"finish_reason": None, "delta": {"content": "Access"}}]}
        await asyncio.sleep(SUPER_LONG_TEST_WAIT_TIME)
        yield {"choices": [{"finish_reason": None, "delta": {"content": " granted"}}]}

    mocker.patch(
        "openai.ChatCompletion.acreate",
        fake_async_fn,
    )


@pytest.fixture
def mock_moderation_super_long(mocker) -> None:
    """Takes SUPER_LONG_TEST_WAIT_TIME seconds before responding."""

    async def fake_async_fn(*args, **kwargs):
        await asyncio.sleep(SUPER_LONG_TEST_WAIT_TIME)
        return {"results": [{"flagged": False}]}

    mocker.patch(
        "openai.Moderation.acreate",
        fake_async_fn,
    )


@pytest.fixture
def mock_timeout(settings) -> None:
    """Only wait SHORT_TEST_WAIT_TIME seconds before timing out."""
    settings.API_TIMEOUT_SECONDS = SHORT_TEST_WAIT_TIME


async def make_logged_in_user(async_client: AsyncClient) -> Tuple[User, Account]:
    """Create a logged in user."""
    (
        user,
        account,
    ) = await async_create_user_and_account()
    await async_login(user, async_client)
    return user, account


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_chat_completion_timeout(
    # The chat completion will take SUPER_LONG_TEST_WAIT_TIME seconds to respond.
    # It should time out after API_TIMEOUT_SECONDS seconds, but this is slow
    # For a test, so we mock the timeout to be SHORT_TEST_WAIT_TIME seconds.
    mock_chat_completion_super_long: None,
    mock_timeout: None,
    async_client: AsyncClient,
) -> None:
    _, account = await make_logged_in_user(async_client)

    # Check that the request times out.
    # This is to confirm we mocked it correctly.
    with pytest.raises(asyncio.TimeoutError):
        messages = [
            {"role": "system", "content": "defense_pre_prompt"},
            {"role": "user", "content": "attack_prompt"},
            {"role": "user", "content": "defense_post_prompt"},
        ]
        await asyncio.wait_for(
            openai.ChatCompletion.acreate(
                model="gpt-3.5-turbo-0613",
                temperature=0.0,
                max_tokens=10,
                messages=messages,
                stream=True,
            ),
            timeout=SHORT_TEST_WAIT_TIME,
        )

    # Check that we catch the timeout.
    # If we don't, this test will last forever.
    async def streaming_loop():
        async for result in streaming_llm_output(
            attacking_account=account,
            attack_prompt="attack_prompt",
            defense_pre_prompt="defense_pre_prompt",
            defense_post_prompt="defense_post_prompt",
            llm_choice=LLMChoice.GPT_35_TURBO_0613,
            max_tokens=100,
            enable_gpt_repeat_hack_for_defenses=False,
        ):
            assert "Could not generate AI response" in result[0]

    await asyncio.wait_for(streaming_loop(), timeout=LONGER_TEST_WAIT_TIME)


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_moderation_timeout(
    # The moderation will take SUPER_LONG_TEST_WAIT_TIME seconds to respond.
    # It should time out after API_TIMEOUT_SECONDS seconds, but this is slow
    # For a test, so we mock the timeout to be SHORT_TEST_WAIT_TIME seconds.
    mock_moderation_super_long: None,
    mock_timeout: None,
    async_client: AsyncClient,
) -> None:
    await make_logged_in_user(async_client)

    # Check that the moderation times out.
    # This is to confirm we mocked it correctly.
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(
            openai.Moderation.acreate("message"), timeout=SHORT_TEST_WAIT_TIME
        )

    # Check that we catch the timeout.
    # If we don't, thow a timeout error after LONGER_TEST_WAIT_TIME seconds.
    moderation_result = await asyncio.wait_for(
        content_moderation("message"),
        timeout=LONGER_TEST_WAIT_TIME,
    )
    assert isinstance(moderation_result, ModerationResponseAPIError)


@pytest.mark.django_db
@pytest.mark.asyncio
async def test_chat_streaming_timeout(
    # The ChatCompletion endpoint will take SUPER_LONG_TEST_WAIT_TIME seconds to yield each item.
    # We should time out after API_TIMEOUT_SECONDS seconds, but this is slow
    # For a test, so we mock the timeout to be SHORT_TEST_WAIT_TIME seconds.
    mock_chat_completion_streaming_super_long: None,
    mock_timeout: None,
    async_client: AsyncClient,
) -> None:
    # Make a user
    user, account = await async_create_user_and_account()
    await async_login(user, async_client)

    # Check that the request times out.
    # This is to confirm we mocked it correctly.
    messages = [
        {"role": "system", "content": "defense_pre_prompt"},
        {"role": "user", "content": "attack_prompt"},
        {"role": "user", "content": "defense_post_prompt"},
    ]
    chunks = await asyncio.wait_for(
        openai.ChatCompletion.acreate(
            model="gpt-3.5-turbo-0613",
            temperature=0.0,
            max_tokens=10,
            user=user,
            messages=messages,
            stream=True,
        ),
        timeout=SHORT_TEST_WAIT_TIME,
    )
    with pytest.raises(asyncio.TimeoutError):

        async def chunks_loop():
            async for _ in chunks:
                pass

        await asyncio.wait_for(chunks_loop(), timeout=SHORT_TEST_WAIT_TIME)

    # Check that we catch the timeout.
    # If we don't, it will error out after LONGER_TEST_WAIT_TIME seconds.

    async def get_llm_output():
        async for result in streaming_llm_output(
            attacking_account=account,
            attack_prompt="attack_prompt",
            defense_pre_prompt="defense_pre_prompt",
            defense_post_prompt="defense_post_prompt",
            llm_choice=LLMChoice.GPT_35_TURBO_0613,
            max_tokens=100,
            enable_gpt_repeat_hack_for_defenses=False,
        ):
            assert "Could not generate AI response" in result[0]

    await asyncio.wait_for(get_llm_output(), timeout=LONGER_TEST_WAIT_TIME)
