from Crypto.Util.number import getPrime, inverse
import random


def miller_rabin(n, k=5):
    if n == 2 or n == 3:
        return True
    if n <= 1 or n % 2 == 0:
        return False

    # Write n-1 as d * 2^r
    r = 0
    d = n - 1
    while d % 2 == 0:
        d //= 2
        r += 1

    # Perform the test k times
    for _ in range(k):
        a = random.randint(2, n - 2)
        x = pow(a, d, n)  # x = a^d % n
        if x == 1 or x == n - 1:
            continue

        # Keep squaring x and check for n-1
        for _ in range(r - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                break
        else:
            return False

    return True


# Euclidean algorithm
def extended_gcd(a, b):
    """Computes the GCD of a and b, along with coefficients x and y such that ax + by = gcd(a, b)."""
    if b == 0:
        return a, 1, 0

    g, x1, y1 = extended_gcd(b, a % b)

    x = y1
    y = x1 - (a // b) * y1
    return g, x, y


def mod_inverse(a, m):
    g, x, _ = extended_gcd(a, m)
    if g != 1:  # a and m must be coprime
        raise ValueError(f"{a} has no modular inverse modulo {m}")
    return x % m


def my_prime(bits: int = 1024) -> int:
    i = random.randint(1 << (bits - 1), 1 << bits)
    return i if miller_rabin(i) else my_prime(bits)


print(my_prime(8))


# Key Generation
def generate_keys(bits: int = 1024) -> tuple[tuple[int, int], tuple[int, int]]:
    p = my_prime(bits // 2)
    q = my_prime(bits // 2)
    n = p * q
    phi = (p - 1) * (q - 1)
    e = 65537  # Common public exponent
    d = inverse(e, phi)
    return (e, n), (d, n)  # Public, Private keys


# Encryption
def encrypt(msg: int, pubkey: tuple[int, int]) -> int:
    e, n = pubkey
    return pow(msg, e, n)


# Decryption
def decrypt(cipher: int, privkey: tuple[int, int]) -> int:
    d, n = privkey
    return pow(cipher, d, n)


# Example Usage
pub, priv = generate_keys()
pub = (177349751, 2144805071)
priv = (1698859991, 2144805071)
message = 42
ciphertext = encrypt(message, pub)
plaintext = decrypt(ciphertext, priv)

print(f"Message: {message}, Ciphertext: {ciphertext}, Decrypted: {plaintext}")