{-# LANGUAGE ScopedTypeVariables #-}

import System.Random (randomRIO)
import Data.Bits (shiftR, (.&.), testBit)
import Control.Monad (replicateM)

-- Function to generate a random large prime number
generatePrime :: Integer -> IO Integer
generatePrime bitLength = do
    let min = 2^(bitLength - 1)
    let max = 2^bitLength - 1
    candidate <- randomRIO (min, max)
    putStrLn $ "Trying candidate: " ++ show candidate
    isPrimeCandidate <- isProbablyPrime candidate 5
    if isPrimeCandidate
        then do
            putStrLn $ "Found prime: " ++ show candidate
            return candidate
        else generatePrime bitLength

-- Miller-Rabin primality test
isProbablyPrime :: Integer -> Int -> IO Bool
isProbablyPrime n k
    | n < 2 = return False
    | n == 2 || n == 3 = return True
    | even n = return False
    | otherwise = do
        let (r, d) = decompose (n - 1)
        witnesses <- replicateM k (randomRIO (2, n - 2))
        return $ all (millerRabinTest n r d) witnesses

-- Decompose (n-1) to find r and d such that (n-1) = 2^r * d
decompose :: Integer -> (Integer, Integer)
decompose n = go n 0
  where
    go m r
        | m `mod` 2 == 0 = go (m `div` 2) (r + 1)
        | otherwise = (r, m)

-- Perform the Miller-Rabin test with one witness
millerRabinTest :: Integer -> Integer -> Integer -> Integer -> Bool
millerRabinTest n r d a = powMod a d n == 1 || any (\i -> powMod a (d * 2^i) n == n - 1) [0..(r - 1)]

-- Fast modular exponentiation
powMod :: Integer -> Integer -> Integer -> Integer
powMod base exp modulus = powMod' base exp modulus 1
  where
    powMod' _ 0 _ result = result
    powMod' b e m result
      | e .&. 1 == 1 = powMod' (b * b `mod` m) (shiftR e 1) m (result * b `mod` m)
      | otherwise    = powMod' (b * b `mod` m) (shiftR e 1) m result

-- Extended Euclidean Algorithm to find GCD and Bézout coefficients
extendedGCD :: Integer -> Integer -> (Integer, Integer, Integer)
extendedGCD a 0 = (a, 1, 0)
extendedGCD a b = 
    let (g, s, t) = extendedGCD b (a `mod` b)
    in (g, t, s - (a `div` b) * t)

-- Compute the modular multiplicative inverse
modInverse :: Integer -> Integer -> Integer
modInverse e phi = let (_, inv, _) = extendedGCD e phi in inv `mod` phi

-- Encrypt: C = M^e mod n
encrypt :: Integer -> Integer -> Integer -> Integer
encrypt m e n = powMod m e n

-- Decrypt: M = C^d mod n
decrypt :: Integer -> Integer -> Integer -> Integer
decrypt c d n = powMod c d n

main :: IO ()
main = do
    -- Generate large primes p and q
    putStrLn "Generating prime p (1024 bits)..."
    p <- generatePrime 1024
    putStrLn "Generating prime q (1024 bits)..."
    q <- generatePrime 1024
    let n = p * q
    let phi = (p - 1) * (q - 1)
    let e = 65537  -- Commonly used public exponent

    -- Ensure e is coprime with phi
    putStrLn "Computing modular inverse d..."
    let d = modInverse e phi

    putStrLn $ "Prime p: " ++ show p
    putStrLn $ "Prime q: " ++ show q
    putStrLn $ "Modulus n: " ++ show n
    putStrLn $ "Euler's totient phi(n): " ++ show phi
    putStrLn $ "Public exponent e: " ++ show e
    putStrLn $ "Private exponent d: " ++ show d

    -- Encrypt and decrypt a message
    let message = 1234567890  -- Example message
    putStrLn $ "Encrypting message: " ++ show message
    let ciphertext = encrypt message e n
    putStrLn $ "Encrypted message: " ++ show ciphertext
    let decryptedMessage = decrypt ciphertext d n
    putStrLn $ "Decrypted message: " ++ show decryptedMessage
    putStrLn $ "Original Message: " ++ show message
    putStrLn $ "Encrypted Message: " ++ show ciphertext
    putStrLn $ "Decrypted Message: " ++ show decryptedMessage