########################################################################################################################
# IMPORTS

import logging
import warnings
from typing import List, Optional


class PiiDependenciesMissingError(ImportError):
    pass


class SpacyModelNotFoundError(ImportError):
    pass


try:
    import spacy  # type: ignore
    from presidio_analyzer import AnalyzerEngine, RecognizerRegistry  # type: ignore
    from presidio_analyzer.nlp_engine import NlpEngineProvider  # type: ignore
    from presidio_analyzer.predefined_recognizers import GLiNERRecognizer  # type: ignore
    from presidio_anonymizer import AnonymizerEngine  # type: ignore
except ImportError as e:
    raise PiiDependenciesMissingError(
        "One or more PII anonymization dependencies are missing. "
        "Please install them by running: pip install datamarket[pii]\n"
        f"Original error: {e}"
    ) from e


########################################################################################################################
# SETTINGS

logger = logging.getLogger(__name__)
logging.getLogger("presidio-analyzer").setLevel(logging.ERROR)

warnings.filterwarnings(
    "ignore",
    message=r"\[W108\]",
    category=UserWarning,
    module="spacy.pipeline.lemmatizer",
)

# GLiNER entity labels → Presidio entity types
_GLINER_ENTITY_MAPPING = {
    "person": "PERSON",
    "name": "PERSON",
    "phone number": "PHONE_NUMBER",
    "telephone": "PHONE_NUMBER",
    "mobile phone number": "PHONE_NUMBER",
    "landline phone number": "PHONE_NUMBER",
    "fax number": "PHONE_NUMBER",
    "email": "EMAIL_ADDRESS",
    "email address": "EMAIL_ADDRESS",
    "username": "USERNAME",
    "social media handle": "USERNAME",
}

########################################################################################################################
# CLASSES


class PiiAnonymizer:
    # Backbone spaCy model used purely for tokenization — GLiNER handles NER for all languages
    _SPACY_MODEL = "en_core_web_sm"
    # Language registered with Presidio — must match the backbone; GLiNER ignores it
    _PRESIDIO_LANG = "en"

    def __init__(self):
        if not spacy.util.is_package(self._SPACY_MODEL):
            raise SpacyModelNotFoundError(
                f"spaCy model '{self._SPACY_MODEL}' not found. "
                f"Please install it by running: python -m spacy download {self._SPACY_MODEL}"
            )

        self.anonymizer = AnonymizerEngine()
        self.analyzer = self._load_analyzer_engine()

    @staticmethod
    def _nlp_config() -> dict:
        # Single lightweight English model used only as a tokenizer backbone.
        # GLiNER is multilingual and does not rely on spaCy NER or language-specific features.
        return {
            "nlp_engine_name": "spacy",
            "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}],
        }

    def _load_analyzer_engine(self) -> AnalyzerEngine:
        provider = NlpEngineProvider(nlp_configuration=self._nlp_config())
        nlp_engine = provider.create_engine()

        registry = RecognizerRegistry(supported_languages=[self._PRESIDIO_LANG])
        registry.load_predefined_recognizers(
            nlp_engine=nlp_engine,
            languages=[self._PRESIDIO_LANG],
        )

        # Remove spaCy's built-in NER recognizer — GLiNER replaces it entirely
        registry.remove_recognizer("SpacyRecognizer")

        # GLiNER multilingual PII model — zero-shot, CPU-friendly (~460 MB), Apache 2.0
        gliner_recognizer = GLiNERRecognizer(
            model_name="urchade/gliner_multi_pii-v1",
            entity_mapping=_GLINER_ENTITY_MAPPING,
            flat_ner=False,
            multi_label=True,
            map_location="cpu",
        )
        registry.add_recognizer(gliner_recognizer)

        return AnalyzerEngine(
            registry=registry,
            nlp_engine=nlp_engine,
            supported_languages=[self._PRESIDIO_LANG],
        )

    def anonymize_text(self, text: str, entities: Optional[List[str]] = None) -> str:
        """Anonymize PII in *text* using GLiNER. Handles any language without configuration."""
        analyzer_result = self.analyzer.analyze(
            text=text,
            entities=entities,
            language=self._PRESIDIO_LANG,
        )
        anonymizer_result = self.anonymizer.anonymize(text=text, analyzer_results=analyzer_result)
        return anonymizer_result.text
