from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import OperatorConfig


class PIIGuardrail:
    """
    Bidirectional PII detection and redaction using Microsoft Presidio
    - INPUT: Detect PII in user queries before sending to LLM
    - OUTPUT: Redact PII in LLM responses before showing to user
    """

    def __init__(self, language='es'):
        self.analyzer = AnalyzerEngine()
        self.anonymizer = AnonymizerEngine()
        self.language = language

        # Entity types to detect (configurable per use case)
        self.entity_types = [
            "PERSON",
            "EMAIL_ADDRESS",
            "PHONE_NUMBER",
            "LOCATION",
            "CREDIT_CARD",
            "IBAN_CODE",
            "IP_ADDRESS",
            "US_SSN",
            "MEDICAL_LICENSE",
            "DATE_TIME",
            "NRP"  # NRP = Spanish national ID
        ]

    def detect_pii(self, text: str) -> list:
        """
        Detect PII entities in text
        
        Returns:
            List of detected entities with type, score, and location
        """
        results = self.analyzer.analyze(
            text=text,
            entities=self.entity_types,
            language=self.language
        )
        return results

    def redact_pii(self, text: str, redaction_strategy='replace') -> dict:
        """
        Redact PII from text
        
        Strategies: 'replace' (with placeholder), 'mask' (with ***), 'hash'
        """
        analyzer_results = self.detect_pii(text)

        if not analyzer_results:
            return {
                'redacted_text': text,
                'entities_found': [],
                'pii_detected': False
            }

        # Anonymize based on strategy
        if redaction_strategy == 'replace':
            operators = {
                "PERSON": OperatorConfig("replace", {"new_value": "[NOMBRE_REDACTADO]"}),
                "EMAIL_ADDRESS": OperatorConfig("replace", {"new_value": "[EMAIL_REDACTADO]"}),
                "PHONE_NUMBER": OperatorConfig("replace", {"new_value": "[TELEFONO_REDACTADO]"}),
                "CREDIT_CARD": OperatorConfig("replace", {"new_value": "[TARJETA_REDACTADA]"}),
                "DEFAULT": OperatorConfig("replace", {"new_value": "[PII_REDACTADO]"})
            }
        elif redaction_strategy == 'mask':
            operators = {
                "DEFAULT": OperatorConfig("mask", {"masking_char": "*", "chars_to_mask": 100})
            }
        elif redaction_strategy == 'hash':
            operators = {"DEFAULT": OperatorConfig("hash")}

        anonymized_result = self.anonymizer.anonymize(
            text=text,
            analyzer_results=analyzer_results,
            operators=operators
        )

        return {
            'redacted_text': anonymized_result.text,
            'entities_found': [
                {
                    'type': item.entity_type,
                    'start': item.start,
                    'end': item.end,
                    'score': item.score
                }
                for item in analyzer_results
            ],
            'pii_detected': True
        }

    def validate_safe_for_llm(self, text: str, threshold: float = 0.85) -> dict:
        """
        Check if text is safe to send to LLM (no high-confidence PII)
        """
        pii_entities = self.detect_pii(text)
        high_confidence_pii = [
            entity for entity in pii_entities
            if entity.score >= threshold
        ]

        return {
            'is_safe': len(high_confidence_pii) == 0,
            'pii_count': len(pii_entities),
            'high_confidence_pii': high_confidence_pii
        }


# ============================================================
# USO EN PRODUCCIÓN
# ============================================================

guardrail = PIIGuardrail(language='es')

# EJEMPLO 1: Validar input del usuario ANTES de enviar a LLM
user_input = "Mi nombre es Juan García y mi email es juan.garcia@empresa.com. Mi DNI es 12345678Z."

safety_check = guardrail.validate_safe_for_llm(user_input)

if not safety_check['is_safe']:
    print(f"⚠️ PII DETECTADO EN INPUT: {safety_check['high_confidence_pii']}")

    # Opción 1: Rechazar query
    # Opción 2: Redactar automáticamente antes de enviar a LLM
    redacted = guardrail.redact_pii(user_input, 'replace')
    print(f"Input redactado: {redacted['redacted_text']}")
    # Enviar redacted_text al LLM en lugar del original
else:
    print("✅ Input seguro, sin PII detectado")


# EJEMPLO 2: Redactar output del LLM ANTES de mostrar al usuario
llm_response = "Basándome en la información, el cliente Juan García (juan.garcia@empresa.com) vive en Madrid..."

output_validation = guardrail.redact_pii(llm_response, 'replace')

if output_validation['pii_detected']:
    print(f"⚠️ PII en output LLM - Redactando automáticamente")
    print(f"Response segura: {output_validation['redacted_text']}")
    # Mostrar al usuario el texto redactado
else:
    print(f"✅ Output limpio: {llm_response}")


# EJEMPLO 3: Auditoría y logging
detected_entities = guardrail.detect_pii(user_input)

print("📊 Audit Log:")
for entity in detected_entities:
    print(f"   - {entity.entity_type}: confidence {entity.score:.2f}")