import Data.Maybe (fromJust,isJust)
import Data.List.Split (chunksOf)
import Data.Tuple (swap)
import Data.Tuple.Extra (first,second,dupe)
import System.Environment (getArgs)
import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString.Lazy.Char8 as C
import qualified Data.Map as M

-- Example Usage: echo "hello world" | ./base64  | ./base64 -d

-- TODO
-- cleanup/simplify
-- unit testing (cross check with cli base64, different lengths)
-- flag for line widht / check flags / filename on command-line 
-- (see man base64)
-- compare performance . named map? array?

main = do arg<-getArgs
          dat<-B.getContents
          B.putStr.B.pack.encode.B.unpack $  dat
{-
          if length arg == 0 
             then putStr . (++"\n") . encode64 . map fromIntegral . B.unpack $ dat
             else B.putStr. B.pack . map fromIntegral . decode64 .  C.unpack $ dat
-}

encode = map ((+65).(`mod` 24))

table64    = zip [0..] (['A'..'Z']++['a'..'z']++['0'..'9']++['+','/'])

enc64 k = M.lookup k mp where mp = M.fromList table64
dec64 k = M.lookup k mp where mp = M.fromList $ map swap table64

encode64 :: [Int] -> [Char]
encode64 = map (fromJust.enc64.flip mod 64)  -- concat . map (pad . first enc . second length . dupe) . chunksOf 3
           where enc = map (fromJust . enc64) . sumC
                 pad (v,l) =take 4 $ take (1+l) v ++ "==="
                 sumC = map fst . reverse . take 4 . drop 1 .iterate to64 . (,) 0 . sum  . map (uncurry (*)) . zip mult
                 mult = map (256^) [2,1,0]
                 to64 (r,v) = let r' = v `mod` 64 in (r',(v-r')`div`64)

decode64 :: [Char]->[Int]
decode64 = map fst .  concat . map (rem . first (reverse . take 3 . drop 1 . iterate to256 . (,) 0 . dec. map (fromJust) . filter (isJust) . map dec64).second (length.filter(=='=')). dupe) . chunksOf 4 . filter (/='\n')
           where dec = sum . map (uncurry (*)) . zip (map (64^) [3,2..]) 
                 to256 (r,v) = let r' = v `mod` 256 in (r',(v-r')`div`256)
                 rem (v,l) = take (3-l) v