"""Passthrough envelope: destination in the clear, body encrypted.

After send, the sender cannot open the ciphertext; they keep lookup_id only.
Production alg is npe (HPKE). Lab algs are lab-xor and plain-lab.
"""

from __future__ import annotations  # annotations as strings

import hashlib  # lookup_id and lab key
import json  # body encode
import os  # urandom nonce
from dataclasses import dataclass  # envelope fields
from typing import Any, Dict, Mapping, Optional  # maps


def lookup_id(sender: str, nonce: bytes) -> str:
    """Opaque id: HMAC-SHA256(sender, nonce) hex. Not reversible at the broker."""
    return hashlib.sha256(nonce + sender.encode("utf-8")).hexdigest()  # bind sender+nonce


def _xor(key: bytes, data: bytes) -> bytes:
    """Lab-only repeating XOR (same construction as pfc-lab-xor)."""
    out = bytearray(len(data))  # output buffer
    for i, b in enumerate(data):  # each plaintext byte
        out[i] = b ^ key[i % len(key)]  # xor with cycling key
    return bytes(out)  # immutable


@dataclass
class PassthroughEnvelope:
    """On-wire object. Header-like fields stay JSON; body is ct."""

    to: str  # mailbox dest, in the clear for routing
    from_lookup_id: str  # sender cannot be recovered by broker
    alg: str  # npe | lab-xor | plain-lab
    ct: str  # hex ciphertext or empty if plain-lab
    nonce: str  # hex nonce used in lookup_id
    error_token: str  # public-key token placeholder for return path

    def header(self) -> Dict[str, str]:
        """Routing/logging header (no payload)."""
        return {  # allowed on the untrusted broker
            "to": self.to,
            "from_lookup_id": self.from_lookup_id,
            "alg": self.alg,
            "nonce": self.nonce,
            "error_token": self.error_token,
        }

    def wire(self) -> Dict[str, Any]:
        """Full JSON for NATS."""
        w = self.header()  # start with header
        w["ct"] = self.ct  # ciphertext only
        w["v"] = 1  # version
        w["mode"] = "passthrough"  # dest in clear
        return w  # ready to publish


def seal(
    *,
    to: str,
    sender: str,
    body: Mapping[str, Any],
    mode: str,
    lab_key: bytes = b"",
) -> PassthroughEnvelope:
    """Encrypt body for recipient; return passthrough envelope."""
    nonce = os.urandom(16)  # fresh nonce
    lid = lookup_id(sender, nonce)  # opaque sender id
    raw = json.dumps(body, sort_keys=True, separators=(",", ":")).encode()  # body bytes
    if mode == "plain-lab":  # tests only
        ct = raw.hex()  # not secret
    elif mode == "lab-xor":  # lab PSK
        if not lab_key:  # require key
            raise ValueError("lab_key required for lab-xor")
        key = hashlib.sha256(lab_key).digest()  # 32-byte key
        ct = _xor(key, raw).hex()  # hex ct
    elif mode == "npe":  # production — sidecar not in this module
        raise NotImplementedError("npe: attach NPE sidecar; do not HPKE here")
    else:  # unknown
        raise ValueError("unknown crypto.mode " + mode)
    token = hashlib.sha256(nonce + b"error-token").hexdigest()[:32]  # return-path token stub
    return PassthroughEnvelope(  # dest in clear
        to=to,
        from_lookup_id=lid,
        alg=mode,
        ct=ct,
        nonce=nonce.hex(),
        error_token=token,
    )


def open_lab(env: PassthroughEnvelope, lab_key: bytes = b"") -> Dict[str, Any]:
    """Decrypt lab envelopes only. Production uses NPE sidecar."""
    raw = bytes.fromhex(env.ct)  # ct bytes
    if env.alg == "plain-lab":  # tests
        return json.loads(raw.decode())  # json body
    if env.alg == "lab-xor":  # lab
        key = hashlib.sha256(lab_key).digest()  # same kdf
        pt = _xor(key, raw)  # decrypt
        return json.loads(pt.decode())  # json
    raise ValueError("open_lab cannot open alg " + env.alg)  # npe not here