"""ElGamal — 2048-bit ElGamal encryption.

Ported from net.i2p.crypto.ElGamalEngine.
Uses Python's native int arithmetic (arbitrary precision).

I2P-specific format:
- Max plaintext: 222 bytes
- Ciphertext: always 514 bytes (two 257-byte components)
- Prepends random nonzero byte + SHA-256 hash before encryption
- Verifies hash on decryption
"""

import hashlib
import os

# I2P's 2048-bit ElGamal prime (from CryptoConstants.java)
ELGAMAL_PRIME = int(
    "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
    "29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
    "EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
    "E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
    "EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"
    "C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"
    "83655D23DCA3AD961C62F356208552BB9ED529077096966D"
    "670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"
    "E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"
    "DE2BCBF6955817183995497CEA956AE515D2261898FA0510"
    "15728E5A8AACAA68FFFFFFFFFFFFFFFF", 16
)

# Generator
ELGAMAL_GENERATOR = 2

# Max plaintext size: prime_size(256) - 1(random) - 32(hash) - 1(format) = 222
MAX_CLEARTEXT = 222

# Component size in output (257 bytes each, zero-padded)
COMPONENT_SIZE = 257

# Total ciphertext size
CIPHERTEXT_SIZE = COMPONENT_SIZE * 2  # 514


class ElGamalEngine:
    """ElGamal encryption/decryption with I2P-specific formatting.

    Uses 2048-bit prime from CryptoConstants.
    Not recommended for new applications — prefer X25519/ChaCha20.
    """

    @staticmethod
    def encrypt(data: bytes, public_key: bytes) -> bytes:
        """Encrypt up to 222 bytes using ElGamal.

        Args:
            data: plaintext, max 222 bytes
            public_key: 256-byte ElGamal public key

        Returns:
            514-byte ciphertext (two 257-byte components)

        Raises:
            ValueError: if data exceeds 222 bytes
        """
        if len(data) > MAX_CLEARTEXT:
            raise ValueError(f"Data too long: {len(data)} > {MAX_CLEARTEXT}")

        # Build plaintext: [random_nonzero_byte][sha256(data)][data]
        rand_byte = 0
        while rand_byte == 0:
            rand_byte = os.urandom(1)[0]
        hash_val = hashlib.sha256(data).digest()
        plaintext = bytes([rand_byte]) + hash_val + data

        # Pad to 256 bytes
        plaintext = plaintext.ljust(256, b"\x00")

        # Convert to integer (big-endian)
        m = int.from_bytes(plaintext, "big")

        # Convert public key to integer
        y = int.from_bytes(public_key[:256], "big")

        # Pick random k, compute y_val = g^k mod p, d = m * y^k mod p
        k = int.from_bytes(os.urandom(256), "big") % (ELGAMAL_PRIME - 2) + 1
        y_val = pow(ELGAMAL_GENERATOR, k, ELGAMAL_PRIME)
        d = (m * pow(y, k, ELGAMAL_PRIME)) % ELGAMAL_PRIME

        # Encode as 257-byte components (zero-padded big-endian)
        y_bytes = y_val.to_bytes(COMPONENT_SIZE, "big")
        d_bytes = d.to_bytes(COMPONENT_SIZE, "big")

        return y_bytes + d_bytes

    @staticmethod
    def decrypt(encrypted: bytes, private_key: bytes) -> bytes | None:
        """Decrypt ElGamal ciphertext.

        Args:
            encrypted: 514-byte ciphertext
            private_key: 256-byte ElGamal private key

        Returns:
            Decrypted data, or None if hash verification fails.
        """
        if len(encrypted) != CIPHERTEXT_SIZE:
            return None

        # Extract components
        y_bytes = encrypted[:COMPONENT_SIZE]
        d_bytes = encrypted[COMPONENT_SIZE:]

        y_val = int.from_bytes(y_bytes, "big")
        d = int.from_bytes(d_bytes, "big")

        # Convert private key to integer
        x = int.from_bytes(private_key[:256], "big")

        # Compute m = d * y^(-x) mod p = d * y^(p-1-x) mod p
        m = (d * pow(y_val, ELGAMAL_PRIME - 1 - x, ELGAMAL_PRIME)) % ELGAMAL_PRIME

        # Convert back to bytes (256 bytes)
        plaintext = m.to_bytes(256, "big")

        # Parse: [random_byte][32-byte hash][data_with_padding]
        hash_val = plaintext[1:33]
        padded_data = plaintext[33:]

        # Try all possible data lengths (longest first) using the hash
        # to find the correct boundary between data and zero-padding.
        for length in range(len(padded_data), -1, -1):
            candidate = padded_data[:length]
            if hashlib.sha256(candidate).digest() == hash_val:
                return candidate

        return None

    @staticmethod
    def generate_keypair() -> tuple[bytes, bytes]:
        """Generate ElGamal key pair.

        Returns:
            (public_key, private_key) — both 256 bytes
        """
        # Random private key x in [1, p-2]
        x = int.from_bytes(os.urandom(256), "big") % (ELGAMAL_PRIME - 2) + 1
        # Public key y = g^x mod p
        y = pow(ELGAMAL_GENERATOR, x, ELGAMAL_PRIME)

        pub = y.to_bytes(256, "big")
        priv = x.to_bytes(256, "big")
        return pub, priv