"""Blindfold PII validator for Guardrails AI.

Detects PII in LLM outputs using the Blindfold API and optionally fixes it
via tokenization, redaction, masking, hashing, synthesis, or encryption.
"""

import os
from typing import Any, Dict, List, Optional

from guardrails.validators import (
    FailResult,
    PassResult,
    Validator,
    register_validator,
)
from guardrails.classes.validation.validation_result import ErrorSpan

VALID_PII_METHODS = ("tokenize", "redact", "mask", "hash", "synthesize", "encrypt")


@register_validator(name="blindfold/pii_protection", data_type="string")
class BlindfoldPII(Validator):
    """Detects and protects PII using the Blindfold API.

    **Key Arguments**

    | Argument         | Type        | Default     | Description |
    |------------------|-------------|-------------|-------------|
    | `policy`         | `str`       | `"basic"`   | Detection policy |
    | `pii_method`     | `str`       | `"tokenize"`| How to fix PII |
    | `region`         | `str`       | `None`      | `"eu"` or `"us"` |
    | `entities`       | `list`      | `None`      | Entity types to detect |
    | `score_threshold`| `float`     | `None`      | Confidence threshold |
    | `api_key`        | `str`       | `None`      | Falls back to env var |
    """

    def __init__(
        self,
        policy: str = "basic",
        pii_method: str = "tokenize",
        region: Optional[str] = None,
        entities: Optional[List[str]] = None,
        score_threshold: Optional[float] = None,
        api_key: Optional[str] = None,
        on_fail: Optional[str] = None,
        **kwargs: Any,
    ):
        super().__init__(on_fail=on_fail, **kwargs)

        if pii_method not in VALID_PII_METHODS:
            raise ValueError(
                f"pii_method must be one of {VALID_PII_METHODS}, got '{pii_method}'"
            )

        self._policy = policy
        self._pii_method = pii_method
        self._region = region
        self._entities = entities
        self._score_threshold = score_threshold
        self._api_key = api_key or os.environ.get("BLINDFOLD_API_KEY")
        self._client = None

    def _get_client(self):
        """Lazy-initialize the Blindfold client."""
        if self._client is None:
            from blindfold import Blindfold

            if not self._api_key:
                raise ValueError(
                    "Blindfold API key is required. Pass api_key= or set "
                    "the BLINDFOLD_API_KEY environment variable."
                )
            self._client = Blindfold(
                api_key=self._api_key,
                region=self._region,
            )
        return self._client

    def _build_kwargs(self) -> Dict[str, Any]:
        """Build common kwargs for Blindfold API calls."""
        kwargs: Dict[str, Any] = {"policy": self._policy}
        if self._entities is not None:
            kwargs["entities"] = self._entities
        if self._score_threshold is not None:
            kwargs["score_threshold"] = self._score_threshold
        return kwargs

    def _validate(self, value: Any, metadata: Dict[str, Any]) -> PassResult | FailResult:
        """Validate that the value does not contain PII.

        If PII is found, returns a FailResult with the protected text as
        fix_value (when on_fail="fix").
        """
        text = str(value)
        client = self._get_client()
        kwargs = self._build_kwargs()

        # Detect PII
        detect_result = client.detect(text, **kwargs)

        if not detect_result.detected_entities:
            return PassResult()

        # Build error spans from detected entities
        error_spans = [
            ErrorSpan(
                start=entity.start,
                end=entity.end,
                reason=f"{entity.type} detected",
            )
            for entity in detect_result.detected_entities
        ]

        # Apply the chosen PII method to get the fix value
        fix_value = self._apply_pii_method(client, text, kwargs)

        entity_types = sorted({e.type for e in detect_result.detected_entities})
        count = len(detect_result.detected_entities)
        error_message = (
            f"Found {count} PII entit{'y' if count == 1 else 'ies'} "
            f"({', '.join(entity_types)})"
        )

        return FailResult(
            error_message=error_message,
            error_spans=error_spans,
            fix_value=fix_value,
            metadata=metadata,
        )

    def _apply_pii_method(
        self, client: Any, text: str, kwargs: Dict[str, Any]
    ) -> str:
        """Apply the configured PII protection method and return the fixed text."""
        method = self._pii_method

        if method == "tokenize":
            result = client.tokenize(text, **kwargs)
            return result.text
        elif method == "redact":
            result = client.redact(text, **kwargs)
            return result.text
        elif method == "mask":
            result = client.mask(text, **kwargs)
            return result.text
        elif method == "hash":
            result = client.hash(text, **kwargs)
            return result.text
        elif method == "synthesize":
            result = client.synthesize(text, **kwargs)
            return result.text
        elif method == "encrypt":
            result = client.encrypt(text, **kwargs)
            return result.text
        else:
            raise ValueError(f"Unknown pii_method: {method}")
