import asyncio
import logging
from dataclasses import dataclass
from enum import Enum
from typing import AsyncGenerator, Tuple

import anthropic
import vertexai
from django.conf import settings
from openai import OpenAIError
from vertexai.language_models import ChatModel

from promptgame.gameui.utils import remove_gpt_token_runs

log = logging.getLogger(__name__)

_cached_openai = None


class LLMChoice(Enum):
    # These give a specific LLM version, not just general vendor info. Ideally
    # we should never remove from this enum unless a model is removed from
    # vendor support; if you want to add a new version of a model, it's ideal to
    # introduce it as a new enum element instead of replacing the old one (and
    # keeping the old code so that existing defenses don't break).
    GPT_35_TURBO_0613 = "GPT_35_TURBO_0613"
    PALM_CHAT_BISON_001 = "PALM_CHAT_BISON_001"
    CLAUDE_INSTANT_1_2 = "CLAUDE_INSTANT_1_2"


@dataclass(frozen=True)
class LLMInput:
    """Defense (+ access code) or attack that we want to run against a model.
    Model classes are responsible for deciding how to serialize this high-level
    representation into a token stream."""

    pre_prompt: str
    attack_or_access_code: str
    post_prompt: str


def openai_module():
    global _cached_openai
    if _cached_openai is None:
        import openai

        openai.api_key = settings.OPENAI_API_KEY
        _cached_openai = openai
    return _cached_openai


async def call_content_moderation_api(message: str) -> Tuple[bool, dict]:
    openai = openai_module()
    log.info(f"Making OpenAI moderation call")
    response = await asyncio.wait_for(
        openai.Moderation.acreate(input=message), timeout=settings.API_TIMEOUT_SECONDS
    )
    log.info(f"Finished OpenAI moderation call (acreate done)")
    assert isinstance(response, dict), f"{type(response)=}"
    return response["results"][0]["flagged"], response


@dataclass
class ModerationResponseOK:
    """The endpoint responded and the message is OK to send."""

    endpoint_metadata: dict


@dataclass
class ModerationResponseInappropriate:
    """The endpoint responded and the message is not okay to send."""

    failure_message: str
    endpoint_metadata: dict


@dataclass
class ModerationResponseAPIError:
    """The endpoint did not respond due to an error (e.g. rate-limiting)."""

    failure_message: str


ModerationResponse = (
    ModerationResponseOK | ModerationResponseInappropriate | ModerationResponseAPIError
)


async def content_moderation(message: str) -> ModerationResponse:
    try:
        content_moderation_failed, response_dict = await call_content_moderation_api(
            message
        )
        if content_moderation_failed:
            log.info(f"OpenAI content moderation failed")
            message_str = (
                "Invalid access attempt: your input did not pass our moderation tools. "
                "You'll need to change it and try again."
            )
            return ModerationResponseInappropriate(message_str, response_dict)
        return ModerationResponseOK(response_dict)
    except (OpenAIError, asyncio.TimeoutError) as ex:
        log.info(
            f"OpenAIError or TimeoutError while querying OpenAI moderation endpoint: {type(ex)} {ex}"
        )
        return ModerationResponseAPIError(
            "⚠️ Trouble with API request; please refresh and try again. "
            "If the error persists, contact robust-llms@berkeley.edu for help. ",
        )
    raise RuntimeError("unreachable")  # this keeps pyright happy


async def call_llm_api(
    user: str,
    llm_input: LLMInput,
    max_tokens: int,
    llm_choice: LLMChoice,
    enable_gpt_repeat_hack_for_defenses: bool,
) -> AsyncGenerator[str, None]:
    match llm_choice:
        case LLMChoice.GPT_35_TURBO_0613:
            return call_openai_api(
                user=user,
                llm_input=llm_input,
                max_tokens=max_tokens,
                enable_gpt_repeat_hack_for_defenses=enable_gpt_repeat_hack_for_defenses,
            )
        case LLMChoice.PALM_CHAT_BISON_001:
            return call_palm_api(
                user=user,
                llm_input=llm_input,
                max_tokens=max_tokens,
            )
        case LLMChoice.CLAUDE_INSTANT_1_2:
            return call_claude_api(
                user=user,
                llm_input=llm_input,
                max_tokens=max_tokens,
            )
    raise NotImplementedError(f"Unknown model {llm_choice}")


async def call_openai_api(
    *,
    user: str,
    llm_input: LLMInput,
    max_tokens: int,
    enable_gpt_repeat_hack_for_defenses: bool,
) -> AsyncGenerator[str, None]:
    openai = openai_module()
    log.info(f"Making OpenAI chat completion request for {user=}")

    pre_prompt = llm_input.pre_prompt
    post_prompt = llm_input.post_prompt
    if enable_gpt_repeat_hack_for_defenses:
        # hacky workaround for OpenAI rejecting long runs of the same token in
        # our saved defenses
        pre_prompt = remove_gpt_token_runs(pre_prompt)
        post_prompt = remove_gpt_token_runs(post_prompt)

    # Messages are formatted system, user, user
    message_dicts = [
        {"role": "system", "content": pre_prompt},
        {"role": "user", "content": llm_input.attack_or_access_code},
        {"role": "user", "content": post_prompt},
    ]

    result = await asyncio.wait_for(
        openai.ChatCompletion.acreate(
            model="gpt-3.5-turbo-0125",
            temperature=0.0,
            max_tokens=max_tokens,
            user=user,
            messages=message_dicts,
            stream=True,
        ),
        timeout=settings.API_TIMEOUT_SECONDS,
    )
    log.info(f"Created OpenAI chat completion request for {user=}, can stream now")
    async for chunk in result:
        choice = chunk["choices"][0]
        next_text: str = choice.get("delta", {}).get("content", "")
        if next_text:
            yield next_text
        if choice["finish_reason"] == "length":
            yield "... (cut off by length limit)"
        if choice["finish_reason"] == "content_filter":
            log.warning(f"Completion truncated by content filter for {user=}")
            yield "... (cut off by content filter)"


_VERTEX_INITIALIZED = False


async def call_palm_api(
    *, user: str, llm_input: LLMInput, max_tokens: int
) -> AsyncGenerator[str, None]:
    # init the API if necessary
    global _VERTEX_INITIALIZED
    if not _VERTEX_INITIALIZED:
        # FIXME(sam): move these to Django-level settings
        vertexai.init(project=settings.VERTEX_AI_PROJECT_ID)
        _VERTEX_INITIALIZED = True

    chat_model = ChatModel.from_pretrained("chat-bison@001")

    def make_nonempty(s: str) -> str:
        """I was getting 400 errors with message 'The text content is empty', so I
        wrapped this around our strings to make sure they're always non-empty."""
        if s == "":
            return " "
        else:
            return s

    chat = chat_model.start_chat(
        context=make_nonempty(llm_input.pre_prompt),
        max_output_tokens=max_tokens,
        temperature=0,
    )
    responses = chat.send_message_streaming_async(
        message=f"{llm_input.attack_or_access_code}\n{llm_input.post_prompt}",
    )

    log.info(f"Created VertexAI chat completion request for {user=}, can stream now")
    async for resp in responses:
        to_yield = resp.text
        # Vertex gives us a way of detecting blocked content, but not of
        # detecting truncated responses
        if resp.is_blocked:
            to_yield = f"️{to_yield} ⚠ Response blocked by upstream content filter"
        yield resp.text


async def call_claude_api(
    *, user: str, llm_input: LLMInput, max_tokens: int
) -> AsyncGenerator[str, None]:
    async with anthropic.AsyncAnthropic(
        api_key=settings.ANTHROPIC_API_KEY
    ) as async_anthropic:
        # Changed 2024-11-28 after model deprecation
        # responses = await async_anthropic.completions.create(
        #     model="claude-instant-1.2",
        #     max_tokens_to_sample=max_tokens,
        #     prompt=f"{anthropic.HUMAN_PROMPT} {llm_input.pre_prompt}\n{llm_input.attack_or_access_code}\n{llm_input.post_prompt}{anthropic.AI_PROMPT}",
        #     temperature=0,
        #     stream=True,
        # )

        # log.info(f"Created Claude completion request for {user=}, can stream now")
        # async for resp in responses:
        #     yield resp.completion
        #     if resp.stop_reason == "max_tokens":
        #         yield "... (cut off by length limit)"

        async with async_anthropic.messages.stream(
            model="claude-3-haiku-20240307",
            max_tokens=max_tokens,
            messages=[
                {
                    "role": "user",
                    "content": llm_input.pre_prompt,
                },
                {
                    "role": "user",
                    "content": llm_input.attack_or_access_code,
                },
                {
                    "role": "user",
                    "content": llm_input.post_prompt,
                },
            ],
            temperature=0,
        ) as stream:
            log.info(f"Created Claude messages stream for {user=}")
            async for event in stream:
                if event.type == "text":
                    yield event.text
                if (
                    event.type == "message_delta"
                    and event.delta.stop_reason == "max_tokens"
                ):
                    yield "... (cut off by length limit)"
