"""
/restore endpoint (T2)  --  egress re-identification for the HNCB demo (step 8).

Fusion SaaS POSTs the agent's token-only answer back here with the session_id.
We:
  1. read THIS session's tokens from the on-prem vault (scoped by session_id, so
     one session can never restore another's identities)
  2. swap any inline tokens in the text back to the original PII
  3. re-attach identity at the envelope level too: prepend the real customer name,
     so the advisor sees who the answer is about even if no token appears inline
  4. return {final} -- the only step where a real name comes back into view

The reversible map lives only in the vault (on-prem zone); the cloud never saw it.
Runs as a Lambda behind the same public API Gateway as /tokenize, guarded by the
same shared secret. Stdlib + boto3 only, per repo conventions.
"""
import base64
import hmac
import json
import os

import boto3
from boto3.dynamodb.conditions import Attr

ddb = boto3.resource("dynamodb")
VAULT = ddb.Table(os.environ["VAULT_TABLE"])

# Shared secret Fusion presents (x-api-key / bearer). Unset -> allow (dev only).
API_KEY = os.environ.get("TOKENIZE_API_KEY")


def _authorized(headers):
    if not API_KEY:
        return True
    presented = headers.get("x-api-key") or ""
    if not presented:
        auth = headers.get("authorization") or ""
        if auth.lower().startswith("bearer "):
            presented = auth[7:]
    return bool(presented) and hmac.compare_digest(presented, API_KEY)


def _parse(event):
    if "body" in event and not isinstance(event.get("body"), dict):
        raw = event.get("body") or "{}"
        if event.get("isBase64Encoded"):
            raw = base64.b64decode(raw).decode("utf-8")
        payload = json.loads(raw) if raw else {}
        headers = {k.lower(): v for k, v in (event.get("headers") or {}).items()}
        return payload, headers
    return event, {}


def _reply(status, obj):
    return {
        "statusCode": status,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(obj, ensure_ascii=False),
    }


def _session_tokens(session_id):
    # Demo-grade: scan the vault filtered by session_id (small table, no GSI).
    # A production vault would query a session_id index instead.
    items, kwargs = [], {"FilterExpression": Attr("session_id").eq(session_id)}
    while True:
        res = VAULT.scan(**kwargs)
        items.extend(res.get("Items", []))
        lek = res.get("LastEvaluatedKey")
        if not lek:
            return items
        kwargs["ExclusiveStartKey"] = lek


def lambda_handler(event, _context):
    payload, headers = _parse(event)

    if not _authorized(headers):
        return _reply(401, {"error": "unauthorized"})

    session_id = (payload.get("session_id") or "").strip()
    text = payload.get("text")
    if not session_id or text is None:
        return _reply(400, {"error": "session_id and text are required"})

    tokens = _session_tokens(session_id)

    # Inline: swap every token for this session back to its original value.
    # Longest token first so no token is a prefix of another mid-replace.
    restored = text
    for it in sorted(tokens, key=lambda t: len(t["token"]), reverse=True):
        restored = restored.replace(it["token"], it["value"])

    # Envelope: prepend the customer's real name (the PERSON mapping), if any.
    name = next((it["value"] for it in tokens if it.get("type") == "PERSON"), None)
    final = f"（客戶：{name}）\n{restored}" if name else restored

    return _reply(200, {"final": final, "restored": len(tokens)})