"""
Chat message encryption and sender authentication.

Design A of `docs/chat-sender-keys.md`, decided 2026-09-07. What it is, and what
it deliberately is not, in the order the decisions were made:

**Not a ratchet.** `senderkeys.py` implements Signal-style sender keys and is
unused by production. With sender keys distributed under the group key and a
node that serves history to devices which were not present, the node must retain
and hand out each chain's *earliest* key — and a chain key at iteration *i*
yields every message key from *i* onward by pure HKDF. Forward secrecy is then
zero, and the ratchet is computing HKDF over a value every member already holds.
The property is given up on the record rather than inherited by accident.

**A key per group, per epoch, per device.** The node generates an epoch key and
delivers it to members wrapped under the current group key. Each device derives
its *own* subkey from it, by name, so:

* two devices never share an AES key, and nonce reuse across devices is
  impossible without any coordination — the property per-device ratchet chains
  were wanted for, obtained by derivation instead of by mutable state (which is
  C1 one level down, and is exactly what `GroupSenderKeyStore` got wrong);
* a receiver derives any sender's subkey from the epoch key it already has, so
  nothing is distributed per device and there is no per-device state to
  persist, migrate or lose;
* history keeps working, for new members and new devices alike, because the
  epoch key does not move as messages are sent.

**Epochs, not rotation of the archive.** The epoch key is wrapped under the
group key *at delivery*, never stored under it — so rotating the group key
(which is the documented step after removing a member) is a re-wrap and costs
nothing. A new epoch is opened when the set of devices that may read *future*
messages shrinks; old epochs are kept and still delivered to current members,
which is what keeps the history readable to the people who could already read it.

**Signing is separate from encryption**, and is what actually establishes who
said something. The signature is over the *ciphertext*, so it can be checked
before decryption and by anyone holding the roster, and it names the device key
the node pinned — not a fresh key the sender invented, which is what made the
sender-key distribution format forgeable by any member.

What none of this protects against, stated per the v5/v6 convention: the node
operator and every current group member hold the group key and therefore the
epoch keys. This is the same boundary as file access, by design. It protects
against someone who obtains the node's storage without the keystore password.
"""

from __future__ import annotations

import os

import msgpack
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

EPOCH_KEY_LEN = 32
# 96-bit, the WebCrypto AES-GCM standard — and the same reasoning as
# `groupbox.py`: one key per device per epoch puts the NIST SP 800-38D ceiling
# of 2**32 random nonces far out of reach of a person typing.
NONCE_LEN = 12
SIG_LEN = 64

_DEVICE_KEY_INFO = "meshbay:chat:dev:v1"
_SIG_PREFIX = b"meshbay:chat:v1"


def new_epoch_key() -> bytes:
    """A fresh epoch key. Generated by the node — never by a member (C5b)."""
    return os.urandom(EPOCH_KEY_LEN)


def device_key(epoch_key: bytes, group_id: str, device_b64: str) -> bytes:
    """
    The AES-256 subkey one device encrypts under, in one group, in one epoch.

    Derived by name, so every member can compute every device's subkey from the
    epoch key and nobody has to distribute anything per device. `salt=None` here
    matches `salt: new Uint8Array(0)` in crypto.js — RFC 5869 extracts with a
    zero key either way, which `deriveChunkKey` already relies on in production
    and `test_js_python_parity` holds.
    """
    if not epoch_key:
        raise ValueError("no epoch key")
    if not device_b64:
        raise ValueError("no device")
    info = f"{_DEVICE_KEY_INFO}|{group_id}|{device_b64}".encode()
    return HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None, info=info,
    ).derive(epoch_key)


def associated_data(group_id: str, epoch: int) -> bytes:
    """
    What a chat ciphertext is bound to.

    The group stops a message being moved between two groups on one node; the
    epoch stops one being re-presented as belonging to a later key, which would
    otherwise let a member who kept an old epoch key make an old message look
    current. Same reasoning as `groupbox.associated_data`, one field longer.
    """
    return f"chat_msg|{group_id}|{epoch}".encode()


def signing_transcript(group_id: str, epoch: int, device: bytes, nonce: bytes,
                       ct: bytes) -> bytes:
    """
    What the sending device signs — the ciphertext, never the plaintext.

    Over the ciphertext so a receiver can establish authorship *before* it
    decrypts, and so anyone holding the roster can verify a stored message
    without the epoch key. Length-prefixed and domain-separated, per L4, like
    every other transcript in this codebase: without the lengths, a message
    could be re-cut into a different one with the same bytes.

    Note what is *not* in here: this connection's nonce. Every other transcript
    binds one, and this one cannot — a receiver reading history has no access to
    the connection the message arrived on. Replay is therefore refused by
    storage instead, on the unique `(device, nonce)` pair.
    """
    out = bytearray(_SIG_PREFIX)
    for field in (group_id.encode(), str(epoch).encode(), device, nonce, ct):
        out += len(field).to_bytes(4, "big")
        out += field
    return bytes(out)


def seal(epoch_key: bytes, group_id: str, epoch: int, device_b64: str,
         device_raw: bytes, sk: Ed25519PrivateKey, payload: dict) -> dict:
    """
    One message, sealed and signed: `{nonce, ct, sig}` for the caller to merge.

    The routing and authentication fields — `format`, `epoch`, `device` — stay
    in clear, because a receiver has to select a key and check a signature
    before it can decrypt, and because the node routes on them without being
    able to read anything.
    """
    key = device_key(epoch_key, group_id, device_b64)
    # Random per message. Never derived from the payload: two identical messages
    # under one device's key would then reuse it, and AES-GCM's failure under
    # nonce reuse is not graceful.
    nonce = os.urandom(NONCE_LEN)
    ct = AESGCM(key).encrypt(
        nonce, msgpack.packb(payload, use_bin_type=True),
        associated_data(group_id, epoch))
    sig = sk.sign(signing_transcript(group_id, epoch, device_raw, nonce, ct))
    return {"nonce": nonce, "ct": ct, "sig": sig}


def verify(group_id: str, epoch: int, device_raw: bytes, nonce: bytes,
           ct: bytes, sig: bytes) -> bool:
    """Whether `sig` is this device's signature over this ciphertext."""
    try:
        pk = Ed25519PublicKey.from_public_bytes(device_raw)
        pk.verify(sig, signing_transcript(group_id, epoch, device_raw, nonce, ct))
        return True
    except Exception:
        return False


def open_message(epoch_key: bytes, group_id: str, epoch: int, device_b64: str,
                 nonce: bytes, ct: bytes) -> dict:
    """
    Open a sealed message. Raises on anything that does not open.

    Never a partial result and never a default — an unopenable message is not an
    empty one, and rendering it as blank would make a message nobody can read
    indistinguishable from a message nobody wrote. The caller marks it as
    unreadable and shows the gap, which is the honest thing for a reader to see.
    """
    key = device_key(epoch_key, group_id, device_b64)
    plain = AESGCM(key).decrypt(bytes(nonce), bytes(ct),
                                associated_data(group_id, epoch))
    payload = msgpack.unpackb(plain, raw=False)
    if not isinstance(payload, dict):
        raise ValueError("chat: sealed payload is not a map")
    return payload