import re

# Complete regex patterns for sensitive data detection
SENSITIVE_PATTERNS = {
    "WINDOWS_KEY": r"\b([A-Z0-9]{5}-[A-Z0-9]{5}-[A-Z0-9]{5}-[A-Z0-9]{5}-[A-Z0-9]{5})\b",
    "IP_ADDRESS": r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b",
    "API_KEY": r"(?i)(?:api_key|token|secret)[\s:=]+['\"]([a-zA-Z0-9_\-]{20,})['\"]"
}


def sanitize_prompt(user_input):
    """
    Scan the incoming prompt and mask sensitive data.
    Goal: Replace sensitive data in the original text with safe tokens
    and create an audit log of the found items.
    """
    sanitized_text = user_input
    audit_log = []

    for data_type, pattern in SENSITIVE_PATTERNS.items():
        # Find all matches using re.findall()
        matches = re.findall(pattern, sanitized_text)

        if matches:
            # Replace sensitive data with safe tokens using re.sub()
            sanitized_text = re.sub(pattern, f"[REDACTED_{data_type}]", sanitized_text)

            # Log each finding into the audit log
            for match in matches:
                # If the regex uses capturing groups, re.findall might return tuples; handle accordingly
                original_value = match[0] if isinstance(match, tuple) else match
                audit_log.append({"type": data_type, "original_value": original_value})

    return sanitized_text, audit_log


# --- TESTING ---
test_prompts = [
    "Please explain this code: api_key = 'AKIAIOSFODNN7EXAMPLE'",
    "The IP address of the internal database server is 192.168.1.150, how should I optimize its access?",
    "I'm sorry to hear about the loss of your grandmother. Here's a Windows key: VK7JG-NPHTM-C97JM-9MPGT-3V66T"
]

print("--- Data Sanitization Middleware Test ---\n")
for i, prompt in enumerate(test_prompts):
    safe_prompt, findings = sanitize_prompt(prompt)

    print(f"[Test Case {i + 1}]")
    print(f"Original prompt: {prompt}")
    print(f"Sanitized prompt: {safe_prompt}")
    print(f"Audit log: {findings}\n")