src/shield/redact.rs

use tracing::instrument;
use serde::{Serialize, Deserialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PiiMatch {
    pub entity_type: String,
    pub original: String,
    pub placeholder: String,
}

#[instrument(skip_all)]
pub fn redact_pii(raw_text: &str) -> (String, Vec<PiiMatch>) {
    let custom_patterns: Vec<(&str, &str)> = vec![
        (r"\b(?:Dr\.|Dr|Professor|Prof\.)\s+[A-Z][a-z]+\b", "PROVIDER_NAME"),
        // Match 2+ capitalized words (e.g., "John Smith", "Patient John Smith")
        (r"\b[A-Z][a-z]+(?: [A-Z][a-z]+)+\b", "PERSON_NAME"),
        (r"\b\d{1,2}/\d{1,2}/\d{2,4}\b", "DATE"),
        (r"\b\d{3}-\d{2}-\d{4}\b", "SSN"),
        (r"\b\d{10}\b", "PHONE"),
        (r"\b[A-Z]{2}\d{2}\s?\d{2}\s?\d{2}\s?\d\b", "NHS_NUMBER"),
        (r"\b\d{3}-\d{4}\b", "ZIP"),
        (r"\b\d{2,3}\s?(?:years?|yo)\b", "AGE"),
        (r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b", "EMAIL"),
    ];

    let compiled: Vec<(regex::Regex, &str)> = custom_patterns
        .iter()
        .filter_map(|(p, label)| regex::Regex::new(p).ok().map(|re| (re, *label)))
        .collect();

    let mut pii_matches = Vec::new();
    let mut redacted = raw_text.to_string();
    let mut counter: std::collections::HashMap<String, usize> = std::collections::HashMap::new();

    let raw_bytes = raw_text.as_bytes();
    let mut all_matches: Vec<(usize, usize, &regex::Regex, &str)> = Vec::new();
    for (re, entity_type) in &compiled {
        for mat in re.find_iter(raw_text) {
            all_matches.push((mat.start(), mat.end(), re, *entity_type));
        }
    }

    // Sort by start, then longest match first (to prefer "Patient John Smith" over "Patient John")
    all_matches.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| b.1.cmp(&a.1)));

    // Remove overlapping matches
    let mut used = vec![false; raw_bytes.len()];
    let mut filtered = Vec::new();
    for (start, end, re, entity_type) in all_matches {
        let range = start..end;
        if used[range.clone()].iter().any(|&b| b) {
            continue;
        }
        used[range].fill(true);
        filtered.push((start, end, re, entity_type));
    }

    // Process in reverse order for safe replacement
    filtered.sort_by(|a, b| b.0.cmp(&a.0));
    for (start, end, _re, entity_type) in filtered {
        let original = &raw_text[start..end];
        let count = counter.entry(entity_type.to_string()).or_insert(0);
        *count += 1;
        let placeholder = format!("[{}_{}]", entity_type, count);
        pii_matches.push(PiiMatch {
            entity_type: entity_type.to_string(),
            original: original.to_string(),
            placeholder: placeholder.clone(),
        });
        redacted.replace_range(start..end, &placeholder);
    }

    tracing::debug!("Redacted {} PII tokens", pii_matches.len());
    (redacted, pii_matches)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn redact_pii_basic() {
        let text = "Patient John Smith, 45 yo, phone 5551234567, SSN 123-45-6789.";
        let (redacted, matches) = redact_pii(text);

        // All PII tokens must be gone
        assert!(!redacted.contains("John"));
        assert!(!redacted.contains("Smith"));
        assert!(!redacted.contains("5551234567"));
        assert!(!redacted.contains("123-45-6789"));

        // Placeholders present
        assert!(redacted.contains("[PERSON_NAME_1]"));
        assert!(redacted.contains("[PHONE_1]"));
        assert!(redacted.contains("[SSN_1]"));

        // Verify we captured the full name
        assert!(matches.iter().any(|m| m.original.contains("John Smith") || m.original.contains("Patient John Smith")));
    }
}