"""PII (Personally Identifiable Information) scanner for Tru-Guard.

Uses Microsoft Presidio with spaCy for PII detection.
"""

from typing import Any

from trusys.scanners import BaseScanner, register_scanner
from trusys.types import (
    ScanContext,
    ScanResult,
    ScanType,
    Severity,
    Violation,
)

# Default PII entities that Presidio can detect
DEFAULT_ENTITIES = [
    "PHONE_NUMBER",
    "EMAIL_ADDRESS",
    "CREDIT_CARD",
    "PERSON",
    "LOCATION",
    "DATE_TIME",
    "NRP",  # Nationality, Religious or Political group
    "MEDICAL_LICENSE",
    "URL",
    "IP_ADDRESS",
    "US_SSN",
    "US_BANK_NUMBER",
    "US_DRIVER_LICENSE",
    "US_ITIN",
    "US_PASSPORT",
    "UK_NHS",
    "IBAN_CODE",
    "CRYPTO",
]


@register_scanner(
    name="pii",
    scan_types=["input", "output"],
    description="Detects PII (Personally Identifiable Information) using Microsoft Presidio with spaCy",
    validator_type="pii",
)
class PIIScanner(BaseScanner):
    """Scanner that detects PII using Microsoft Presidio.

    Uses spaCy as the NLP engine for entity recognition.
    Supports configurable entity types to detect.

    Args:
        config: Configuration dictionary with:
               - entities (list[str], optional): List of entity types to detect.
                 Defaults to common PII types like PHONE_NUMBER, EMAIL_ADDRESS, etc.
               - language (str, optional): Language code for analysis (default: "en")
               - score_threshold (float, optional): Minimum confidence score to report (default: 0.5)
        on_fail: Action to take when PII is detected
        enabled: Whether this scanner is enabled
    """

    # Lazy-loaded analyzer instance (shared across scanner instances)
    _analyzer = None
    _analyzer_language = None

    def __init__(
        self,
        config: dict[str, Any] | None = None,
        on_fail: str | None = None,
        enabled: bool | None = None,
    ) -> None:
        """Initialize the PII scanner.

        Args:
            config: Configuration dictionary (or list of entity strings) containing:
                   - entities (list[str], optional): Entity types to detect
                   - language (str, optional): Language code (default: "en")
                   - score_threshold (float, optional): Minimum confidence score (default: 0.5)
                   If config is a list of strings, it is treated as the entities list.
            on_fail: Action to take when PII is detected
            enabled: Whether this scanner is enabled
        """
        # Normalize config: API may send list but we expect dict
        if config is None:
            config = {}
        elif not isinstance(config, dict):
            if isinstance(config, list):
                if config and isinstance(config[0], dict):
                    # List of config objects (e.g. from API) - use first
                    config = config[0]
                elif config and all(isinstance(x, str) for x in config):
                    # List of entity type strings
                    config = {"entities": config}
                else:
                    config = {}
            else:
                config = {}

        super().__init__(config=config, on_fail=on_fail, enabled=enabled)

        # Extract configuration with defaults
        self.entities = config.get("entities", DEFAULT_ENTITIES)
        self.language = config.get("language", "en")
        self.score_threshold = config.get("score_threshold", 0.5)

        # Validate entities is a list
        if not isinstance(self.entities, list):
            raise ValueError("'entities' must be a list of entity type strings")
        
        if len(self.entities) == 0:
            raise ValueError("'entities' must contain at least one entity type")

        # Validate score_threshold
        if not isinstance(self.score_threshold, (int, float)):
            raise ValueError("'score_threshold' must be a number between 0 and 1")
        if not 0 <= self.score_threshold <= 1:
            raise ValueError("'score_threshold' must be between 0 and 1")

    @classmethod
    def _get_analyzer(cls, language: str = "en"):
        """Get or create the Presidio analyzer (lazy initialization).

        Args:
            language: Language code for analysis

        Returns:
            AnalyzerEngine instance
        """
        if cls._analyzer is None or cls._analyzer_language != language:
            try:
                from presidio_analyzer import AnalyzerEngine
            except ImportError as e:
                raise ImportError(
                    "presidio-analyzer is required for PII scanning. "
                    "Install it with: pip install presidio-analyzer && "
                    "python -m spacy download en_core_web_lg"
                ) from e

            # Avoid Presidio's INFO log "nlp_engine not provided, creating default."
            import logging
            logging.getLogger("presidio-analyzer").setLevel(logging.WARNING)
            cls._analyzer = AnalyzerEngine()
            cls._analyzer_language = language

        return cls._analyzer

    def scan(self, content: dict[str, Any], context: ScanContext) -> ScanResult:
        """Perform PII scan on content using Presidio.

        Args:
            content: JSON dictionary containing content to scan.
                   Should contain 'prompt' for input scans or 'response' for output scans.
            context: The scan context

        Returns:
            ScanResult with detected PII as violations
        """
        # Extract the actual text to scan based on scan type
        if context.scan_type == ScanType.INPUT:
            text_to_scan = content.get("prompt", "")
        elif context.scan_type == ScanType.OUTPUT:
            text_to_scan = content.get("response", "")
        else:
            # Fallback: try prompt first, then response
            text_to_scan = content.get("prompt") or content.get("response", "")

        if not text_to_scan:
            return ScanResult(
                passed=True,
                scanner_name=self.name,
                scan_type=context.scan_type,
                violations=[],
                metadata={
                    "entities_checked": self.entities,
                    "pii_found": [],
                    "assertion": {"type": "pii", "output": ""},
                    "score": 1.0,
                    "reason": "",
                    "latency": 0.0,
                    "action": {"type": self.on_fail.value if self.on_fail else "ignore"},
                },
            )

        # Get analyzer and perform analysis
        analyzer = self._get_analyzer(self.language)
        
        try:
            results = analyzer.analyze(
                text=text_to_scan,
                entities=self.entities,
                language=self.language,
                score_threshold=self.score_threshold,
            )
        except Exception as e:
            # Return error as violation
            return ScanResult(
                passed=False,
                scanner_name=self.name,
                scan_type=context.scan_type,
                violations=[
                    Violation(
                        rule="pii_analysis_error",
                        severity=Severity.HIGH,
                        message=f"PII analysis failed: {e}",
                    )
                ],
                metadata={
                    "error": str(e),
                    "assertion": {"type": "pii", "output": f"error: {e}"},
                    "score": 0.0,
                    "reason": f"PII analysis failed: {e}",
                    "latency": 0.0,
                    "action": {"type": self.on_fail.value if self.on_fail else "ignore"},
                },
            )

        # Remove overlapping Presidio results (keep longest span; on tie keep higher score)
        results = self._deduplicate_results(results)

        # Convert Presidio results to violations
        violations = []
        pii_found = []

        for result in results:
            # Extract the matched text
            matched_text = text_to_scan[result.start:result.end]
            
            pii_info = {
                "entity_type": result.entity_type,
                "start": result.start,
                "end": result.end,
                "score": result.score,
                "matched_text": matched_text,
            }
            pii_found.append(pii_info)

            violations.append(
                Violation(
                    rule=f"pii_detected_{result.entity_type.lower()}",
                    severity=Severity.HIGH,
                    message=f"PII detected: {result.entity_type} (confidence: {result.score:.2f})",
                    span=(result.start, result.end),
                    metadata={
                        "entity_type": result.entity_type,
                        "confidence": result.score,
                        "matched_text": matched_text,
                    },
                )
            )

        passed = len(violations) == 0

        # Build reason from violations
        reason = ""
        if not passed:
            entity_types = list(set(v.metadata.get("entity_type", "") for v in violations))
            reason = f"PII detected: {', '.join(entity_types)}"

        # Build assertion output - list all matches
        assertion_output = ""
        if violations:
            parts = []
            for v in violations:
                entity_type = v.metadata.get("entity_type", "")
                matched_text = v.metadata.get("matched_text", "")
                if matched_text:
                    parts.append(f"entity_type': '{entity_type}', 'matched_text': '{matched_text}'")
            assertion_output = "; ".join(parts) if parts else ""

        # Calculate score (1.0 if passed, otherwise based on number of violations)
        score = 1.0 if passed else max(0.0, 1.0 - (len(violations) * 0.1))

        # Build metadata
        metadata = {
            "entities_checked": self.entities,
            "language": self.language,
            "score_threshold": self.score_threshold,
            "pii_found": pii_found,
            "total_pii_count": len(pii_found),
            "assertion": {"type": "pii", "output": assertion_output},
            "score": score,
            "reason": reason,
            "latency": 0.0,  # Will be set by execution runner
            "action": {"type": self.on_fail.value if self.on_fail else "ignore"},
        }

        # Generate suggested fix by masking PII
        suggested_fix = None
        if violations:
            suggested_fix = self._generate_masked_content(text_to_scan, results)

        return ScanResult(
            passed=passed,
            scanner_name=self.name,
            scan_type=context.scan_type,
            violations=violations,
            suggested_fix=suggested_fix,
            metadata=metadata,
        )

    @staticmethod
    def _deduplicate_results(pii_results: list) -> list:
        """Remove overlapping Presidio results, keeping the longest span.

        When two results overlap (e.g. EMAIL_ADDRESS at 9-25 and URL at 14-25),
        only the longer span is kept.  On equal length the higher-confidence
        result wins.

        Args:
            pii_results: List of Presidio RecognizerResult objects

        Returns:
            Filtered list with no overlapping spans
        """
        # Sort by span length descending, then by confidence descending
        sorted_results = sorted(
            pii_results,
            key=lambda x: (-(x.end - x.start), -x.score),
        )

        accepted: list[tuple[int, int]] = []  # (start, end) of kept spans
        kept = []
        for result in sorted_results:
            start, end = result.start, result.end
            # Skip if this span overlaps any already-accepted span
            if any(
                not (end <= a_start or start >= a_end)
                for a_start, a_end in accepted
            ):
                continue
            accepted.append((start, end))
            kept.append(result)

        return kept

    def _generate_masked_content(self, text: str, pii_results: list) -> str:
        """Generate content with PII masked.

        Args:
            text: Original text
            pii_results: List of *deduplicated* Presidio RecognizerResult objects

        Returns:
            Text with PII replaced by entity type markers
        """
        # Apply replacements right-to-left so indices stay valid
        sorted_results = sorted(pii_results, key=lambda x: x.start, reverse=True)
        masked_text = text
        for result in sorted_results:
            replacement = f"[{result.entity_type}]"
            masked_text = masked_text[:result.start] + replacement + masked_text[result.end:]

        return masked_text

    def get_config(self) -> dict[str, Any]:
        """Get the scanner configuration.

        Returns:
            Dictionary of configuration values
        """
        config = super().get_config()
        config["entities"] = self.entities
        config["language"] = self.language
        config["score_threshold"] = self.score_threshold
        return config

    def __repr__(self) -> str:
        """String representation of the scanner."""
        entities_repr = ", ".join(f"'{e}'" for e in self.entities[:3])
        if len(self.entities) > 3:
            entities_repr += f", ... (+{len(self.entities) - 3} more)"
        return (
            f"{self.__class__.__name__}("
            f"entities=[{entities_repr}], "
            f"language={self.language!r}, "
            f"score_threshold={self.score_threshold}, "
            f"on_fail={self.on_fail.value!r})"
        )
