import json
from enum import Enum
from typing import Dict

from google.cloud import aiplatform
from tenacity import retry, stop_after_attempt, wait_exponential_jitter
from ttdata.common import DetectionModelInput, LLMPromptType, ModelAdapter, ModelInput
from vertexai.preview.language_models import ChatModel

# Based on https://github.com/googleapis/python-aiplatform/blob/8271301454814b233a630d1c18ebe5e4833fcec2/vertexai/language_models/_language_models.py#L1428
USER_AUTHOR = "user"
MODEL_AUTHOR = "model"


class VertexAIModel(Enum):
    PaLM_2_Chat = "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


class VertexAIAdapter(ModelAdapter):
    def __init__(self, model: VertexAIModel) -> None:
        super().__init__()
        self.model = model
        aiplatform.init(
            project="prompt-ad-game",
        )

    # There are multiple ways to structure the prompt for VertexAI.
    # See: https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-chat
    # Based on the experiments in ../notebooks/tt_test_vertexai_prompts.ipynb, the following approach ended up with the
    # highest accuracy on existing data.
    def _get_prediction_instance(
        self,
        model_input: ModelInput | DetectionModelInput,
        prompt_type: LLMPromptType = LLMPromptType.PREDICT,
    ) -> dict[str, list[dict[str, str]]]:
        if prompt_type == LLMPromptType.DETECT_FEW_SHOT:
            raise NotImplementedError()
        elif prompt_type == LLMPromptType.DETECT_ZERO_SHOT:
            raise NotImplementedError()
        assert prompt_type == LLMPromptType.PREDICT

        message_structs = self._get_prediction_messages(model_input)
        prediction_instance = {"messages": message_structs}
        return prediction_instance

    def _get_prediction_messages(self, model_input: ModelInput) -> list[Dict[str, str]]:
        message_structs = [
            {"user": MODEL_AUTHOR, "content": make_nonempty(model_input.pre_prompt)},
            {
                "user": USER_AUTHOR,
                "content": make_nonempty(model_input.attack_or_access_code),
            },
            {"user": USER_AUTHOR, "content": make_nonempty(model_input.post_prompt)},
        ]
        return message_structs

    def get_cache_key(
        self,
        *,
        model_input: ModelInput | DetectionModelInput,
        temperature: float,
        max_length_tokens: int,
        prompt_type: LLMPromptType = LLMPromptType.PREDICT,
    ) -> str:
        return json.dumps(
            (
                self.model.value,
                self._get_prediction_instance(model_input, prompt_type),
                temperature,
                max_length_tokens,
            ),
        )

    @retry(wait=wait_exponential_jitter(initial=5, max=20), stop=stop_after_attempt(10))
    def _predict(
        self,
        *,
        model_input: ModelInput,
        temperature: float,
        max_length_tokens: int,
        prompt_type: LLMPromptType,
    ) -> str:
        chat_model = ChatModel.from_pretrained(self.model.value)
        chat = chat_model.start_chat(
            context=model_input.pre_prompt,
        )

        prediction_instance = self._get_prediction_instance(model_input, prompt_type)
        parameters = {
            "temperature": temperature,
            "maxDecodeSteps": max_length_tokens,
        }
        # The public Vertex API (chat.send_message) limits us to sending one message at a time,
        # so we use the private methods here to predict with multiple messages.
        prediction_response = chat._model._endpoint.predict(
            instances=[prediction_instance],
            parameters=parameters,
        )
        response_obj = chat._parse_chat_prediction_response(
            prediction_response=prediction_response
        )

        return response_obj.text
