"""The sealed envelope — the unit of everything, and the trust boundary itself.

Wire format v2 (docs/V4-FEDERATION.md §2.3), one JSON object::

    {"v": 2, "id": "<uuid4 hex>", "from": "ax1…", "to": "ax1…",
     "ts": "2026-08-31T04:12:09Z", "kind": "intents", "sealed": "<b64>"}

``sealed`` is ``nacl.public.Box(sender_priv, recipient_pub).encrypt(plaintext)``
(the nonce is prepended by ``encrypt`` and consumed by ``decrypt``).

**Every header is authenticated — the outer JSON is only a routing copy.** This
is the load-bearing correction from the P1 security review. A NaCl ``Box`` seals
the *plaintext*, not the surrounding JSON, and its shared key is **symmetric**:
``Box(A_priv, B_pub) == Box(B_priv, A_pub)``. So the outer ``from``/``to``/
``id``/``ts``/``kind`` headers, which a hostile relay can rewrite freely, must
never be trusted on their own — earlier drafts that did were open to three
distinct attacks (reflection, replay, kind-relabel; see the review). The fix is
one general rule instead of three special cases: the sealed plaintext carries
its **own** copy of the header tuple, and ``open`` accepts the envelope only
when that authenticated copy matches the transport headers *and* names this
brain as recipient and the opening peer as sender::

    plaintext = {"hdr": {"v", "id", "from", "to", "ts", "kind"},  # authenticated
                 "<kind>": <body>}                                # authenticated

Consequences, each an attack the binding closes:

- **Reflection.** A relay copies A's outbound envelope back into A's own mailbox
  with ``from``/``to`` swapped. It decrypts (symmetric key), but the sealed
  ``hdr.from`` still says A while A opens it as if from B — mismatch → reject.
- **Replay.** A relay re-deposits one captured envelope under fresh ``id``s. The
  sealed ``hdr.id`` can't change, so ``hdr.id != outer id`` → reject; and a
  *successful* open guarantees ``env["id"]`` equals the authenticated id, so the
  receiver may dedupe on it safely.
- **Kind relabel / dual body.** ``open`` requires the plaintext to be exactly
  ``{"hdr", <kind>}`` with ``hdr.kind`` equal to the transport kind, so a
  relabeled header or a body smuggling two kind keys is rejected.

A post-decrypt header/structure mismatch is therefore **positive evidence of
tampering**, not malformed input, and is raised as :class:`OpenFailed` (the
tamper family) so the receive loop's audit can tell it apart from pre-crypto
junk (:class:`BadEnvelope`).

The hard size cap (8 MiB on the serialized envelope — room for a full 1M-token
context window, ~4 MB of text, after the seal's base64) is enforced three times on
purpose — at stage (refuse clearly, the sender can split), at the relay (413),
and on receive (reject + audit) — so no single implementation slip re-opens it.
"""
import base64
import datetime as dt
import json
import uuid

from alexandria_vault.federation import sodium
from alexandria_vault.federation import identity

ENVELOPE_VERSION = 2
KINDS = ("intents", "receipt", "ping")
# Mirrors relay/core.MAX_BLOB_BYTES — move them together.
MAX_ENVELOPE_BYTES = 8 * 1024 * 1024
# The transport headers, in the order they are bound into the sealed plaintext.
_HEADER_FIELDS = ("v", "id", "from", "to", "ts", "kind")


class EnvelopeError(ValueError):
    """Base for every way an envelope can be refused."""


class BadEnvelope(EnvelopeError):
    """Malformed shape/schema — refused before (or without) touching crypto."""


class OpenFailed(EnvelopeError):
    """Gate-2 reject: not sealed by that peer for us, tampered, or a header the
    relay rewrote (reflection / replay / relabel) that the sealed copy exposes."""


class EnvelopeTooLarge(EnvelopeError):
    """Over the 8 MiB wire cap — split the payload into smaller pushes."""


# ---- plaintext body builders -------------------------------------------------
def intents_body(intents):
    """Body for kind='intents': the V1 intent dicts, schema unchanged."""
    return {"intents": list(intents)}


def receipt_body(of_id, status, applied=(), parked=()):
    """Body for kind='receipt' — delivery status only, never vault content."""
    if status not in ("received", "applied", "declined"):
        raise ValueError(f"receipt status must be received|applied|declined, got {status!r}")
    return {"receipt": {"of": of_id, "status": status,
                        "applied": list(applied), "parked": list(parked)}}


def ping_body():
    """Body for kind='ping' — the pairing round-trip proof."""
    return {"ping": {}}


# ---- seal / serialize / parse / open ------------------------------------------
def _utc_ts(now=None):
    """An RFC3339 'Z' timestamp in UTC.

    Accepts ``None`` (wall clock), an epoch ``float``/``int`` (the seconds-based
    clock the outbox/worker inject), or a ``datetime``. A naive datetime is
    *assumed* to be UTC (our single convention) rather than silently reinterpreted
    as local wall-clock — the earlier ``astimezone``-only form stamped local time
    with a 'Z', skewing every ts by the host's offset.
    """
    if now is None:
        now = dt.datetime.now(dt.timezone.utc)
    elif isinstance(now, (int, float)):
        now = dt.datetime.fromtimestamp(now, dt.timezone.utc)
    if now.tzinfo is None:
        now = now.replace(tzinfo=dt.timezone.utc)
    else:
        now = now.astimezone(dt.timezone.utc)
    return now.strftime("%Y-%m-%dT%H:%M:%SZ")


def _as_public_key(key):
    """Coerce raw 32 bytes (as peers are stored in config) to a PublicKey; pass a
    PublicKey through unchanged."""
    nacl = sodium()
    return key if isinstance(key, nacl.public.PublicKey) else nacl.public.PublicKey(bytes(key))


def _as_private_key(key):
    nacl = sodium()
    return key if isinstance(key, nacl.public.PrivateKey) else nacl.public.PrivateKey(bytes(key))


def _header(kind, eid, sender_fp, recipient_fp, ts):
    """The authenticated header tuple, one dict used for BOTH the sealed copy and
    the transport copy so they are equal by construction."""
    return {"v": ENVELOPE_VERSION, "id": eid, "from": sender_fp,
            "to": recipient_fp, "ts": ts, "kind": kind}


def seal(sender_priv, recipient_pub, kind, body, *, sender_fp, recipient_fp,
         eid=None, now=None):
    """Build a sealed envelope dict, binding the headers and enforcing the cap.

    The header tuple is sealed *inside* the plaintext (``hdr``) and mirrored in
    the transport JSON; ``open`` later requires the two to agree. ``eid`` lets
    the outbox keep ONE envelope id across deposit retries (sealing happens at
    deposit time; a fresh id per attempt would defeat the receiver's dedupe).
    """
    nacl = sodium()
    if kind not in KINDS:
        raise BadEnvelope(f"unknown envelope kind {kind!r} (expected one of {KINDS})")
    # Body must be EXACTLY {kind: ...} — exclusivity, not mere membership, so a
    # body can never smuggle a second kind key past the relabel defense.
    if not isinstance(body, dict) or set(body) != {kind}:
        raise BadEnvelope(f"body for kind={kind} must be exactly {{{kind!r}: ...}} "
                          "(use the *_body builders)")
    hdr = _header(kind, eid or uuid.uuid4().hex, sender_fp, recipient_fp, _utc_ts(now))
    plaintext = json.dumps({"hdr": hdr, kind: body[kind]},
                           ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    box = nacl.public.Box(_as_private_key(sender_priv), _as_public_key(recipient_pub))
    sealed = box.encrypt(plaintext)  # nonce prepended
    env = dict(hdr)
    env["sealed"] = base64.b64encode(bytes(sealed)).decode("ascii")
    raw = serialize(env)
    if len(raw) > MAX_ENVELOPE_BYTES:
        raise EnvelopeTooLarge(
            f"envelope is {len(raw)} bytes — the wire cap is {MAX_ENVELOPE_BYTES} "
            "(8 MiB); split the payload into smaller pushes")
    return env


def serialize(env):
    """The canonical wire bytes (compact, sorted keys) — what the relay stores."""
    return json.dumps(env, ensure_ascii=False, separators=(",", ":"),
                      sort_keys=True).encode("utf-8")


def parse(raw):
    """Wire bytes/str -> schema-validated transport envelope dict. Every violation
    is BadEnvelope (or EnvelopeTooLarge), raised BEFORE any crypto is attempted.
    Header *authentication* happens later, in :func:`open_envelope`."""
    if isinstance(raw, str):
        raw = raw.encode("utf-8")
    if len(raw) > MAX_ENVELOPE_BYTES:
        raise EnvelopeTooLarge(
            f"envelope is {len(raw)} bytes — over the {MAX_ENVELOPE_BYTES} byte cap")
    try:
        env = json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, ValueError) as exc:
        raise BadEnvelope(f"not JSON: {exc}") from exc
    if not isinstance(env, dict):
        raise BadEnvelope("envelope must be a JSON object")
    if env.get("v") != ENVELOPE_VERSION:
        raise BadEnvelope(f"unsupported envelope version {env.get('v')!r}")
    eid = env.get("id")
    if not isinstance(eid, str) or not (1 <= len(eid) <= 128):
        raise BadEnvelope("envelope id must be a short string")
    for field in ("from", "to"):
        # Shared, fully-anchored fingerprint validator (identity owns the format):
        # rejects a trailing newline and any over-length string a relay might inject.
        if not identity.is_fingerprint(env.get(field)):
            raise BadEnvelope(f"'{field}' must be an ax1… fingerprint")
    if not isinstance(env.get("ts"), str) or len(env["ts"]) > 64:
        raise BadEnvelope("'ts' must be a timestamp string")
    if env.get("kind") not in KINDS:
        raise BadEnvelope(f"unknown envelope kind {env.get('kind')!r}")
    if not isinstance(env.get("sealed"), str) or not env["sealed"]:
        raise BadEnvelope("'sealed' must be a non-empty base64 string")
    return env


def open_envelope(env, recipient_priv, sender_pub):
    """Decrypt + fully authenticate one parsed envelope; return its ``{kind: body}``.

    ``sender_pub`` is the public key of the peer the caller looked up by the
    envelope's ``from`` fingerprint (gate 1, the allowlist). Beyond decrypting,
    this enforces the header binding: the sealed ``hdr`` must equal the transport
    headers AND name this brain (``to``) as recipient and ``sender_pub`` (``from``)
    as sender — closing reflection, replay, and kind-relabel in one check. Any
    failure past a successful decrypt is :class:`OpenFailed` (tamper evidence);
    only pre-crypto shape errors are :class:`BadEnvelope`.
    """
    nacl = sodium()
    recipient_priv = _as_private_key(recipient_priv)
    sender_pub = _as_public_key(sender_pub)
    try:
        blob = base64.b64decode(env["sealed"], validate=True)
    except Exception as exc:  # binascii.Error and friends
        raise BadEnvelope(f"'sealed' is not valid base64: {exc}") from exc
    box = nacl.public.Box(recipient_priv, sender_pub)
    try:
        plaintext = box.decrypt(blob)
    except Exception as exc:  # nacl.exceptions.CryptoError: wrong key or tampered
        raise OpenFailed("envelope failed to open — not sealed by that peer for "
                         "this brain, or tampered in transit") from exc
    try:
        obj = json.loads(plaintext.decode("utf-8"))
    except (UnicodeDecodeError, ValueError) as exc:
        # Authenticated bytes that aren't our JSON: a sender bug, not relay junk.
        raise OpenFailed(f"sealed plaintext is not JSON: {exc}") from exc
    kind = env.get("kind")
    # Exclusivity: exactly the header key plus the one declared kind key.
    if not isinstance(obj, dict) or set(obj) != {"hdr", kind}:
        raise OpenFailed("sealed plaintext structure does not match the declared "
                         "kind (relabeled header or a body smuggling two kinds)")
    # Header binding: the sealed copy must equal the transport headers, and it
    # must name us as recipient and the opening peer as sender (direction).
    expected = _header(kind, env.get("id"),
                       identity.fingerprint(sender_pub),
                       identity.fingerprint(recipient_priv.public_key),
                       env.get("ts"))
    if obj["hdr"] != expected:
        raise OpenFailed("envelope headers are not authenticated — reflected, "
                         "replayed under a new id, or tampered in transit")
    return {kind: obj[kind]}
