# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""PII/PHI de-identification via Presidio.

Integrated into KliraOTLPSpanExporter as a pre-export hook.
Scans protobuf span attributes for PII/PHI and applies de-identification
using the configured method (redact/mask/replace/hash).

Works for any PII type — not limited to healthcare.
"""

from __future__ import annotations

import logging
from typing import Any

from klira.sdk.contracts.phi_pipeline import (
    PhiMethod,
    PhiScanResult,
)

logger = logging.getLogger(__name__)

SPACY_MODEL = "en_core_web_lg"


class PhiScanner:
    """Scans text for PII/PHI using Presidio.

    Presidio is an optional dependency (install with klira[anonymization]).
    When anonymization=True, Presidio scans and de-identifies text.
    When False, this class is not invoked (zero runtime cost).
    """

    def __init__(self, method: PhiMethod = PhiMethod.REDACT) -> None:
        self._method = method
        self._analyzer: Any = None
        self._anonymizer: Any = None
        self._operator_config_cls: Any = None
        self._initialized = False

    def _ensure_spacy_model(self) -> None:
        """Download the spaCy NLP model if not already installed."""
        import spacy.util  # type: ignore[import-untyped,import-not-found]

        if spacy.util.is_package(SPACY_MODEL):
            return

        logger.info(
            "spaCy model '%s' not found — downloading (this is a one-time operation)...",
            SPACY_MODEL,
        )
        try:
            from spacy.cli import download  # type: ignore[import-untyped,import-not-found]

            download(SPACY_MODEL)
        except Exception as e:
            raise RuntimeError(
                f"Failed to download spaCy model '{SPACY_MODEL}'. "
                f"Install manually with: python -m spacy download {SPACY_MODEL}"
            ) from e

    def _ensure_initialized(self) -> None:
        """Lazy-initialize Presidio engines.

        Raises:
            ImportError: If presidio packages are not installed.
            RuntimeError: If Presidio engines fail to initialize
                or spaCy model cannot be downloaded.
        """
        if self._initialized:
            return

        try:
            from presidio_analyzer import AnalyzerEngine  # type: ignore[import-untyped,import-not-found]
            from presidio_anonymizer import AnonymizerEngine  # type: ignore[import-untyped,import-not-found]
            from presidio_anonymizer.entities import OperatorConfig  # type: ignore[import-untyped,import-not-found]
        except ImportError as e:
            raise ImportError(
                "Presidio is required for anonymization. "
                "Install with: pip install 'klira[anonymization]'"
            ) from e

        self._ensure_spacy_model()

        try:
            self._analyzer = AnalyzerEngine()
            self._anonymizer = AnonymizerEngine()
            self._operator_config_cls = OperatorConfig
            self._initialized = True
        except Exception as e:
            raise RuntimeError(
                "Presidio initialization failed. "
                "Install with: pip install 'klira[anonymization]'"
            ) from e

    def scan(self, text: str) -> tuple[PhiScanResult, list[Any]]:
        """Scan text for PII/PHI entities.

        Returns a tuple of (PhiScanResult, raw_analyzer_results).
        The raw results can be passed to deidentify() to avoid re-analysis.
        """
        self._ensure_initialized()

        results = self._analyzer.analyze(text=text, language="en")

        if not results:
            return (
                PhiScanResult(
                    detected=False,
                    entity_count=0,
                    entity_types=(),
                ),
                results,
            )

        entity_types = tuple(sorted({r.entity_type for r in results}))
        return (
            PhiScanResult(
                detected=True,
                entity_count=len(results),
                entity_types=entity_types,
            ),
            results,
        )

    def _build_operators(self) -> dict[str, Any]:
        """Build Presidio operators dict for the configured PhiMethod."""
        OperatorConfig = self._operator_config_cls
        if self._method == PhiMethod.REDACT:
            return {"DEFAULT": OperatorConfig("redact", {})}
        if self._method == PhiMethod.MASK:
            return {
                "DEFAULT": OperatorConfig(
                    "mask",
                    {"masking_char": "*", "chars_to_mask": 100, "from_end": False},
                )
            }
        if self._method == PhiMethod.HASH:
            return {"DEFAULT": OperatorConfig("hash", {"hash_type": "sha256"})}
        # PhiMethod.REPLACE — Presidio default: <ENTITY_TYPE>
        return {"DEFAULT": OperatorConfig("replace", {})}

    def deidentify(
        self,
        text: str,
        analyzer_results: list[Any] | None = None,
    ) -> str:
        """De-identify text by applying the configured PHI method.

        Args:
            text: Text to de-identify.
            analyzer_results: Pre-computed Presidio analyzer results from scan().
                If None, runs analysis again (backward compatible).
        """
        self._ensure_initialized()

        if analyzer_results is None:
            analyzer_results = self._analyzer.analyze(text=text, language="en")
        if not analyzer_results:
            return text

        anonymized = self._anonymizer.anonymize(
            text=text,
            analyzer_results=analyzer_results,
            operators=self._build_operators(),
        )
        return anonymized.text  # type: ignore[no-any-return]
