"""E2E encryption for federation payloads.

X25519 key exchange + ChaCha20-Poly1305 for symmetric encryption.
Uses the `cryptography` library (already a koi-net dependency).
"""

import base64
import os

from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305


def generate_keypair() -> tuple[bytes, bytes]:
    """Generate X25519 keypair. Returns (private_key_bytes, public_key_bytes)."""
    private_key = X25519PrivateKey.generate()
    public_key = private_key.public_key()
    return (
        private_key.private_bytes_raw(),
        public_key.public_bytes_raw(),
    )


def derive_shared_key(our_private: bytes, their_public: bytes) -> bytes:
    """Derive shared secret from X25519 key exchange."""
    private_key = X25519PrivateKey.from_private_bytes(our_private)
    public_key = X25519PublicKey.from_public_bytes(their_public)
    return private_key.exchange(public_key)


def encrypt(shared_key: bytes, plaintext: bytes) -> bytes:
    """Encrypt with ChaCha20-Poly1305. Returns nonce + ciphertext."""
    aead = ChaCha20Poly1305(shared_key)
    nonce = os.urandom(12)
    ciphertext = aead.encrypt(nonce, plaintext, None)
    return nonce + ciphertext


def decrypt(shared_key: bytes, data: bytes) -> bytes:
    """Decrypt ChaCha20-Poly1305. Expects nonce (12 bytes) + ciphertext."""
    aead = ChaCha20Poly1305(shared_key)
    nonce = data[:12]
    ciphertext = data[12:]
    return aead.decrypt(nonce, ciphertext, None)


def public_key_to_b64(key_bytes: bytes) -> str:
    return base64.b64encode(key_bytes).decode()


def b64_to_public_key(b64: str) -> bytes:
    return base64.b64decode(b64)