import json
import logging
import os
import time
from enum import Enum

import requests
from requests.adapters import HTTPAdapter
from ttdata.common import DetectionModelInput, LLMPromptType, ModelAdapter, ModelInput
from urllib3.util import Retry

log = logging.getLogger(__name__)


class DeepInfraModel(Enum):
    """
    This is not an exhaustive list of the models that deep infra has,
    but they are probably the most relevant ones for our purposes
    """

    # Disabling non-chat model for now because it seems to be more expensive
    # than the rest  (charging per second instead of per token)
    # LLAMA_2_7B = "meta-llama/Llama-2-7b-hf"
    LLAMA_2_7B_CHAT = "meta-llama/Llama-2-7b-chat-hf"
    # this is a bit more expensive than the rest
    LLAMA_2_13B_CHAT = "meta-llama/Llama-2-13b-chat-hf"
    LLAMA_2_70B_CHAT = "meta-llama/Llama-2-70b-chat-hf"
    CODELLAMA_34B_INSTRUCT = "codellama/CodeLlama-34b-Instruct-hf"
    # TODO: uncomment this line when we support Falcon7B
    # FALCON_7B = "tiiuae/falcon-7b"


class DeepInfraAdapter(ModelAdapter):
    endpoint = "https://api.deepinfra.com/v1/inference/"
    headers = {
        "Content-Type": "application/json",
    }

    def __init__(self, model: DeepInfraModel) -> None:
        super().__init__()
        self.model = model
        self.endpoint = self.endpoint + model.value

    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_prompt(model_input, prompt_type),
                temperature,
                max_length_tokens,
            )
        )

    def _get_prompt(self, model_input: ModelInput, prompt_type: LLMPromptType) -> str:
        # based on actual llama generation: https://github.com/facebookresearch/llama/blob/main/llama/generation.py#L327
        #   test results on n=49 --> 19/49 = 38.78% replay success rate
        if prompt_type == LLMPromptType.DETECT_FEW_SHOT:
            raise NotImplementedError()
        elif prompt_type == LLMPromptType.DETECT_ZERO_SHOT:
            raise NotImplementedError()
        assert prompt_type == LLMPromptType.PREDICT

        return (
            f"[INST] <<SYS>>\n"
            f"{model_input.pre_prompt}\n"
            "<</SYS>>\n"
            "\n"
            f"{model_input.attack_or_access_code} [/INST] [INST] {model_input.post_prompt} [/INST]"
        )

    def _predict(
        self,
        model_input: ModelInput | DetectionModelInput,
        temperature: float,
        max_length_tokens: int,
        prompt_type: LLMPromptType = LLMPromptType.PREDICT,
    ) -> str:
        """
        Given a single instance of an attack attempt,
        we replay the same attack on the model and return the output

        model_input: class ModelInput

        Example DeepInfra API Documentation: https://deepinfra.com/meta-llama/Llama-2-70b-chat-hf/api
        The important thing to note is that input is expected to be one large string
        """
        self.headers["Authorization"] = "Bearer " + os.environ.get("DEEPINFRA_API_KEY")
        # Deepinfra throws an error if temp is exactly 0
        # (using values <= 1e-4 seems to make the model only output empty
        # strings, possibly because we're overflowing the int8s or float16s or
        # whatever they're using)
        temperature = max(1e-2, temperature)
        input = self._get_prompt(model_input, prompt_type)
        request_body = {
            "input": input,
            "temperature": temperature,
            "max_new_tokens": max_length_tokens,
        }

        max_retries = 10
        retry_strategy = Retry(
            total=max_retries,
            method_whitelist=["POST"],
            backoff_factor=0.5,  # this will sleep {backoff factor} * (2 ** ({number of total retries} - 1)) seconds
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session = requests.Session()
        session.mount("https://", adapter)

        start_time = time.time()
        response = session.post(self.endpoint, json=request_body, headers=self.headers)
        log.info(f"DeepInfra API query took {(time.time() - start_time) * 1000} ms")

        if response.status_code != 200:
            log.error(
                f"After {max_retries} retries, DeepInfra API returned error code {response.status_code} with message {response.text}"
            )
            raise Exception(response.text)

        return response.json()["results"][0]["generated_text"]
