"""Public package surface for Universal PII Firewall."""

import base64
from typing import Optional

from .config import UPFConfig
from .llm import SecureLLMResult, secure_llm_call
from .pipeline.text_pipeline import HighRiskBlockedError, TextPipeline, TextSanitizeResult
from .redaction.pseudonym import PseudonymSession
from .redaction.text import RedactionMode


def _image_pipeline(config):
    from .pipeline.image_pipeline import ImagePipeline
    return ImagePipeline(config)


try:
    from .pipeline.image_pipeline import ImageSanitizeResult
except ImportError:
    ImageSanitizeResult = None  # type: ignore[assignment,misc]


def sanitize_text(text: str, config: Optional[UPFConfig] = None) -> str:
    """Return text with detected PII redacted according to *config*."""
    return TextPipeline(config).sanitize(text)


def sanitize_text_with_details(text: str, config: Optional[UPFConfig] = None) -> TextSanitizeResult:
    """Return sanitized text plus entities, risk score, and risk level."""
    return TextPipeline(config).sanitize_with_details(text)


def sanitize_image(ocr_text: str, config: Optional[UPFConfig] = None) -> str:
    """Sanitize OCR text using the same text pipeline as `sanitize_text`."""
    return _image_pipeline(config).sanitize_ocr_text(ocr_text)


def sanitize_image_bytes(
    image_bytes: bytes, ocr_text: Optional[str] = None, config: Optional[UPFConfig] = None
):
    """Sanitize image bytes and return text + redacted image payload."""
    return _image_pipeline(config).sanitize_image_bytes(image_bytes, ocr_text=ocr_text)


def sanitize_image_base64(
    image_base64: str, ocr_text: Optional[str] = None, config: Optional[UPFConfig] = None
):
    """Decode base64 image payload, then sanitize it as bytes."""
    return sanitize_image_bytes(base64.b64decode(image_base64, validate=True), ocr_text=ocr_text, config=config)


__all__ = [
    "sanitize_text",
    "sanitize_text_with_details",
    "sanitize_image",
    "sanitize_image_bytes",
    "sanitize_image_base64",
    "UPFConfig",
    "RedactionMode",
    "PseudonymSession",
    "secure_llm_call",
    "SecureLLMResult",
    "HighRiskBlockedError",
    "ImageSanitizeResult",
    "TextSanitizeResult",
]
