{-# LANGUAGE OverloadedStrings #-}
module ASCII85 where
import Data.ByteString (ByteString, pack, unpack, concat)
import Data.Word (Word8)
import Data.Char (digitToInt, chr, ord, isSeparator, intToDigit)
import Data.List.Split (chunksOf, splitWhen)
import Data.List (intercalate)
import Numeric (showIntAtBase)

binStringToInt :: String -> Int
binStringToInt = foldl (\ accu x -> accu * 2 + digitToInt x) 0

intToFixedBinString :: Int -> Int -> String
intToFixedBinString n bitcount = extraPadding ++ binstr
  where binstr = showIntAtBase 2 intToDigit n ""
        extraPadding = take (bitcount - length binstr) $ repeat ('0')

wordToBinString :: Word8 -> String
wordToBinString word = intToFixedBinString (fromIntegral $ toInteger word) 8

modsForIntTo85 :: Int -> Int -> [Word8] -> [Word8]
modsForIntTo85 _ (-1) res = reverse res
modsForIntTo85 n left res = modsForIntTo85 (n `div` 85) (left - 1) (res ++ [newmod])
  where newmod = fromIntegral $ toInteger (n `mod` 85)

chunkToAscii85 :: [Word8] -> [Word8]
chunkToAscii85 [0, 0, 0, 0] = [122]
chunkToAscii85 chunk = ascii85Chunk
  where fullChunk = chunk ++ take (4 - length chunk) (repeat 0)
        fullChunkAsBinString = foldr (++) "" (map wordToBinString fullChunk)
        fullChunkAsInt = binStringToInt fullChunkAsBinString
        ascii85Chunk = take (length chunk + 1) $ map (+33) (modsForIntTo85 fullChunkAsInt 4 [])

toAscii85 :: ByteString -> ByteString
toAscii85 bytes = Data.ByteString.concat [prefix, ascii85Data, suffix]
  where prefix = "<~"
        bytesInChunks = chunksOf 4 (unpack bytes)
        ascii85Data = pack $ foldl (++) [] $ map chunkToAscii85 bytesInChunks
        suffix = "~>"

wordsBackToInt :: [Word8] -> Int -> Int
wordsBackToInt [] _ = 0
wordsBackToInt (w:rest) cnt = ((fromIntegral $ toInteger w) - 33) * (85^cnt) + wordsBackToInt rest (cnt - 1)

binStringToWord :: String -> Word8
binStringToWord = foldl (\ accu x -> accu * 2 + (fromInteger $ toInteger (digitToInt x))) 0

decodeChunk :: [Word8] -> [Word8]
decodeChunk chunk = decodedChunk
  where fullChunk = chunk ++ take (5 - length chunk) (repeat 117)
        fullChunkAsInt = wordsBackToInt fullChunk 4
        fullChunkAsBinaryComponents = chunksOf 8 $ intToFixedBinString fullChunkAsInt 32
        characters = map binStringToWord fullChunkAsBinaryComponents
        decodedChunk = take (length chunk - 1) characters

fromAscii85 :: ByteString -> ByteString
fromAscii85 bytes = pack $ foldl (++) [] $ map decodeChunk (chunksOf 5 bytesWithNoZeds)
  where unpackedBytes = unpack bytes
        isValidChar chr = chr >= 33 && (chr <= 117 || chr == 122)
        cleanBytes = filter isValidChar $ take (length unpackedBytes - 4) (drop 2 unpackedBytes)
        longhandNulls = take 5 $ repeat 33
        bytesWithNoZeds = intercalate longhandNulls $ splitWhen (==122) cleanBytes