import re
import spacy
from typing import List, Dict, Tuple

# Load spaCy model for NER (en_core_web_sm)
nlp = spacy.load("en_core_web_sm")

# Regex patterns for PII fields
PATTERNS = {
    "email": re.compile(r"[a-zA-Z0-9+._%-]+@[a-zA-Z0-9._%-]+\.[a-zA-Z]{2,6}"),
    "phone_number": re.compile(r"\b(?:\+?\d{1,3}[ -]?)?(?:\(\d{3}\)|\d{3})[ -]?\d{3}[ -]?\d{4}\b"),
    "dob": re.compile(r"\b\d{2}[/-]\d{2}[/-]\d{4}\b"),
    "aadhar_num": re.compile(r"\b\d{4}[ -]?\d{4}[ -]?\d{4}\b"),
    "credit_debit_no": re.compile(r"\b\d{4}[ -]?\d{4}[ -]?\d{4}[ -]?\d{4}\b"),
    "cvv_no": re.compile(r"\b\d{3}\b"),
    "expiry_no": re.compile(r"\b(0[1-9]|1[0-2])\/[0-9]{2}\b"),
}

Entity = Dict[str, object]


def mask_pii(text: str) -> Tuple[str, List[Entity]]:
    """
    Masks PII in the text. Returns masked_text and list of entities with their positions and original values.
    """
    entities: List[Entity] = []
    masked_text = text

    # 1) Regex-based masking
    for ent_type, pattern in PATTERNS.items():
        for match in pattern.finditer(masked_text):
            start, end = match.span()
            original = match.group()
            entities.append({
                "position": [start, end],
                "classification": ent_type,
                "entity": original,
            })
            masked_text = masked_text[:start] + f"[{ent_type}]" + masked_text[end:]

    # 2) spaCy NER for PERSON and DATE
    doc = nlp(masked_text)
    offset = 0
    for ent in doc.ents:
        label = ent.label_
        if label in ("PERSON", "DATE"):
            ent_type = "full_name" if label == "PERSON" else "dob"
            start, end = ent.start_char, ent.end_char
            original = masked_text[start:end]
            entities.append({
                "position": [start, end],
                "classification": ent_type,
                "entity": original,
            })
            masked_text = masked_text[:start] + f"[{ent_type}]" + masked_text[end:]

    # Sort entities by start index
    entities.sort(key=lambda x: x["position"][0])
    return masked_text, entities


def demask_pii(masked_text: str, entities: List[Entity]) -> str:
    """
    Restores original PII values into the masked_text using stored entity list.
    """
    result = masked_text
    # Apply replacements in reverse order to preserve indices
    for ent in sorted(entities, key=lambda x: x["position"][0], reverse=True):
        start, end = ent["position"]
        original = ent["entity"]
        # Replace the [ent_type] placeholder back to original
        result = result[:start] + original + result[end:]
    return result