"""Named Entity Recognition (NER) based PII detection method."""

from __future__ import annotations

from typing import Any

from mankinds_eval.core import MethodResult, Sample
from mankinds_eval.methods.base import Method

# Check for optional dependencies
try:
    from transformers import pipeline as transformers_pipeline

    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    transformers_pipeline = None  # type: ignore[misc, assignment]

try:
    import torch

    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    torch = None  # type: ignore[assignment]


class PIIDetection(Method):
    """Detect PII/named entities in text.

    This method uses a NER model to detect personally identifiable information
    (PII) such as names, emails, phone numbers, addresses, etc.

    Attributes:
        name: Method identifier.
        required_fields: Fields required in the Sample.
    """

    name: str = "PIIDetection"
    version: str = "0.1.0"
    required_fields: list[str] = ["input", "output"]

    def __init__(
        self,
        model: str | None = None,
        entities: list[str] | None = None,
        fail_on_detection: bool = True,
        target: str = "output",
        device: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialize the PII detection method.

        Args:
            model: Model name/path. If None, uses default from config.
            entities: List of entity types to detect. If None, detects all.
                Common types: EMAIL, PHONE, NAME, ADDRESS, SSN, CREDIT_CARD, etc.
            fail_on_detection: If True, passed=False when PII is found.
                If False, passed is always True (just reports findings).
            target: Which field to analyze ("output" or "input").
            device: Device to use ("cuda", "cpu", "mps", or None for auto).
            **kwargs: Additional configuration passed to parent.
        """
        super().__init__(**kwargs)

        if target not in ("output", "input"):
            raise ValueError(f"Invalid target: {target}. Must be 'output' or 'input'")

        self.model_name = model
        self.entities = [e.upper() for e in entities] if entities else None
        self.fail_on_detection = fail_on_detection
        self.target = target
        self.device = device
        self._pipeline: Any = None  # Lazy loaded

        # Store config for serialization
        self.config = {
            "model": model,
            "entities": entities,
            "fail_on_detection": fail_on_detection,
            "target": target,
            "device": device,
            **kwargs,
        }

    def _get_device_id(self) -> int | str:
        """Determine the device ID for transformers pipeline.

        Returns:
            Device ID: -1 for CPU, 0+ for GPU, or "mps" for Apple Silicon.
        """
        if self.device is not None:
            if self.device == "cpu":
                return -1
            elif self.device == "cuda":
                return 0
            elif self.device == "mps":
                return "mps"
            else:
                return -1

        if not TORCH_AVAILABLE:
            return -1

        if torch.cuda.is_available():
            return 0
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
        return -1

    def _load_pipeline(self) -> Any:
        """Lazy load NER pipeline.

        Returns:
            Loaded transformers pipeline.

        Raises:
            ImportError: If transformers is not installed.
        """
        if self._pipeline is None:
            if not TRANSFORMERS_AVAILABLE:
                raise ImportError(
                    "transformers is required for PIIDetection. "
                    "Install it with: pip install mankinds-eval[ml]"
                )

            from mankinds_eval.config import get_model

            model_name = get_model("ner", self.model_name)
            device = self._get_device_id()
            self._pipeline = transformers_pipeline(
                "ner",
                model=model_name,
                device=device,
                aggregation_strategy="simple",
            )

        return self._pipeline

    def _normalize_entity_type(self, entity_type: str) -> str:
        """Normalize entity type to uppercase standard format.

        Args:
            entity_type: Raw entity type from model.

        Returns:
            Normalized entity type.
        """
        # Remove common prefixes (B-, I-, etc. from BIO tagging)
        normalized = entity_type.upper()
        for prefix in ("B-", "I-", "B_", "I_", "S-", "S_", "E-", "E_"):
            if normalized.startswith(prefix):
                normalized = normalized[len(prefix) :]
                break

        return normalized

    def _filter_entities(self, entities: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Filter entities by type if entity filter is specified.

        Args:
            entities: List of detected entities from the model.

        Returns:
            Filtered list of entities.
        """
        if self.entities is None:
            return entities

        filtered = []
        for entity in entities:
            entity_type = self._normalize_entity_type(entity.get("entity_group", entity.get("entity", "")))
            if entity_type in self.entities:
                filtered.append(entity)

        return filtered

    async def evaluate(self, sample: Sample) -> MethodResult:
        """Detect PII entities in the target text.

        Args:
            sample: The sample to evaluate.

        Returns:
            MethodResult with PII detection results.
        """
        pipeline = self._load_pipeline()

        # Get target text
        text = sample.output if self.target == "output" else sample.input

        # Run NER pipeline
        raw_entities = pipeline(text)

        # Filter by entity types if specified
        filtered_entities = self._filter_entities(raw_entities)

        # Build entity list for metadata
        entity_list = []
        entity_types_found: set[str] = set()
        for entity in filtered_entities:
            entity_type = self._normalize_entity_type(entity.get("entity_group", entity.get("entity", "")))
            entity_types_found.add(entity_type)
            start = entity.get("start", 0)
            end = entity.get("end", 0)
            entity_text = text[start:end] if start < end else entity.get("word", "")
            entity_list.append(
                {
                    "type": entity_type,
                    "text": entity_text,
                    "score": entity.get("score", 0.0),
                    "start": start,
                    "end": end,
                }
            )

        # Calculate score (number of PII entities found)
        num_entities = len(entity_list)
        score = float(num_entities)

        # Determine passed status
        if self.fail_on_detection:
            passed = num_entities == 0
        else:
            passed = True

        # Build reason
        if num_entities == 0:
            reason = "No PII detected"
        else:
            types_str = ", ".join(sorted(entity_types_found))
            reason = f"Found {num_entities} PII entit{'y' if num_entities == 1 else 'ies'}: {types_str}"

        from mankinds_eval.config import get_model

        actual_model = get_model("ner", self.model_name)

        return MethodResult(
            method_name=self.name,
            score=score,
            passed=passed,
            reason=reason,
            metadata={
                "entities": entity_list,
                "entity_types": sorted(entity_types_found),
                "entity_count": num_entities,
                "model": actual_model,
                "target": self.target,
                "filter_entities": self.entities,
            },
        )
