"""ChaCha20 stream cipher.

Ported from net.i2p.crypto.ChaCha20.
Implements RFC 7539 ChaCha20 encryption/decryption,
wrapping Python's ``cryptography`` library.

Key: 32 bytes. Nonce: 12 bytes (standard) or 16 bytes
(4-byte little-endian counter prefix + 12-byte nonce).
"""

from __future__ import annotations

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms


class ChaCha20:
    """ChaCha20 stream cipher wrapper. Thread-safe, stateless (all static methods)."""

    # RFC 7539 specifies initial counter = 1 for encryption
    _INITIAL_COUNTER = 1

    def __init__(self) -> None:  # pragma: no cover
        raise TypeError("ChaCha20 should not be instantiated; use static methods")

    @staticmethod
    def encrypt(key: bytes, nonce: bytes, plaintext: bytes) -> bytes:
        """Encrypt plaintext with ChaCha20.

        Args:
            key: 32-byte key.
            nonce: 12-byte nonce (RFC 7539) or 16-byte nonce
                   (first 4 bytes = little-endian counter, next 12 = nonce).

        Returns:
            Ciphertext of the same length as plaintext.

        Raises:
            ValueError: if key or nonce lengths are invalid.
        """
        return ChaCha20._apply(key, nonce, plaintext)

    @staticmethod
    def decrypt(key: bytes, nonce: bytes, ciphertext: bytes) -> bytes:
        """Decrypt ciphertext with ChaCha20.

        ChaCha20 is a stream cipher — decryption is identical to encryption.

        Args:
            key: 32-byte key.
            nonce: 12-byte nonce (RFC 7539) or 16-byte nonce
                   (first 4 bytes = little-endian counter, next 12 = nonce).

        Returns:
            Plaintext of the same length as ciphertext.

        Raises:
            ValueError: if key or nonce lengths are invalid.
        """
        return ChaCha20._apply(key, nonce, ciphertext)

    @staticmethod
    def _apply(key: bytes, nonce: bytes, data: bytes) -> bytes:
        """Core ChaCha20 XOR operation (encrypt == decrypt).

        The ``cryptography`` library expects a 16-byte nonce composed of
        a 4-byte little-endian initial counter followed by the 12-byte
        RFC 7539 nonce.
        """
        if len(key) < 32:
            raise ValueError(f"Key must be at least 32 bytes, got {len(key)}")

        k = key[:32]

        if len(nonce) == 12:
            # Standard 12-byte nonce: prepend counter = 1 (little-endian)
            full_nonce = ChaCha20._INITIAL_COUNTER.to_bytes(4, "little") + nonce
        elif len(nonce) == 16:
            # 16-byte nonce already includes counter prefix
            full_nonce = nonce
        else:
            raise ValueError(
                f"Nonce must be 12 or 16 bytes, got {len(nonce)}"
            )

        algorithm = algorithms.ChaCha20(k, full_nonce)
        cipher = Cipher(algorithm, mode=None)
        encryptor = cipher.encryptor()
        return encryptor.update(data) + encryptor.finalize()