"""MCP server exposing AmbientMeta PII sanitization tools.

Tools:
    sanitize_text  — Detect and replace PII with placeholders. Returns sanitized
                     text + session_id for later rehydration.
    rehydrate_text — Restore original values from placeholders using a session_id.
    check_pii      — Detect PII in text and return what was found (uses redact mode
                     internally so no session is created).
"""

from __future__ import annotations

import os
import sys

from mcp.server.fastmcp import FastMCP

from ambientmeta import AmbientMeta, AmbientMetaError

_API_KEY = os.environ.get("AMBIENTMETA_API_KEY", "")
_BASE_URL = os.environ.get("AMBIENTMETA_BASE_URL", "https://api.ambientmeta.com")

server = FastMCP(
    "AmbientMeta Privacy Gateway",
    instructions="Sanitize PII before sending text to AI models. Rehydrate after.",
)

# Reuse a single client instance for connection pooling across tool calls.
_client: AmbientMeta | None = None


def _get_client() -> AmbientMeta:
    global _client
    if _client is not None:
        return _client
    if not _API_KEY:
        raise ValueError(
            "AMBIENTMETA_API_KEY environment variable is required. "
            "Get your key at https://dashboard.ambientmeta.com"
        )
    _client = AmbientMeta(api_key=_API_KEY, base_url=_BASE_URL)
    return _client


@server.tool()
def sanitize_text(text: str) -> dict:
    """Detect and replace PII (names, emails, phones, SSNs, etc.) with safe placeholders.

    Use this BEFORE sending user text to any AI model or external service.
    Returns sanitized text and a session_id needed to restore originals later.

    Args:
        text: The text to sanitize. Can be any length.

    Returns:
        dict with:
            sanitized: Text with PII replaced by placeholders like [PERSON_1], [EMAIL_1]
            session_id: ID to pass to rehydrate_text to restore originals (valid 24h)
            entities_found: Number of PII entities detected
            entities: List of detected entities with type and confidence
    """
    client = _get_client()
    try:
        result = client.sanitize(text)
        return {
            "sanitized": result.sanitized,
            "session_id": result.session_id,
            "entities_found": result.entities_found,
            "entities": [
                {"placeholder": e.placeholder, "type": e.type, "confidence": e.confidence}
                for e in result.entities
            ],
        }
    except AmbientMetaError as e:
        return {"error": str(e)}
    except Exception as e:
        return {"error": f"Unexpected error: {type(e).__name__}: {e}"}


@server.tool()
def rehydrate_text(text: str, session_id: str) -> dict:
    """Restore original PII values from placeholders.

    Use this AFTER receiving a response from an AI model to put back the real
    names, emails, etc. that were replaced during sanitization.

    Args:
        text: Text containing placeholders like [PERSON_1], [EMAIL_1].
        session_id: The session_id returned by sanitize_text.

    Returns:
        dict with:
            text: Original text with real values restored
            entities_restored: Number of placeholders replaced
    """
    client = _get_client()
    try:
        result = client.rehydrate(text, session_id)
        return {
            "text": result.text,
            "entities_restored": result.entities_restored,
        }
    except AmbientMetaError as e:
        return {"error": str(e)}
    except Exception as e:
        return {"error": f"Unexpected error: {type(e).__name__}: {e}"}


@server.tool()
def check_pii(text: str) -> dict:
    """Detect PII in text and report what was found.

    Runs the full PII detection pipeline and returns the entity list.
    No session is created — use sanitize_text if you need to replace and
    later restore the PII. This costs one API call (same as sanitize_text).

    Args:
        text: The text to scan.

    Returns:
        dict with:
            has_pii: Whether any PII was detected
            entities_found: Number of PII entities detected
            entities: List of detected entities with type and confidence
    """
    client = _get_client()
    try:
        # Use redact mode if the SDK supports it, otherwise fall back to sanitize
        try:
            result = client.sanitize(text, mode="redact")
        except TypeError:
            result = client.sanitize(text)
        return {
            "has_pii": result.entities_found > 0,
            "entities_found": result.entities_found,
            "entities": [
                {"type": e.type, "confidence": e.confidence}
                for e in result.entities
            ],
        }
    except AmbientMetaError as e:
        return {"error": str(e)}
    except Exception as e:
        return {"error": f"Unexpected error: {type(e).__name__}: {e}"}


def main() -> None:
    """Entry point for the MCP server."""
    if not _API_KEY:
        print(
            "Error: AMBIENTMETA_API_KEY environment variable is required.\n"
            "Get your key at https://dashboard.ambientmeta.com\n\n"
            "Usage:\n"
            "  AMBIENTMETA_API_KEY=am_... ambientmeta-mcp",
            file=sys.stderr,
        )
        sys.exit(1)
    try:
        server.run()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()
