User:Geraldus

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)