#!/usr/bin/runghc

module AffineCipher (encryptString, decryptString, m) where

import CryptoMath
import Data.List
import Data.Maybe
import Test.QuickCheck

symbolSet :: String
symbolSet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz1234567890 !?."

m :: Int
m = length symbolSet

-- |
-- >>> encodeChar 'A'
-- Just 0
encodeChar :: Char -> Maybe Int
encodeChar symbol = elemIndex symbol symbolSet

-- |
-- >>> decodeChar 0
-- 'A'
decodeChar :: Int -> Char
decodeChar index = symbolSet !! index

-- |
-- >>> encKeyPoint 5 2 4
-- 22

-- |
-- >>> encKeyPoint 5 2 65
-- 63
encKeyPoint :: Int -> Int -> Int -> Int
encKeyPoint keyA keyB codePoint = mod (codePoint * keyA + keyB) m

-- |
-- >>> decKeyPoint 5 2 22
-- 4

-- |
-- >>> decKeyPoint 5 2 63
-- 65
decKeyPoint :: Int -> Int -> Int -> Int
decKeyPoint keyA keyB codePoint = mod ((codePoint - keyB) * fromJust (modInv keyA m)) m

-- fromJust is reasonable above, because we assume valid keys were selected for encryption.

-- |
-- >>> encryptChar 5 2 'A'
-- 'C'
encryptChar :: Int -> Int -> Char -> Char
encryptChar keyA keyB c = case encodeChar c of
  Nothing -> c
  Just codePoint -> decodeChar (encKeyPoint keyA keyB codePoint)

-- |
-- >>> decryptChar 5 2 'B'
-- 'N'
decryptChar :: Int -> Int -> Char -> Char
decryptChar keyA keyB c = case encodeChar c of
  Nothing -> c
  Just codePoint -> decodeChar (decKeyPoint keyA keyB codePoint)

-- |
-- >>> encryptString 5 2 "zed."
-- "8UP!"
encryptString :: Int -> Int -> String -> String
encryptString keyA keyB = map (encryptChar keyA keyB)

-- |
-- >>> decryptString 5 2 "8UP!"
-- "zed."
decryptString :: Int -> Int -> String -> String
decryptString keyA keyB = map (decryptChar keyA keyB)

main :: IO ()
main = do
  let message = "A computer would deserve to be called intelligent if it could deceive a human into believing that it was human. -Alan Turing"
  let keyA = 43
  let keyB = 56
  putStrLn "Your encrypted message is:"
  putStrLn (encryptString keyA keyB message)
  putStrLn "Your decrypted ciphertext is:"
  putStrLn (decryptString keyA keyB (encryptString keyA keyB message))