import random
from sympy import isprime, gcd

# Define the 27-character alphabet mapping
ALPHABET = " ABCDEFGHIJKLMNOPQRSTUVWXYZ"
ALPHABET_MAP = {char: idx for idx, char in enumerate(ALPHABET)}
REVERSE_ALPHABET_MAP = {idx: char for idx, char in enumerate(ALPHABET)}


# Extended Euclidean Algorithm to find modular inverse
def modular_inverse(a, m):
    """Returns the modular inverse of a under modulo m, or None if it doesn't exist."""
    m0, x0, x1 = m, 0, 1
    while a > 1:
        q = a // m
        a, m = m, a % m
        x0, x1 = x1 - q * x0, x0
    if x1 < 0:
        x1 += m0
    return x1


# RSA key generation function
def generate_rsa_keys(bit_length=512):
    p = generate_large_prime(bit_length)
    q = generate_large_prime(bit_length)

    n = p * q
    phi = (p - 1) * (q - 1) # euler's function

    e = 65537
    while gcd(e, phi) != 1:
        e = random.randint(2, phi - 1)

    d = modular_inverse(e, phi)
    if d is None:
        raise ValueError("Modular inverse for the chosen e does not exist.")

    return (n, e), (n, d)


def generate_large_prime(bit_length):
    """Generates a random prime number of approximately bit_length bits."""
    while True:
        candidate = random.getrandbits(bit_length)
        if isprime(candidate):
            return candidate


# RSA encryption
def rsa_encrypt(plaintext, public_key):
    n, e = public_key
    # Convert plaintext to numeric format using the alphabet map
    numeric_plaintext = [ALPHABET_MAP[char] for char in plaintext if char in ALPHABET_MAP]
    # Encrypt each character in the numeric plaintext
    ciphertext = [pow(char, e, n) for char in numeric_plaintext]
    return ciphertext


# RSA decryption
def rsa_decrypt(ciphertext, private_key):
    n, d = private_key
    # Decrypt each character in the ciphertext
    decrypted_numbers = [pow(char, d, n) for char in ciphertext]
    # Convert decrypted numbers back to characters
    plaintext = ''.join(REVERSE_ALPHABET_MAP[num] for num in decrypted_numbers)
    return plaintext


# Validation functions
def validate_plaintext(plaintext):
    """Validates that the plaintext contains only characters from the 27-character alphabet."""
    for char in plaintext:
        if char not in ALPHABET_MAP:
            raise ValueError(f"Invalid character in plaintext: {char}")


def validate_ciphertext(ciphertext, n):
    """Validates that each part of the ciphertext is a number less than n."""
    for part in ciphertext:
        if not (0 <= part < n):
            raise ValueError(f"Invalid ciphertext value: {part}")


# Example usage
if __name__ == "__main__":
    # Generate keys
    public_key, private_key = generate_rsa_keys()

    # Plaintext to encrypt
    plaintext = "HELLO WORLD"
    validate_plaintext(plaintext)

    # Encrypt the plaintext
    ciphertext = rsa_encrypt(plaintext, public_key)
    print("Ciphertext:", ciphertext)

    # Validate ciphertext and decrypt
    validate_ciphertext(ciphertext, public_key[0])
    decrypted_text = rsa_decrypt(ciphertext, private_key)
    print("Decrypted Text:", decrypted_text)