#!/usr/bin/env python3
"""Protect context with the ShinrAI PII API v2, call an OpenAI-compatible model, and restore its reply locally.

Environment:
  SHINRAI_API_KEY   your ShinrAI key
  SHINRAI_API_URL   optional, default https://api.getshinrai.com (sandbox: https://api-sbx.getshinrai.com)
  MODEL_ENDPOINT    OpenAI-compatible base URL, for example https://model.example/v1
  MODEL_NAME        model ID at that endpoint
  MODEL_API_KEY     key for that endpoint

Standard library only (Python 3.8+).
"""

import json
import os
import urllib.request


def env(*names, default=None):
    for name in names:
        if os.environ.get(name):
            return os.environ[name]
    if default is None:
        raise SystemExit("Set " + " or ".join(names) + ".")
    return default


def post(url, token, payload):
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=120) as response:
        return json.load(response)


context = "Summarize Emma Weber's account. Email emma@example.com."

# 1. Protect: realistic surrogates, and the mapping of originals to surrogates for this request.
shinrai_url = env("SHINRAI_API_URL", "SHINRAI_BASE_URL", default="https://api.getshinrai.com").rstrip("/")
protected = post(
    shinrai_url + "/v2/protect",
    env("SHINRAI_API_KEY"),
    {"text": context, "policy": {"preset": "pseudonymize"}, "output": {"include": ["entities", "mapping"]}},
)
result = protected["results"][0]
if result.get("status") != "ok":
    raise SystemExit("ShinrAI did not protect the text; nothing was sent to the model.")
protected_text = result["output"]
mapping = [entry for entry in protected.get("mapping", {}).get("delta", []) if entry.get("reversible")]

# 2. Only protected text crosses the model boundary. The mapping stays in this process.
completion = post(
    env("MODEL_ENDPOINT", "MODEL_BASE_URL").rstrip("/") + "/chat/completions",
    env("MODEL_API_KEY"),
    {"model": env("MODEL_NAME", "MODEL"), "messages": [{"role": "user", "content": protected_text}]},
)
reply = completion["choices"][0]["message"]["content"]

# 3. Restore locally: put the originals back, longest surrogate first.
# (POST /v2/restore with {"mapping": {"known": [...]}, "inputs": [{"id": "1", "text": reply}]} does the same on the server.)
for entry in sorted(mapping, key=lambda item: len(item["replacement"]), reverse=True):
    reply = reply.replace(entry["replacement"], entry["original"])

print("Protected model input:", protected_text)
print("Restored reply:", reply)