#!/usr/bin/runghc

module SimpleSubCipher where

import qualified Data.List as L
import Data.Maybe

letters :: String
letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"

encodeChar :: Char -> Maybe Int
encodeChar symbol = L.elemIndex symbol letters

decodeChar :: Int -> Char
decodeChar index = letters !! index

addKeyPoint :: Int -> Int -> Int
addKeyPoint key codePoint = mod (codePoint + key) (length letters)

subKeyPoint :: Int -> Int -> Int
subKeyPoint key codePoint = mod (codePoint - key) (length letters)

translateChar :: (Int -> Int -> Int) -> Char -> Char -> Char
translateChar f keyChar plainChar =
  let keyPoint = fromJust (encodeChar keyChar)
   in case encodeChar plainChar of
        Nothing -> plainChar
        Just codePoint -> decodeChar (f keyPoint codePoint)

encryptString :: String -> String -> String
encryptString = zipWith (translateChar addKeyPoint)

decryptString :: String -> String -> String
decryptString = zipWith (translateChar subKeyPoint)

keyCycle :: Int -> String -> String
keyCycle len key = take len (cycle key)

main :: IO ()
main = do
  let message = "THIS IS A SHORT MESSAGE"
  let keyString = keyCycle (length message) "PIZZA"
  putStrLn "Your encrypted message is:"
  putStrLn (encryptString keyString message)
  putStrLn "Your decrypted message is:"
  putStrLn (decryptString keyString (encryptString keyString message))