import tldextract

class DataFilterManager:
    def __init__(self,target,fake_target):
        self.real_target = target
        self.real_domain = ""
        self.fake_target = fake_target
        self.init_DataFilter()

    def init_DataFilter(self):
        if self.fake_target == "czzfkjxx":
            try:
                extracted = tldextract.extract(self.real_target)
                self.real_domain = extracted.domain
            except Exception as e:
                print(f"{self.real_target}不合法！")

    def filter_prompt(self,prompt):
        if self.real_domain:
            fake_prompt = prompt.replace(self.real_domain, self.fake_target)
        else:
            fake_prompt = prompt.replace(self.real_target,self.fake_target)
        return fake_prompt

    def filter_instruction(self,instruction):
        if self.real_domain:
            real_instruction = instruction.replace(self.fake_target,self.real_domain)
        else:
            real_instruction = instruction.replace(self.fake_target, self.real_target)
        return real_instruction

    def filter_result(self,instr,result):
        if self.real_domain:
            fake_instr = instr.replace(self.real_domain, self.fake_target)
            fake_result = result.replace(self.real_domain, self.fake_target)
        else:
            fake_instr = instr.replace(self.real_target,self.fake_target)
            fake_result = result.replace(self.real_target,self.fake_target)
        return fake_instr,fake_result