module Main where
import Data.Char (chr, toUpper, ord)
import Data.List (elemIndex, intersperse, findIndices)
import Data.Maybe (fromJust, isNothing)
data Suit = Clubs | Diamonds | Hearts | Spades
deriving(Eq, Ord, Enum, Show)
data Rank = Ace | Two | Three | Four | Five | Six | Seven | Eight | Nine | Ten | Knave | Queen | King
deriving (Eq, Ord, Enum, Show)
data Card = NormalCard Rank Suit
| JokerA
| JokerB
deriving (Eq, Ord, Show)
type Cards = [Card]
isJoker (NormalCard r s) = False
isJoker _ = True
cardVal :: Card -> Int
cardVal (NormalCard r s) = rankVal r + suitVal s
where rankVal = (1+) . fromEnum
suitVal = (*13) . fromEnum
cardVal _ = 53 -- For case of any Joker
-- Letter (Char) from card value
cardLtr c = ltrFromEnum v
where v = 1 + mod (cardVal c - 1) 26
ltrFromEnum n = chr $ 64 + n
unkeydDeck :: Cards
unkeydDeck = allNormals ++ [JokerA, JokerB]
where allNormals = [NormalCard r s | s <- [(Clubs)..], r <-[(Ace)..]]
{-
showDeck :: Cards -> String
showDeck = unwords . map showCard
where showCard JokerA = "A"
showCard JokerB = "B"
showCard c = show $ cardVal c
-}
-- Splits deck at some card into two parts excluding the card itself.
splitAtCard :: Cards -> Card -> (Cards, Cards)
splitAtCard deck c
| isNothing ci = (deck, [])
| otherwise =
let (bc, ac') = splitAt (fromJust ci) deck
ac = drop 1 ac' in
(bc, ac)
where ci = elemIndex c deck
-- Moves card at some position in the deck right (to bottom) by some offset.
-- Deck is circular.
moveCardAt :: Int -> Int -> Cards -> Cards
moveCardAt _ _ [] = []
moveCardAt _ 0 d = d
moveCardAt pos offset deck
| shift > lRgt = moveCardAt 0 (shift - lRgt) $ mvCrd : lft ++ rgt
| otherwise = lft ++ take shift rgt ++ (mvCrd:drop shift rgt)
where shift = offset `mod` lDck
lDck = length deck
lRgt = length rgt
rlPos = pos `mod` lDck
mvCrd = deck!!rlPos
(lft, rgt) = splitAtCard deck mvCrd
-- Moves some card, if it present in the deck.
moveCard crd n dck
| isNothing pos = dck
| otherwise = moveCardAt (fromJust pos) n dck
where pos = crd `elemIndex` dck
-- First step
moveRightJA = moveCard JokerA 1
-- Second step
moveRightJB = moveCard JokerB 2
tripleCut deck
| length ids /= 2 = deck
| otherwise = splR ++ mid ++ splL
where ids = findIndices isJoker deck
(splL, rst) = splitAt iFst deck
(mid, splR) = splitAt iSnd rst
iFst = head ids
iSnd = last ids - iFst + 1
countCut deck = init (drop n deck) ++ take n deck ++ [lastCard]
where lastCard = last deck
n = cardVal lastCard
getKey :: Cards -> Maybe Char
getKey deck
| isJoker crd = Nothing
| otherwise = Just $ cardLtr crd
where crd = deck !! cardVal (head deck)
-- Generates infinite keystream from some deck
keystream deck
| isNothing key = keystream d'
| otherwise = (\ (Just c) -> c) key : keystream d'
where d' = countCut $ tripleCut $ moveRightJB $ moveRightJA deck
key = getKey d'
cleanIn :: String -> String
cleanIn = map toUpper . filterLetters
where filterLetters = filter isLetter
isLetter c = c `elem` ['A'..'Z'] ++ ['a'..'z']
appendXs :: String -> String
appendXs str = str ++ replicate n 'X'
where n = if m == 0 then m else 5 - m
m = length str `mod` 5
splitBy :: Int -> [a] -> [[a]]
splitBy _ [] = []
splitBy n lst = take n lst: splitBy n rest
where rest = drop n lst
bakeIn :: String -> [String]
bakeIn = splitBy 5 . appendXs . cleanIn
convertToInts :: [String] -> [[Int]]
convertToInts = map $ map (\ c -> ord c - 64)
makeIntPairs = zipWith zip
-- Encrypts sanitized message with some keystream.
-- Both message and keystream are represented as [[Int]].
-- Message and keystream are splitted in groups of 5 ints, where each number represents letter's alphabetic order number.
encrypt m k = encryptByIntLists $ makeIntPairs m k
where encryptByIntLists = map encryptPairsList
encryptPairsList = map encryptPair
encryptPair (msgVal, keyVal) = ltrFromEnum $ 1 + (msgVal + keyVal - 1) `mod` 26
-- Decrypts message with some keystream.
decrypt e k = decryptByIntLists $ makeIntPairs e k
where decryptByIntLists = map decryptPairsList
decryptPairsList = map decryptPair
decryptPair (encVal, keyVal) = ltrFromEnum $ 1 + (encVal + 25 - keyVal) `mod` 26
-- | The main entry point.
main :: IO ()
main = do
let msg = "Haskell is Awesome!"
let prp = bakeIn msg
let keys = splitBy 5 $ take (5 * length prp) $ keystream unkeydDeck
putStrLn $ "Prepared > " ++ show prp
putStrLn $ "Keys > " ++ show keys
let msgInts = convertToInts prp
let keyInts = convertToInts keys
let encoded = encrypt msgInts keyInts
let encInts = convertToInts encoded
putStrLn "----------------------------------------------------"
putStrLn $ "Encoded > " ++ show encoded
putStrLn $ "Decoded > " ++ show (decrypt encInts keyInts)
-- example from RubyQuiz
putStrLn $ "Example > " ++ show (decrypt (convertToInts (splitBy 5 "ABVAWLWZSYOORYKDUPVH")) keyInts)