from langchain.tools import BaseTool
from experiments.defense_wrapper import wrap_external_document
from utils.constants import SANDWICH, DELIMITING_DATA, MARKING


class EmailReaderTool(BaseTool):
    name: str = "email_reader"
    description: str = "Read an incoming email and return its contents."
    attack_algorithm: str = ""
    user_instruction: str = ""
    defense_type: str = ""
    counterattack: bool = False

    def __init__(
        self,
        user_instruction: str,
        attack_algorithm: str,
        defense_type: str = "",
        counterattack: bool = False,
    ):
        super().__init__()
        self.user_instruction = user_instruction
        self.attack_algorithm = attack_algorithm
        self.defense_type = defense_type
        self.counterattack = counterattack

    def _run(self) -> str:
        """Synchronously fetch and return website content."""
        email = self.attack_algorithm.replace("\t", "")
        defensed_contents = wrap_external_document(
            self.defense_type, email, self.user_instruction, self.counterattack
        )
        return defensed_contents

    def _arun(self, url: str):
        raise NotImplementedError(
            "SafeWebsiteVisitorTool does not support async"
        )
