#!/usr/bin/env python3
"""Protect context, call an OpenAI-compatible model, and restore its reply locally."""

import json
import os
import urllib.request


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=60) as response:
        return json.load(response)


shinrai_url = os.getenv("SHINRAI_BASE_URL", "https://api.shinrai.innovius.io").rstrip("/")
protected = post(
    shinrai_url + "/v1/redact",
    os.environ["SHINRAI_API_KEY"],
    {"text": "Summarize Emma Weber's account. Email emma@example.com.", "mode": "replace"},
)

# Only protected text crosses the model boundary. Keep the mapping local.
completion = post(
    os.environ["MODEL_BASE_URL"].rstrip("/") + "/chat/completions",
    os.environ["MODEL_API_KEY"],
    {"model": os.environ["MODEL"], "messages": [{"role": "user", "content": protected["text"]}]},
)
reply = completion["choices"][0]["message"]["content"]
for original, replacement in sorted(protected.get("mapping", {}).items(), key=lambda item: len(item[1]), reverse=True):
    reply = reply.replace(replacement, original)
print(reply)