import json
from enum import Enum

import httpx
from anthropic import AI_PROMPT, HUMAN_PROMPT, Anthropic
from ttdata.common import DetectionModelInput, LLMPromptType, ModelAdapter, ModelInput


class AnthropicModel(Enum):
    CLAUDE_2 = "claude-2.0"
    CLAUDE_INSTANT = "claude-instant-1.2"


class AnthropicAdapter(ModelAdapter):
    def __init__(self, model: AnthropicModel):
        super().__init__()
        self._model = 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,
                self._get_prompt(model_input, prompt_type),
                temperature,
                max_length_tokens,
            )
        )

    def _get_prompt(
        self,
        model_input: ModelInput | DetectionModelInput,
        prompt_type: LLMPromptType = LLMPromptType.PREDICT,
    ) -> str:
        # Claude expects a prompt as a single string. User and Assistant messages should
        # start with "User: " and "Assistant: " respectively, with a blank line between messages.
        # The HUMAN_PROMPT is "\n\nHuman:" and the AI_PROMPT is "\n\nAssistant:"
        # Reference: https://docs.anthropic.com/claude/reference/getting-started-with-the-api
        # Note that Olivia came up with this prompt based on the Claude-2
        # experiments in ../notebooks/tt_test_anthropic_prompts.ipynb.

        if prompt_type == LLMPromptType.DETECT_FEW_SHOT:
            raise NotImplementedError()
        elif prompt_type == LLMPromptType.DETECT_ZERO_SHOT:
            raise NotImplementedError()
        assert prompt_type == LLMPromptType.PREDICT
        return self._get_prediction_messages(model_input)

    def _get_prediction_messages(self, model_input: ModelInput) -> str:
        return (
            f"{HUMAN_PROMPT} {model_input.pre_prompt}"
            f"\n\n<user_input>{model_input.attack_or_access_code}</user_input>"
            f"\n\n{model_input.post_prompt}{AI_PROMPT}"
        )

    def _predict(
        self,
        model_input: ModelInput | DetectionModelInput,
        temperature: float,
        max_length_tokens: int,
        prompt_type: LLMPromptType = LLMPromptType.PREDICT,
    ) -> str:
        # we start a new client each time because Anthropic was previously using
        # tonnes of file handles when we didn't close it
        limits = httpx.Limits(max_connections=1, max_keepalive_connections=1)
        with Anthropic(connection_pool_limits=limits, max_retries=12) as anthropic:
            # anthropic does exponential backoff (unsure what the delay is); 12 was
            # found by increasing the number of retries until we stopped getting
            # errors
            completion = anthropic.completions.create(
                model=self._model,
                prompt=self._get_prompt(model_input, prompt_type),
                temperature=temperature,
                max_tokens_to_sample=max_length_tokens,
            )
        return completion.completion
