import Data.Char (chr, ord, toUpper)


-- A bit of self documentation help
type Key = String
type Msg = String


key :: Key
key = "TSTING"
msg :: Msg
msg = "I'm not even mad... This is impressive!"


-- | Checks if character is valid for encoding
isValid :: Char -> Bool
isValid c = let cUp = toUpper c :: Char
             in 'A' <= cUp && cUp <= 'Z'


-- | Given 'key' & 'msg' generate a list of [Maybe Int] indices
-- to map 'msg' from 'key', skipping invalid characters
toIdx :: Key -> Msg -> [Maybe Int]
toIdx k m = map (flip mod keyN <$>) $ toIdx_ 0 m
  where keyN = length k :: Int
        toIdx_ :: Int -> Msg -> [Maybe Int]
        toIdx_ _ "" = []
        toIdx_ acc (c:cs)
          | isValid c = Just acc : toIdx_ (acc + 1) cs
          | otherwise = Nothing : toIdx_ acc cs


-- | Given 'key' & 'msg' generate a list of numbers representing
-- the amount to shift 'msg' characters based on 'key'
toShifts :: Key -> Msg -> [Int]
toShifts k m = map toKey (toIdx k m)
  where kUp = map toUpper k :: Key
        toKey :: Maybe Int -> Int
        toKey Nothing  = 0
        toKey (Just x) = ord (kUp!!x) - ord 'A'


-- | Given 'by' & 'c', shift the Char 'c' by amount 'by'. 'by' can be both
-- positive & negative as well as 0.
shift :: Int -> Char -> Char
shift by c
  | isValid c && c >= 'a' = shift_ $ ord 'a'
  | isValid c && c >= 'A' = shift_ $ ord 'A'
  | otherwise = c
  where cONorm    = ord (toUpper c) - ord 'A' :: Int
        azN       = ord 'Z' - ord 'A' :: Int
        shift_ :: Int -> Char
        shift_ aO = chr $ aO + mod (by + cONorm) azN


-- Encode & decode a message using the given key.
vigenere, unVigenere :: Key -> Msg -> Msg
vigenere   k m = zipWith shift (toShifts k m) m
unVigenere k m = zipWith shift (map negate $ toShifts k m) m