"""
StripLLM core class — public API surface.
"""

from typing import Optional, Union

from .injection import detect_injection, InjectionResult
from .pii import redact as _redact, rehydrate as _rehydrate, RedactionResult
from .output import enforce as _enforce, ValidationResult
from .audit import audit as _audit, AuditReport


class InjectionDetectedError(Exception):
    """Raised by clean() when injection signals exceed the threshold."""
    def __init__(self, message: str, risk_score: float, matched_patterns: list):
        super().__init__(message)
        self.risk_score = risk_score
        self.matched_patterns = matched_patterns


class OutputValidationError(Exception):
    """Raised by enforce() when validation errors are found."""
    def __init__(self, message: str, errors: list, warnings: list):
        super().__init__(message)
        self.errors = errors
        self.warnings = warnings


class StripLLM:
    """
    StripLLM — LLM input/output sanitization SDK.

    Usage::

        from stripllm import StripLLM
        strip = StripLLM()

        safe = strip.clean(user_message)
        redacted, mapping = strip.redact("Email me at john@acme.com")
        validated = strip.enforce(llm_response, schema="json")
        report = strip.audit(conversation)
    """

    def __init__(self, threshold: float = 0.3):
        """
        Args:
            threshold: Injection risk score (0.0–1.0) above which clean() raises.
                       Default 0.3. Lower = more sensitive.
        """
        self.threshold = threshold

    # ------------------------------------------------------------------
    # clean()
    # ------------------------------------------------------------------

    def clean(self, text: str) -> str:
        """
        Sanitize user input against prompt injection and jailbreak attacks.

        Raises InjectionDetectedError if the risk score meets or exceeds threshold.
        Otherwise returns the original text (no destructive modification — callers
        retain the original phrasing; detection-and-block is the intended pattern).

        Args:
            text: Raw user message / prompt.

        Returns:
            The original text if no injection detected.

        Raises:
            InjectionDetectedError: If injection signals are found above threshold.
        """
        result: InjectionResult = detect_injection(text)
        if result.risk_score >= self.threshold:
            raise InjectionDetectedError(
                f"Potential prompt injection detected (risk_score={result.risk_score:.3f})",
                risk_score=result.risk_score,
                matched_patterns=result.matched_patterns,
            )
        return text

    def scan(self, text: str) -> InjectionResult:
        """
        Non-raising version of clean() — returns the full InjectionResult for inspection.
        """
        return detect_injection(text)

    # ------------------------------------------------------------------
    # redact()
    # ------------------------------------------------------------------

    def redact(self, text: str, entities: Optional[list] = None) -> tuple:
        """
        Redact PII from text, replacing each occurrence with a typed placeholder.

        Args:
            text: Text that may contain PII.
            entities: Optional list of entity types to target, e.g. ["EMAIL", "PHONE"].
                      Supported types: EMAIL, PHONE, SSN, CREDIT_CARD, IP_ADDRESS,
                      DRIVERS_LICENSE, DATE_OF_BIRTH, PASSPORT, URL.
                      If None, all types are checked.

        Returns:
            Tuple of (redacted_text: str, mapping: dict).
            The mapping can be passed to strip.rehydrate() to restore originals.
        """
        result: RedactionResult = _redact(text, entities=entities)
        return result.redacted_text, result.mapping

    def rehydrate(self, text: str, mapping: dict) -> str:
        """
        Restore original PII values in LLM output using a mapping from redact().

        Args:
            text: LLM output that may contain [EMAIL_1] style placeholders.
            mapping: The mapping dict returned by strip.redact().

        Returns:
            Text with placeholders replaced by original values.
        """
        return _rehydrate(text, mapping)

    # ------------------------------------------------------------------
    # enforce()
    # ------------------------------------------------------------------

    def enforce(
        self,
        text: str,
        schema: Optional[Union[str, dict]] = None,
        raise_on_error: bool = True,
    ) -> str:
        """
        Validate and clean LLM output.

        Args:
            text: Raw LLM response string.
            schema: "json" to require valid JSON, or a dict mapping key→expected_type
                    for structured validation. None = safety checks only.
            raise_on_error: If True (default), raise OutputValidationError on hard errors.
                            If False, return the best-effort output and surface errors as warnings.

        Returns:
            Cleaned output string (JSON is re-serialized for canonical formatting).

        Raises:
            OutputValidationError: If schema validation fails and raise_on_error=True.
        """
        result: ValidationResult = _enforce(text, schema=schema)
        if result.errors and raise_on_error:
            raise OutputValidationError(
                f"Output validation failed: {'; '.join(result.errors)}",
                errors=result.errors,
                warnings=result.warnings,
            )
        return result.output

    def validate(self, text: str, schema: Optional[Union[str, dict]] = None) -> ValidationResult:
        """
        Non-raising version of enforce() — returns the full ValidationResult.
        """
        return _enforce(text, schema=schema)

    # ------------------------------------------------------------------
    # audit()
    # ------------------------------------------------------------------

    def audit(self, conversation: list) -> AuditReport:
        """
        Perform a full security audit of a multi-turn conversation.

        Args:
            conversation: List of turn dicts:
                [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ...]

        Returns:
            AuditReport with .risk_score (0.0–1.0), .findings, and .recommendations.
        """
        return _audit(conversation)
