"""HPKE-Base X25519-HKDF-SHA256-ChaCha20-Poly1305 (RFC 9180).

Production E2E for NATS *content*. Routing fields stay in the clear.
The broker never gets the recipient private key.
"""

from __future__ import annotations  # annotations

import os  # eph nonce
from typing import Tuple  # enc, ct

from cryptography.hazmat.primitives import hashes, hmac  # HKDF pieces
from cryptography.hazmat.primitives.asymmetric.x25519 import (  # KEM
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305  # AEAD
from cryptography.hazmat.primitives.serialization import (  # raw keys
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
)

KEM_ID = 0x0020  # DHKEM(X25519, HKDF-SHA256)
KDF_ID = 0x0001  # HKDF-SHA256
AEAD_ID = 0x0003  # ChaCha20Poly1305
MODE_BASE = 0x00  # HPKE Base
SUITE_ID = b"HPKE" + KEM_ID.to_bytes(2, "big") + KDF_ID.to_bytes(2, "big") + AEAD_ID.to_bytes(2, "big")
N_ENC = 32  # X25519 public
N_PK = 32
N_SK = 32
N_NOISE = 12  # ChaCha nonce length in HPKE
N_K = 32  # key
AAD_DEFAULT = b"verae.npe.v1"  # bind content to this system


def _hmac(key: bytes, data: bytes) -> bytes:
    h = hmac.HMAC(key, hashes.SHA256())  # RFC HMAC
    h.update(data)
    return h.finalize()


def _extract(salt: bytes, ikm: bytes) -> bytes:
    if not salt:  # HMAC key
        salt = b"\x00" * 32
    return _hmac(salt, ikm)


def _expand(prk: bytes, info: bytes, length: int) -> bytes:
    out = b""  # T(0)|T(1)...
    t = b""
    i = 1
    while len(out) < length:
        t = _hmac(prk, t + info + bytes([i]))
        out += t
        i += 1
    return out[:length]


def _labeled_extract(salt: bytes, label: bytes, ikm: bytes) -> bytes:
    labeled = b"HPKE-v1" + SUITE_ID + label + ikm  # RFC 9180
    return _extract(salt, labeled)


def _labeled_expand(prk: bytes, label: bytes, info: bytes, length: int) -> bytes:
    labeled = length.to_bytes(2, "big") + b"HPKE-v1" + SUITE_ID + label + info
    return _expand(prk, labeled, length)


def generate_keypair() -> Tuple[bytes, bytes]:
    """Return (raw_sk, raw_pk) 32+32 bytes."""
    sk = X25519PrivateKey.generate()  # CSPRNG
    pk = sk.public_key()
    return (
        sk.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption()),
        pk.public_bytes(Encoding.Raw, PublicFormat.Raw),
    )


def _dh(sk: bytes, pk: bytes) -> bytes:
    priv = X25519PrivateKey.from_private_bytes(sk)
    pub = X25519PublicKey.from_public_bytes(pk)
    return priv.exchange(pub)


def _extract_and_expand(dh: bytes, kem_context: bytes) -> bytes:
    eae_prk = _labeled_extract(b"", b"eae_prk", dh)
    return _labeled_expand(eae_prk, b"shared_secret", kem_context, 32)


def _key_schedule(shared: bytes, info: bytes) -> Tuple[bytes, bytes]:
    psk_id_hash = _labeled_extract(b"", b"psk_id_hash", b"")
    info_hash = _labeled_extract(b"", b"info_hash", info)
    ks_ctx = bytes([MODE_BASE]) + psk_id_hash + info_hash
    secret = _labeled_extract(shared, b"secret", b"")
    key = _labeled_expand(secret, b"key", ks_ctx, N_K)
    base_nonce = _labeled_expand(secret, b"base_nonce", ks_ctx, N_NOISE)
    return key, base_nonce


def seal(recipient_pk: bytes, plaintext: bytes, aad: bytes = AAD_DEFAULT) -> Tuple[bytes, bytes]:
    """HPKE-Base seal. Returns (encapped_key, ciphertext)."""
    eph_sk, eph_pk = generate_keypair()  # ephemeral
    dh = _dh(eph_sk, recipient_pk)
    kem_context = eph_pk + recipient_pk
    shared = _extract_and_expand(dh, kem_context)
    key, nonce = _key_schedule(shared, b"")
    ct = ChaCha20Poly1305(key).encrypt(nonce, plaintext, aad)
    return eph_pk, ct


def open_ct(recipient_sk: bytes, enc: bytes, ciphertext: bytes, aad: bytes = AAD_DEFAULT) -> bytes:
    """HPKE-Base open. Raises on auth failure."""
    rec_sk = X25519PrivateKey.from_private_bytes(recipient_sk)
    rec_pk = rec_sk.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
    dh = _dh(recipient_sk, enc)
    kem_context = enc + rec_pk
    shared = _extract_and_expand(dh, kem_context)
    key, nonce = _key_schedule(shared, b"")
    return ChaCha20Poly1305(key).decrypt(nonce, ciphertext, aad)