#!/usr/bin/env python3
"""SafeDoc MCP server — anonymisez vos documents AVANT que l'IA ne les lise.

Serveur MCP (Model Context Protocol) stdio sans dépendance externe :
fonctionne avec Claude Desktop, Claude Code et tout client MCP.

Configuration Claude Desktop (claude_desktop_config.json) :
{
  "mcpServers": {
    "safedoc": {
      "command": "python3",
      "args": ["/chemin/vers/safedoc_mcp.py"],
      "env": { "SAFEDOC_API_KEY": "sd_votre_cle" }
    }
  }
}

Claude Code : claude mcp add safedoc -e SAFEDOC_API_KEY=sd_xxx -- python3 safedoc_mcp.py
La clé API se génère depuis votre profil sur https://safe-doc.ai/app/.
"""
import json
import os
import sys
import urllib.parse
import urllib.request

BASE_URL = os.environ.get("SAFEDOC_URL", "https://safe-doc.ai").rstrip("/")
API_KEY = os.environ.get("SAFEDOC_API_KEY", "")

SERVER_INFO = {"name": "safedoc", "version": "0.1.0"}
PROTOCOL_VERSION = "2024-11-05"

TOOLS = [
    {
        "name": "anonymize_text",
        "description": (
            "Anonymise un texte (noms, adresses, identifiants, téléphones, IBAN…) "
            "sur les serveurs européens SafeDoc avant tout traitement par une IA. "
            "Retourne le texte anonymisé et la table de correspondance jeton→valeur "
            "(utilisable pour ré-identifier localement après coup)."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "text": {"type": "string", "description": "Texte à anonymiser"},
                "level": {
                    "type": "string", "enum": ["N1", "N2", "N3"], "default": "N1",
                    "description": "N1 jetons typés, N2 généralisation, N3 maximal",
                },
                "mode": {
                    "type": "string", "enum": ["anonymize", "pseudonymize", "fake_data"],
                    "default": "anonymize",
                },
            },
            "required": ["text"],
        },
    },
    {
        "name": "restore_text",
        "description": (
            "Ré-identifie LOCALEMENT un texte anonymisé à partir de la table de "
            "correspondance retournée par anonymize_text (aucun appel réseau : la "
            "restauration se fait dans ce processus)."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "text": {"type": "string", "description": "Texte contenant des jetons [PERSONNE_1]…"},
                "mapping": {
                    "type": "object",
                    "description": "Table {valeur_originale: jeton} renvoyée par anonymize_text",
                },
            },
            "required": ["text", "mapping"],
        },
    },
]


def call_safedoc(text: str, level: str, mode: str) -> dict:
    if not API_KEY:
        raise RuntimeError("SAFEDOC_API_KEY manquante (générez votre clé sur safe-doc.ai/app → Profil)")
    body = urllib.parse.urlencode(
        {"text": text, "level": level, "mode": mode, "stateless": "true"}
    ).encode()
    req = urllib.request.Request(
        BASE_URL + "/api_v2/anonymize-text",
        data=body,
        headers={"X-API-Key": API_KEY, "Content-Type": "application/x-www-form-urlencoded"},
    )
    with urllib.request.urlopen(req, timeout=180) as r:
        return json.loads(r.read().decode())


def tool_call(name: str, args: dict) -> str:
    if name == "anonymize_text":
        res = call_safedoc(args["text"], args.get("level", "N1"), args.get("mode", "anonymize"))
        return json.dumps(
            {"anonymized_text": res.get("anonymized_text", ""), "mapping": res.get("mapping", {})},
            ensure_ascii=False,
        )
    if name == "restore_text":
        text = args["text"]
        for original, token in (args.get("mapping") or {}).items():
            text = text.replace(str(token), str(original))
        return text
    raise RuntimeError(f"outil inconnu: {name}")


def handle(req: dict):
    method = req.get("method", "")
    if method == "initialize":
        return {
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": {"tools": {}},
            "serverInfo": SERVER_INFO,
        }
    if method == "tools/list":
        return {"tools": TOOLS}
    if method == "tools/call":
        params = req.get("params", {})
        try:
            out = tool_call(params.get("name", ""), params.get("arguments", {}) or {})
            return {"content": [{"type": "text", "text": out}], "isError": False}
        except Exception as e:
            return {"content": [{"type": "text", "text": f"Erreur SafeDoc: {e}"}], "isError": True}
    if method in ("notifications/initialized", "notifications/cancelled"):
        return None  # notification, pas de réponse
    if method == "ping":
        return {}
    raise RuntimeError(f"méthode non supportée: {method}")


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError:
            continue
        rid = req.get("id")
        try:
            result = handle(req)
        except Exception as e:
            if rid is not None:
                resp = {"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": str(e)}}
                sys.stdout.write(json.dumps(resp) + "\n")
                sys.stdout.flush()
            continue
        if rid is None or result is None:
            continue  # notification
        resp = {"jsonrpc": "2.0", "id": rid, "result": result}
        sys.stdout.write(json.dumps(resp, ensure_ascii=False) + "\n")
        sys.stdout.flush()


if __name__ == "__main__":
    main()