import time
from typing import Dict, Any

import presidio_analyzer as analyzer
import presidio_anonymizer as anonymizer
import spacy
from loguru import logger

from deepdata.core.base_op import Mapper

from .custom_entities import id_recognizer, phone_recognizer, zipcode_recognizer, url_recognizer


class PiiDetector(Mapper):
    custom_ops = True

    def __init__(self, *args, **kwargs):
        super(PiiDetector, self).__init__(*args, **kwargs)
        self.support_language = kwargs.get("support_language", "zh")

        self.nlp_engine = None
        self.text_analyzer = None
        self.anom = None

        self.init_model(*args, **kwargs)

    def init_model(self, *args, **kwargs):
        try:
            spacy.load("zh_core_web_sm")
        except OSError:
            logger.warning(
                "zh_core_web_sm model not found. Please run: python -m spacy download zh_core_web_sm"
            )
            raise

        provider = analyzer.nlp_engine.NlpEngineProvider(
            nlp_configuration={
                "nlp_engine_name": "spacy",
                "models": [
                    {"lang_code": "zh", "model_name": "zh_core_web_sm"}
                ]
            }
        )

        # 创建NLP Engine
        self.nlp_engine = provider.create_engine()

        #  初始化AnalyzerEngine
        self.text_analyzer = analyzer.AnalyzerEngine(nlp_engine=self.nlp_engine, supported_languages=["zh"])
        self.text_analyzer.registry.load_predefined_recognizers()
        for recognizer in [id_recognizer, phone_recognizer, zipcode_recognizer, url_recognizer]:
            self.text_analyzer.registry.add_recognizer(recognizer)

        # 初始化AnonymizerEngine
        self.anom = anonymizer.AnonymizerEngine()

    def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        start = time.time()
        self.read_file_first(sample)
        text = sample.get(self.text_key, "")
        if not text:
            logger.warning(f"fileName: {sample.get(self.filename_key, 'unknown')}, text is empty, skip PII detection.")
            return sample

        try:
            analyzer_results = self.text_analyzer.analyze(text=text, language=self.support_language)
            res = self.anom.anonymize(text=text, analyzer_results=analyzer_results)
            sample[self.text_key] = res.text
        except Exception as e:
            logger.error(
                f"fileName: {sample.get(self.filename_key, 'unknown')}, "
                f"PII detection failed: {e}. Returning original text."
            )

        logger.info(
            f"fileName: {sample.get(self.filename_key, 'unknown')}, "
            f"method: PiiDetector costs {time.time() - start:.6f} s"
        )
        return sample