using Primes, StatsBase

"""
modular_inverse(a::T, n::T) where T <: Integer

Finds the modular inverse of `a` mod `n`
"""
function modular_inverse(a::T, n::T) where T <: Integer
    mod(gcdx(a, n)[2], n)
end # function

"""
𝜑(p::T, q::T) where T <: Integer

Returns Euler's totient function
"""
function 𝜑(p::T, q::T) where T <: Integer
    (p-1)*(q-1)
end # function

function square(x::T) where T <: Integer
    x * x
end

"""
is_legal_public_exponent(e::T, p::T, q::T) where T <: Integer

Returns wether the given exponent is a valid one: 1 < e < 𝜑(p,q) && gcd(e, 𝜑(p,q)) == 1
"""
function is_legal_public_exponent(e::T, p::T, q::T) where T <: Integer
    return one(T) < e && e < 𝜑(p, q) && gcd(e, 𝜑(p,q)) == one(T)
end # function

"""
private_exponent(e::T, p::T, q::T) where T <: Integer

The private exponent is the modular inverse of the public one
"""

function private_exponent(e::T, p::T, q::T) where T <: Integer
    if is_legal_public_exponent(e, p, q)
        return invmod(e, 𝜑(p, q))
    else
        error("Not a legal public exponent for that modulus")
    end
end

"""
encrypt(m::T, e::T, n::T) where T <: Integer

Encrypts the given message with the given public key
"""
function encrypt(m::T, e::T, n::T) where T <: Integer
    if m > n
        error("The modulo is too small to encrypt the message")
    else
        powermod(m, e, n)
    end # if
end # function

"""
decrypt(c::T, d::T, n::T) where T <: Integer

Decrypts the given message with the given private key
"""
function decrypt(c::T, d::T, n::T) where T <: Integer
    return powermod(c, d, n)
end # function


# Generating efficiently primes
pr = primes(4000, 5000)

# Randmoly choosing two prime numbers
p = BigInt(sample(pr))
q = BigInt(sample(pr))

# Defining the needed values for RSA
n = p * q

# Choosing an e such that is coprime with 𝜑
e = BigInt(sample(primes(Int(𝜑(p, q)))))
while !is_legal_public_exponent(e, p, q)
    global e = BigInt(sample(primes(Int(𝜑(p, q)))))
end # while

d = private_exponent(e, p, q)

@time begin
    # Message to be sent
    message = collect("Ciao come va? Tutto bene, grazie, e tu?")
    # @show message

    l = length(message)

    char_message = Array{BigInt}(undef, l)
    # Convert the message in an array of integers. Note: no UNICODE support
    for i in eachindex(message)
        char_message[i] = BigInt(Int64(message[i]))
    end # for
    # @show char_message

    char_encrypted = Array{BigInt}(undef, l)
    # Encrypting the message, char by char, with the formula c ≡ m^e mod n
    for i in eachindex(char_message)
        char_encrypted[i] = encrypt(char_message[i], e, n)
    end # for
    # @show char_encrypted

    char_decrypted = Array{BigInt}(undef, l)
    # Decrypting the message, char by char, with the formula m ≡ c^d mod n
    for i in eachindex(char_encrypted)
        char_decrypted[i] = decrypt(char_encrypted[i], d, n)
    end # for
    # @show char_decrypted

    message_after = Array{Char}(undef, l)
    for i in eachindex(char_decrypted)
        message_after[i] = Char(char_decrypted[i])
    end # for
    # @show message_after
end

@time decrypt(encrypt(BigInt(4000002), e, n), d, n)

# First Compilation:
# 0.101560 seconds (25.54 k allocations: 1.236 MiB)
# 0.000056 seconds (15 allocations: 360 bytes)

# Second Compilation:
# 0.005980 seconds (3.33 k allocations: 152.551 KiB)
# 0.000017 seconds (15 allocations: 360 bytes)