-- Note to self: I probably way overcomplicated this solution
module Vigenere where

import qualified Data.Map  as M
import qualified Data.List as L
import Control.Monad.Reader

type CharSet = [Char] -- alias for string, used to distinguish from text input
type CipherTable = M.Map Char CharSet
type ReadEnv a = ReaderT CipherEnv Maybe a

data CipherEnv = Cipher { cipherChars :: CharSet
                        , cipherTable :: CipherTable
                        } deriving (Show)

-- | Creates a character set with no duplicate entries
mkCharSet :: String -> CharSet
mkCharSet = L.nub

-- | Generates a row in the character table
genRow :: Char -> CharSet -> CharSet
genRow char charset = b ++ a
  where (a, b) = break (== char) charset

-- | Constructs a table from a character set
mkTable :: CharSet -> CipherTable
mkTable charset = M.fromList $ map generator charset
  where generator c = (c, genRow c charset)

-- | generates a cryptographic key based on a message and a keyword
--   if the message is "hello there", and the keyword is "apple";
--   then the resulting key will be "appleapplea"
genKey :: String -> String -> String
genKey msg kw = kw `padTo` length msg
  where padTo xs size = take size $ cycle xs

-- | gets a row from the table at a specific key k
--   if no row exists, Nothing is returned
getRow ::  Char -> ReadEnv String
getRow k = do
  env <- ask
  row <- lift $ M.lookup k $ cipherTable env
  return row

-- | Cryptographic function
--   gets the vth letter on the kth row on the table
encryptLetter :: Char -> Char -> ReadEnv Char
encryptLetter k v = do
  env <- ask
  index <- lift $ L.elemIndex v $ cipherChars env
  row <- getRow k
  return $ row !! index

-- | Cryptographic function
--   gets the index of the cth letter on the kth row in the table,
--   then uses that index to find the cth letter in the base
decryptLetter :: Char -> Char -> ReadEnv Char -- base[table[k, c]]
decryptLetter k c = do
  env <- ask
  row <- getRow k
  tableIndex <- lift $ L.elemIndex c row
  let base = cipherChars env
  return $ base !! tableIndex

-- | maps a cryptographic function to two strings
crypt :: (Char -> Char -> ReadEnv Char) -> String -> String -> ReadEnv String
crypt f w kw = sequence zipped
  where key = genKey w kw
        zipped = zipWith f key w

-- encrypts the string using a message and a keyword
encrypt :: String -> String -> ReadEnv String
encrypt = crypt encryptLetter

-- decrypts using a message and a keyword
decrypt :: String -> String -> ReadEnv String
decrypt = crypt decryptLetter

mkEnv :: CharSet -> CipherEnv
mkEnv ch = Cipher c $ mkTable c
  where c = mkCharSet ch

-- | alias for runReaderT
runEnv :: ReaderT r m a -> r -> m a
runEnv = runReaderT

-- testing functions
encodeTest :: IO ()
encodeTest = do
  let characterSet = mkEnv $ " ," ++ ['A' .. 'Z'] ++ ['a' .. 'z']
      message      = "Hello, this is a test"
      keyword      = "bananas"
      cipher       = encrypt message keyword
      result       = runEnv cipher characterSet
  putStrLn $ "Here's the output of running \"" ++ message ++ "\" through the function: "
  putStrLn $ show result

decodeTest :: IO ()
decodeTest = do
  let characterSet = mkEnv $ " ," ++ ['A' .. 'Z'] ++ ['a' .. 'z']
      message      = "kEYLbbsUHVSnIkbAnTRSl"
      keyword      = "bananas"
      cipher       = decrypt message keyword
      result       = runEnv cipher characterSet
  putStrLn $ "Here's the output of running \"" ++ message ++ "\" through the function: "
  putStrLn $ show result