import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification

class ShieldFilter:
    def __init__(self, model_path="LH-Tech-AI/Shield-82M"):
        print(f"Loading Shield-82M from {model_path}...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForTokenClassification.from_pretrained(model_path)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()

        self.group_map = {
            "FIRSTNAME": "PERSON", "MIDDLENAME": "PERSON", "LASTNAME": "PERSON",
            "BUILDINGNUMBER": "ADDRESS", "STREET": "ADDRESS", "CITY": "ADDRESS", 
            "STATE": "ADDRESS", "ZIPCODE": "ADDRESS", "SECONDARYADDRESS": "ADDRESS",
            "EMAIL": "EMAIL", "PHONENUMBER": "PHONE", "PHONEIMEI": "PHONE",
            "DATE": "DOB", "TIME": "DOB"
        }

    def protect(self, text):
        inputs = self.tokenizer(
            text, 
            return_tensors="pt", 
            truncation=True, 
            max_length=512, 
            return_offsets_mapping=True
        ).to(self.device)
        
        offsets = inputs.pop("offset_mapping")[0].cpu().numpy()
        
        with torch.no_grad():
            outputs = self.model(**inputs).logits
            
        predictions = torch.argmax(outputs, dim=2)[0].cpu().numpy()
        id2label = self.model.config.id2label
        
        spans_to_replace = []
        current_group = None
        start_char = -1
        last_char = -1
        
        for idx, (pred_id, offset) in enumerate(zip(predictions, offsets)):
            if offset[0] == 0 and offset[1] == 0:
                continue
                
            label = id2label[pred_id]
            
            if label == "O":
                if current_group is not None:
                    spans_to_replace.append((start_char, last_char, current_group))
                    current_group = None
            else:
                group_tag = self.group_map.get(label, label)
                
                if current_group != group_tag:
                    if current_group is not None:
                        spans_to_replace.append((start_char, last_char, current_group))
                    current_group = group_tag
                    start_char = offset[0]
                    
                last_char = offset[1]
                
        if current_group is not None:
            spans_to_replace.append((start_char, last_char, current_group))
            
        filtered_text = text
        for start, end, tag in sorted(spans_to_replace, key=lambda x: x[0], reverse=True):
            filtered_text = filtered_text[:start] + f"[{tag}]" + filtered_text[end:]
            
        return filtered_text

if __name__ == "__main__":
    shield = ShieldFilter()
    sample = "My name is John Doe. Email: john@example.com. Phone: +49 123 45678."
    print(f"Original: {sample}")
    print(f"Protected: {shield.protect(sample)}")