"""ChaCha20-Blake3 AEAD boxes, counter-nonce channels, and a fast CSPRNG.
Everything supports in-place operation into caller buffers."""

from __future__ import annotations

import hmac
import secrets
import threading
from collections.abc import Buffer
from typing import overload

import blake3
from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20

KEY_LEN = 32
NONCE_LEN = 12
_MAC_KEY_BLOCK = 64


class WireError(ValueError):
    """Failure an unauthenticated peer can cause with crafted packets:
    truncated ciphertext, bad authentication tag. Distinct from TypeError,
    which signals caller misuse (wrong buffer sizes, wrong key length)."""


def nonce_increment(nonce: Buffer) -> None:
    """Increment the nonce in place, little-endian."""
    n = memoryview(nonce).cast("B")
    for i in range(len(n)):
        if n[i] < 255:
            n[i] += 1
            return
        n[i] = 0


def _session(key: Buffer, nonce: Buffer, ad: Buffer | None, maclen: int):
    """ChaCha20 stream + keyed BLAKE3 with ad absorbed (None when maclen=0)."""
    key = memoryview(key).cast("B")
    nonce = memoryview(nonce).cast("B")
    if len(key) != KEY_LEN:
        raise TypeError("key must be 32 bytes")
    if len(nonce) != NONCE_LEN:
        raise TypeError("nonce must be 12 bytes")
    if not 0 <= maclen <= 32:
        raise TypeError("maclen must be in 0..32")
    # IETF layout: 32-bit little-endian block counter || 96-bit nonce
    stream = Cipher(ChaCha20(key, bytes(4) + nonce), mode=None).encryptor()
    if maclen == 0:
        return stream, None
    tag_key = stream.update(bytes(_MAC_KEY_BLOCK))[:KEY_LEN]
    mac = blake3.blake3(key=tag_key)
    ad = memoryview(ad).cast("B") if ad else memoryview(b"")
    mac.update(len(ad).to_bytes(8, "little"))
    mac.update(ad)
    return stream, mac


def _tag(mac, ct_len: int, maclen: int) -> bytes:
    m = mac.copy()
    m.update(ct_len.to_bytes(8, "little"))
    return m.digest()[:maclen]


def encrypt_into(
    out: Buffer,
    msg: Buffer,
    key: Buffer,
    nonce: Buffer | None = None,
    *,
    ad: Buffer | None = None,
    maclen: int = 16,
) -> memoryview:
    """Encrypt message under key, with optional nonce, ad and mac, into out buffer."""
    msg = memoryview(msg).cast("B")
    hdr = NONCE_LEN if nonce is None else 0
    used = hdr + len(msg) + maclen
    mv = memoryview(out).cast("B")
    if len(mv) < used:
        raise TypeError("output buffer too small")
    out = mv[:used]
    if nonce is None:
        nonce = out[:NONCE_LEN]
        random(nonce)
    stream, mac = _session(key, nonce, ad, maclen)
    ct = out[hdr : hdr + len(msg)]
    stream.update_into(msg, ct)
    if mac is not None:
        mac.update(ct)
        out[hdr + len(msg) :] = _tag(mac, len(msg), maclen)
    return out


def decrypt_into(
    out: Buffer,
    ct: Buffer,
    key: Buffer,
    nonce: Buffer | None = None,
    *,
    ad: Buffer | None = None,
    maclen: int = 16,
) -> memoryview:
    """Decrypt ciphertext into out buffer. Other arguments must match those at encryption point."""
    ct = memoryview(ct).cast("B")
    if nonce is None:
        if len(ct) < NONCE_LEN + maclen:
            raise WireError("ciphertext buffer too short")
        nonce, ct = ct[:NONCE_LEN], ct[NONCE_LEN:]
    elif len(ct) < maclen:
        raise WireError("ciphertext buffer too short")
    used = len(ct) - maclen
    mv = memoryview(out).cast("B")
    if len(mv) < used:
        raise TypeError("output buffer too small")
    out = mv[:used]
    stream, mac = _session(key, nonce, ad, maclen)
    if mac is not None:
        mac.update(ct[:used])
        if not hmac.compare_digest(bytes(ct[used:]), _tag(mac, used, maclen)):
            raise WireError("invalid authentication tag")
    stream.update_into(ct[:used], out)
    return out


def encrypt(
    msg: Buffer,
    key: Buffer,
    nonce: Buffer | None = None,
    *,
    ad: Buffer | None = None,
    maclen: int = 16,
) -> bytearray:
    """Encrypt message under key, with optional nonce, ad and mac, returning ciphertext."""
    msg = memoryview(msg).cast("B")
    out = bytearray(len(msg) + maclen + (NONCE_LEN if nonce is None else 0))
    encrypt_into(out, msg, key, nonce, ad=ad, maclen=maclen)
    return out


def decrypt(
    ct: Buffer,
    key: Buffer,
    nonce: Buffer | None = None,
    *,
    ad: Buffer | None = None,
    maclen: int = 16,
) -> bytearray:
    """Decrypt and return message. Other arguments must match those at encryption point."""
    ct = memoryview(ct).cast("B")
    hdr = NONCE_LEN if nonce is None else 0
    if len(ct) < hdr + maclen:
        raise WireError("ciphertext buffer too short")
    out = bytearray(len(ct) - hdr - maclen)
    decrypt_into(out, ct, key, nonce, ad=ad, maclen=maclen)
    return out


class _Sequence:
    def __init__(self, key: Buffer, *, maclen: int = 16):
        key = memoryview(key).cast("B")
        if len(key) != KEY_LEN:
            raise TypeError("key must be 32 bytes")
        self._key = key
        self._maclen = maclen
        self._nonce = bytearray(NONCE_LEN)


class Sender(_Sequence):
    """Encrypting end of a message channel: messages are sent in order,
    each individually authenticated. Pair with Receiver on the same key."""

    def encrypt_into(
        self, out: Buffer, msg: Buffer, *, ad: Buffer | None = None
    ) -> memoryview:
        ret = encrypt_into(out, msg, self._key, self._nonce, ad=ad, maclen=self._maclen)
        nonce_increment(self._nonce)
        return ret

    def encrypt(self, msg: Buffer, *, ad: Buffer | None = None) -> bytearray:
        msg = memoryview(msg).cast("B")
        out = bytearray(len(msg) + self._maclen)
        self.encrypt_into(out, msg, ad=ad)
        return out


class Receiver(_Sequence):
    """Decrypting end of a message channel: messages are verified and
    returned in order. Pair with Sender on the same key."""

    def decrypt_into(
        self, out: Buffer, ct: Buffer, *, ad: Buffer | None = None
    ) -> memoryview:
        ret = decrypt_into(out, ct, self._key, self._nonce, ad=ad, maclen=self._maclen)
        nonce_increment(self._nonce)
        return ret

    def decrypt(self, ct: Buffer, *, ad: Buffer | None = None) -> bytearray:
        ct = memoryview(ct).cast("B")
        if len(ct) < self._maclen:
            raise WireError("ciphertext buffer too short")
        out = bytearray(len(ct) - self._maclen)
        self.decrypt_into(out, ct, ad=ad)
        return out


_random = threading.local()


@overload
def random(nbytes: int, /) -> bytearray: ...
@overload
def random(buf: Buffer, /) -> None: ...
def random(arg: int | Buffer, /):
    """random(nbytes) -> nbytes as bytearray; random(buf) fills in place.
    Seeded from OS random, thread safe."""
    if isinstance(arg, int):
        mv = memoryview(ret := bytearray(arg))
    else:
        mv, ret = memoryview(arg).cast("B"), None
    zeroes = bytes(len(mv))
    try:
        _random.enc.update_into(zeroes, mv)
    except (AttributeError, ValueError):
        _random.enc = Cipher(
            ChaCha20(secrets.token_bytes(KEY_LEN), bytes(16)), mode=None
        ).encryptor()
        _random.enc.update_into(zeroes, mv)
    return ret