import Data.Char (ord)
import Data.Word (Word8)
import Data.Tuple.Extra
import Data.Tuple (swap)
import Data.Array (array,(!),Array)
import Data.Bits (shiftL,shiftR,(.&.))
import System.Environment (getArgs,getProgName,withArgs)
import Test.Hspec (hspec,describe,it,shouldBe)
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Char8 as Char8

main = do args<-getArgs
          let arg=args!!0
          prg<-getProgName
          if length args == 0
            then BL.getContents >>= BL.putStr . encode64
            else if length args == 1 && arg == "-t" 
                    then withArgs [] test 
                    else if length args == 1 && arg == "-d" 
                        then BL.getContents >>= BL.putStr . decode64
                        else putStrLn $ unlines 
                                [ "usage: "++prg++" [FLAG]"
                                 ,"reads from stdin"
                                 ,"default mode is encode"
                                 ,"flags: -d decode, -t unit test"]

encode64 = chunkBL 76 (flip BL.append (Char8.pack "\n")) . chunkBL 3 (from 6 4.to 8)
decode64 = chunkBL 4 (from 8 3.to 6)

tab64enc :: Array Int Word8
tab64enc = array (0,63) $ zip [0..] $ map (fromIntegral.ord) $ 
           ['A'..'Z']++['a'..'z']++['0'..'9']++['+','/']

to s =  (\(v,l)->(v`shiftL`((3-fromIntegral l)*s),l+1)) . first (BL.foldl' f 0) . second (BL.length) . dupe 
       where f a v = a `shiftL` s + fromIntegral v 

from s l (x,m) = BL.take (fromIntegral l) . BL.pack $ (++repeat 61) $ map f (take (fromIntegral m) $ reverse [0..l-1]) 
                 where f v = tab64enc ! (x `shiftR` (v*s) .&. (2^s-1))

chunkBL n f b | BL.null b  = b
              | otherwise = let l = f . BL.take n $ b
                                r = chunkBL n f . BL.drop n $ b
                            in  BL.append l r    

-- |Unit Testing
tests=[  ("any carnal pleasure.","YW55IGNhcm5hbCBwbGVhc3VyZS4=\n")
        ,("any carnal pleasure","YW55IGNhcm5hbCBwbGVhc3VyZQ==\n")
        ,("any carnal pleasur","YW55IGNhcm5hbCBwbGVhc3Vy\n")
        ,("any carnal pleasu","YW55IGNhcm5hbCBwbGVhc3U=\n")
        ,("any carnal pleas","YW55IGNhcm5hbCBwbGVhcw==\n") ]

test = hspec $ do describe "Base64" $ do 
                    it "encode" $ do  mapM_ (tst encode64) tests
                    it "decode" $ do  mapM_ (tst decode64) (map swap tests)
       where tst f (d,e) = (f (Char8.pack d)) `shouldBe` (Char8.pack e)