"""NATS content wrapping: HPKE body, routing/error in the clear.

Clear (broker may see): subject, ``to``, ``from_lookup_id``, ``alg``,
error_code, dest_class, lookup_id.
Never clear: inner content JSON / PHI / object bytes.
"""

from __future__ import annotations  # annotations

import json  # inner body
import os  # urandom
from typing import Any, Dict, Mapping, Optional  # types

from .envelope import lookup_id  # opaque sender handle
from .hpke import open_ct, seal  # HPKE-Base
from .pubkey_dir import PubKeyDir  # directory


def wrap_content(
    *,
    to_handle: str,
    sender: str,
    body: Mapping[str, Any],
    directory: PubKeyDir,
) -> Dict[str, Any]:
    """Encrypt body to the recipient's directory public key."""
    rec = directory.get(to_handle)
    if not rec:  # unknown endpoint
        raise KeyError("directory has no public key for " + to_handle)
    pk = bytes.fromhex(rec["enc_pk"])
    inner = json.dumps(body, sort_keys=True, separators=(",", ":")).encode()
    aad = ("to=" + to_handle).encode()  # bind ciphertext to routing dest
    enc, ct = seal(pk, inner, aad=aad)
    nonce = os.urandom(16)
    return {
        "v": 1,
        "alg": "npe",
        "mode": "passthrough",
        "to": rec.get("inbox") or to_handle,  # routing in the clear
        "to_handle": to_handle,
        "from_lookup_id": lookup_id(sender, nonce),
        "nonce": nonce.hex(),
        "enc": enc.hex(),  # HPKE encapsulated key (not content)
        "ct": ct.hex(),  # content ciphertext
    }


def unwrap_content(env: Mapping[str, Any], *, handle: str, directory: PubKeyDir) -> Dict[str, Any]:
    """Open content as the named endpoint. Broker cannot do this."""
    sk = directory.load_sk(handle)
    to_handle = str(env.get("to_handle") or handle)
    aad = ("to=" + to_handle).encode()
    pt = open_ct(sk, bytes.fromhex(env["enc"]), bytes.fromhex(env["ct"]), aad=aad)
    return json.loads(pt.decode())


def routing_view(env: Mapping[str, Any]) -> Dict[str, Any]:
    """What an untrusted router is allowed to look at."""
    return {
        k: env[k]
        for k in ("v", "alg", "mode", "to", "to_handle", "from_lookup_id", "nonce")
        if k in env
    }