"""
import json
from enum import Enum

from ttdata.common import DetectionModelInput, LLMPromptType, ModelAdapter, ModelInput


class RobertaModel(Enum):
    ROBERTA_LARGE = "roberta-large"


class RobertaAdapter(ModelAdapter):
    CONVERSION_DICT = {0: "No", 1: "Yes"}

    def __init__(self, model: RobertaModel):
        super().__init__()
        self._model = model.value

    def get_cache_key(
        self,
        model_input: ModelInput | DetectionModelInput,
        temperature: float = 0,
        max_length_tokens: int = 256,
        prompt_type: LLMPromptType = LLMPromptType.DETECT_ZERO_SHOT,
    ) -> str:
        return json.dumps(
            (
                self._model,
                self._get_messages(model_input, prompt_type),
                temperature,
                max_length_tokens,
            ),
        )

    def _get_messages(self, model_input: DetectionModelInput, prompt_type: LLMPromptType) -> list[str]:
        return [
            "[CLS]" + model_input.access_code + "[SEP]" + model_input.attack_llm_output + "[SEP]"
        ]

    def _predict(self, model_input: ModelInput | DetectionModelInput, prompt_type: LLMPromptType = LLMPromptType.DETECT_ZERO_SHOT) -> str:
        device = (
            "mps"
            if torch.backends.mp.is_available()
            else "cuda"
            if torch.cuda.is_available()
            else "cpu"
        )
        device = torch.device(device)

        model = RobertaForSequenceClassification.from_pretrained(
            "roberta_prompt_extraction_model"
        ).to(device)
        tokenizer = RobertaTokenizer.from_pretrained("roberta-large")

        encoding = self.tokenizer(
            model_input,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt",
        )
        tokenized_example = {
            "input_ids": encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
        }
        input_ids = encoding["input_ids"].squeeze().to(device)
        attention_mask = encoding["attention_mask"].squeeze().to(device)

        outputs = model(input_ids, attention_mask=attention_mask)
        logits = outputs.logits

        _, predicted = torch.max(logits, dim=1)

        return RobertaAdapter.CONVERSION_DICT[predicted.cpu().numpy()[0]]
"""
