import asyncio
from contextlib import contextmanager
from typing import AsyncGenerator, Callable
from unittest.mock import MagicMock

import pytest
from django.contrib.auth.models import User
from django.db import transaction

from promptgame.gameui.llm_wrappers import (
    LLMChoice,
    LLMInput,
    ModerationResponseAPIError,
    ModerationResponseInappropriate,
    ModerationResponseOK,
)
from promptgame.gameui.models import AuthAttempt


@pytest.fixture
def mock_streaming_llm_output_granted(mocker) -> None:
    """Mocks streaming_llm_output so that it always gives 'Access granted'."""

    async def fake_async_fn(*args, **kwargs):
        yield ("Access", True)
        yield (" granted", True)

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


@pytest.fixture
def mock_streaming_llm_fail_or_error(mocker) -> None:
    """Mocks streaming_llm_output so that it always gives 'Access granted'."""

    async def fake_async_fn(attacking_account, attack_prompt, *args, **kwargs):
        if "SAY_FAIL" in attack_prompt:
            yield ("fail", True)
        elif "SAY_ERROR" in attack_prompt:
            yield ("error", True)
        elif "SAY_ACCESS_GRANTED" in attack_prompt:
            yield ("Access", True)
            yield (" granted", True)
        else:
            yield (attack_prompt, True)

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


@pytest.fixture
def mock_streaming_llm_append_model_name(mocker) -> None:
    """Mocks the LLM-specific response streaming APIs so that they echo back the
    attack with the model name appended to the end."""

    def make_mock(
        model_choice: LLMChoice,
    ) -> Callable[[User, LLMInput, int], AsyncGenerator[str, None]]:
        async def mock_fn(
            user: User,
            llm_input: LLMInput,
            max_tokens: int,
            enable_gpt_repeat_hack_for_defenses: bool = False,
        ) -> AsyncGenerator[str, None]:
            yield f"{llm_input.attack_or_access_code} {model_choice.name}"

        return mock_fn

    mocker.patch(
        "promptgame.gameui.llm_wrappers.call_openai_api",
        make_mock(LLMChoice.GPT_35_TURBO_0613),
    )
    mocker.patch(
        "promptgame.gameui.llm_wrappers.call_palm_api",
        make_mock(LLMChoice.PALM_CHAT_BISON_001),
    )

    mocker.patch(
        "promptgame.gameui.llm_wrappers.call_claude_api",
        make_mock(LLMChoice.CLAUDE_INSTANT_1_2),
    )


@pytest.fixture
def mock_streaming_llm_output_granted_slow(mocker) -> None:
    """Like the 'access granted' mock, but takes a little bit of time to yield
    each."""

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

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


@pytest.fixture
def mock_streaming_llm_output_denied(mocker) -> None:
    """An 'access denied' mock"""

    async def fake_async_fn(*args, **kwargs):
        yield ("Access", True)
        yield (" denied", True)

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


@pytest.fixture
def mock_content_moderation_noflag(mocker) -> None:
    """Mock content_moderation so that it always returns False (not flagged)."""

    async def fake_async_fn(*args, **kwargs):
        return ModerationResponseOK({})

    mocker.patch(
        "promptgame.gameui.llm_wrappers.content_moderation",
        fake_async_fn,
    )


@pytest.fixture
def mock_content_moderation_flag_on_fail_or_error(mocker) -> None:
    """Mock content_moderation so that it returns True (flagged) iff the input contains 'fail'."""

    async def fake_async_fn(text, *args, **kwargs):
        if "fail" in text:
            return ModerationResponseInappropriate(
                "Your input contained the word 'fail'.", {}
            )
        elif "error" in text:
            return ModerationResponseAPIError("API error")
        return ModerationResponseOK({})

    mocker.patch(
        "promptgame.gameui.llm_wrappers.content_moderation",
        fake_async_fn,
    )


@pytest.fixture
def mock_requests_post_success(mocker) -> None:
    """Mock requests.post with successful captcha request"""

    async def fake_async_fn(*args, **kwargs):
        post_mock = MagicMock()
        post_mock.json = lambda: {"success": True}
        return post_mock

    mocker.patch(
        "httpx.AsyncClient.post",
        fake_async_fn,
    )


@pytest.fixture
def mock_requests_post_failure(mocker) -> None:
    """Mock requests.post with failed captcha request"""

    async def fake_async_fn(*args, **kwargs):
        post_mock = MagicMock()
        post_mock.json = lambda: {"success": False}
        return post_mock

    mocker.patch(
        "httpx.AsyncClient.post",
        fake_async_fn,
    )


@pytest.fixture
def mock_auth_attempt_create(mocker) -> None:
    auth_attempt_mock = mocker.MagicMock(spec=AuthAttempt)
    auth_attempt_mock.save.side_effect = lambda: None  # Mock the save method

    mocker.patch.object(AuthAttempt.objects, "create", return_value=auth_attempt_mock)


@pytest.fixture
def mock_atomic(mocker) -> None:
    """
    Mock Django's transaction.atomic() context manager.
    """

    @contextmanager
    def no_op_cm():
        """A no-op context manager."""
        yield

    mocker.patch.object(transaction, "atomic", no_op_cm)


@pytest.fixture
def mock_check_restriction_enabled(mocker) -> None:
    """Mock check_restriction_enabled so that it always returns False."""

    mocker.patch(
        "promptgame.gameui.restrictions.check_restriction_enabled",
        lambda *args, **kwargs: False,
    )
