import Control.Monad.State
import Crypto.Hash (Digest, MD5, hash)
import Data.ByteString (ByteString)
import Data.ByteString.Char8 (pack)
import qualified Data.Map.Strict as M

salt :: String
salt = "qzyelonm"

-- compute md5 of a bytestring
md5 :: ByteString -> Digest MD5
md5 = hash

-- compute the ith md5 hash
-- the second argument (count) is the number of times the salt+nounce is rehasehd
md5key :: Int -> Int -> String
md5key i count = let
        helper :: String -> Int -> String
        helper s 0 = s
        helper s c = helper (show . md5 $ pack s) (c - 1)
    in
        helper (salt ++ show i) count

-- returns the infinite list [(index, hash)] containing all the hashes
-- starting at the given index
-- the second argument (count) is passed to md5key
hashStream :: Int -> Int -> [(Int, String)]
hashStream startIdx count = map (\i -> (i, md5key i count)) [startIdx..]

-- returns given a count and a string it returns the distinct characters that repeat
-- at least count times consecutively in the string
-- the characters that repeat are returned in the orded that they appear
-- returns an empty string if no repeating characters are found
getRepeating :: Int -> String -> String
getRepeating count str = let
        helper :: String -> String -> String
        helper s rep
            | length s < count = rep
            | otherwise =
                if prefix == replicate count c && c `notElem` rep
                then helper suffix (rep ++ [c])
                else helper (tail s) rep
            where (prefix, suffix) = splitAt count s
                  c = head prefix
    in
        helper str []

-- StreamInfo contains information about the keystream that help with generating
-- the list of valid keys
-- as keys are parsed they are stored in this structure
--
-- streamStartIdx = the index of the next key that should be parsed
-- streamHashCount = how many times a key should be hashed (see md5key)
-- streamTriplets = a map storing the key indices (parsed so far) that contain
--                  triplets and the list of triplet characters in each one
-- streamQuintlets = a map storing for each hexadecimal digit the list of key
--                   indices that contain 5 consecutive appearances of that digit
data StreamInfo = StreamInfo {
    streamStartIdx :: Int,
    streamHashCount :: Int,
    streamTriplets :: M.Map Int String,
    streamQuintlets :: M.Map Char [Int]
} deriving (Show)

-- initial StreamInfo, starts as 0, has no triplets
-- and an empty lists for every hex digit
initStreamInfo :: Int -> StreamInfo
initStreamInfo hc =
    StreamInfo 0 hc M.empty (M.fromList [(c, []) | c <- "0123456789abcdef"])

-- returns the infinite list of parsing the key stream starting form the
-- StreamInfo's startIdx
-- it turns each (idx, key) pair generated by hashStream into a (idx, rep3, rep5)
-- where rep3 = characters that repeat 3 times consecutively in the key
--   and rep5 = characters that repeat 5 times consecutively in the key
parseStream :: StreamInfo -> [(Int, String, String)]
parseStream streamInfo = let
        startIdx = streamStartIdx streamInfo
        hashCount = streamHashCount streamInfo
        parseKey (idx, str) = (idx, getRepeating 3 str, getRepeating 5 str)
    in
        map parseKey $ hashStream startIdx hashCount

-- given a list generated by parseStream it integrates it into the StreamInfo
-- this means incrementing its startIdx and placing the rep3 and rep5 into the
-- triplets and quintlets maps
integrate :: [(Int, String, String)] -> StreamInfo -> StreamInfo
integrate [] si = si
integrate ((idx, r3, r5):xs) (StreamInfo startIdx hashCount triplets quintlets) = let
        triplets' = if null r3 then triplets else M.insert idx r3 triplets
        quintlets' = foldl (flip (M.adjust (idx:))) quintlets r5
    in
        integrate xs (StreamInfo (startIdx + 1) hashCount triplets' quintlets')

-- it continues to parse keys until a key containing a triplet is found
-- the StreamInfo is updated
hashUntilTriplet :: State StreamInfo (Int, String)
hashUntilTriplet = do
    streamInfo <- get
    let stream = parseStream streamInfo
    let (prefix, tripletsPair:_) = span (\(_, r3, _) -> null r3) stream
    let (idx, triplets, _) = tripletsPair

    modify $ integrate prefix
    modify $ integrate [tripletsPair]

    return (idx, triplets)

-- it continues to parse keys until the key at the given index is hashed
-- the StreamInfo is updated
hashUntil :: Int -> State StreamInfo ()
hashUntil idx = do
    streamInfo <- get

    -- how many keys to parse until reaching idx
    let cnt = max 0 (1 + idx - streamStartIdx streamInfo)

    modify $ integrate . take cnt . parseStream $ streamInfo

-- given an index it parses keys until it finds the a triplet at that index
-- or immediately it
-- if StreamInfo already knows about a triplet at idx or after it it
-- returns it and no hashing or parsing is done
getNextTriplet :: Int -> State StreamInfo (Int, String)
getNextTriplet idx = do
    streamInfo <- get

    -- check if there already is a parsed triplet at idx or right after it
    let tripletsPair = M.lookupGE idx $ streamTriplets streamInfo

    -- return tripletsPair if it's found otherwise hash until we find one
    maybe hashUntilTriplet return tripletsPair

-- checks if there is a quintlet of the given character in the
-- hashes [idx + 1, idx + 1000]
-- if StreamInfo already knows about such a quintlet it returns True directly
-- and no hashing or parsing si done
checkTriplet :: Int -> Char -> State StreamInfo Bool
checkTriplet idx chr = do
    let inInterval x = x > idx && x <= idx + 1000
    let check = any inInterval . M.findWithDefault [] chr . streamQuintlets

    pre <- fmap check get
    if pre then return True
    else do
        hashUntil (idx + 1000)
        fmap check get

-- returns an infinite list of valid triplets
validTriplets :: Int -> State StreamInfo [Int]
validTriplets startIdx = do
    (idx, chr:_) <- getNextTriplet startIdx

    check <- checkTriplet idx chr

    -- rest of valid triplets
    rest <- validTriplets $ idx + 1

    -- append to list of the rest if valid
    if check
    then return $ idx:rest
    else return rest

main :: IO ()
main = do
    let streamInfo1 = initStreamInfo 1
    print $ evalState (validTriplets 0) streamInfo1 !! 63

    let streamInfo2 = initStreamInfo 2017
    print $ evalState (validTriplets 0) streamInfo2 !! 63
