module Data.ByteString.Base58
       ( -- * Alphabet
         Alphabet(..)
       , bitcoinAlphabet
       , rippleAlphabet
       , flickrAlphabet
       -- * Encoding and decoding bytestrings
       , encodeBase58
       , decodeBase58
       -- * Encoding and decoding integers
       , encodeBase58I
       , decodeBase58I
       ) where

import Control.Applicative
import Data.Bits
import Data.ByteString ( ByteString )
import Data.Char (chr, ord)
import Data.Map ( Map )
import Data.Maybe
import Data.String
import Data.Typeable ( Typeable )
import Data.Word
import GHC.Generics ( Generic )
import Numeric

import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BC
import qualified Data.List as L
import qualified Data.Map.Strict as M

newtype Alphabet =
    Alphabet
    { unAlphabet  :: ByteString
    } deriving (Ord, Eq, Show, Typeable, Generic, IsString)

bitcoinAlphabet :: Alphabet
bitcoinAlphabet =
    "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"

rippleAlphabet :: Alphabet
rippleAlphabet =
    "rpshnaf39wBUDNEGHJKLM4PQRST7VWXYZ2bcdeCg65jkm8oFqi1tuvAxyz"

flickrAlphabet :: Alphabet
flickrAlphabet =
    "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"

-- | Take 'i' byte from alphabet
b58 :: Alphabet -> Int -> Word8
b58 a i = BS.index (unAlphabet a) i

-- | Lookup position of byte 'w' in alphabet
b58' :: Alphabet -> Word8 -> Maybe Int
b58' a w = BS.elemIndex w (unAlphabet a)

encodeBase58I :: Alphabet -> Integer -> ByteString
encodeBase58I alpha i =
    BC.pack $ showIntAtBase 58 f i ""
  where
    f :: Int -> Char
    f = chr . fromIntegral . b58 alpha . fromIntegral

decodeBase58I :: Alphabet -> ByteString -> Maybe Integer
decodeBase58I alpha s = case go of
    Just (r,[]) -> Just r
    _           -> Nothing
  where
    c = b58' alpha . fromIntegral . ord
    p = isJust . c
    f = fromIntegral . fromJust . c
    go = listToMaybe $ readInt 58 p f (BC.unpack s)

-- | Encode a bytestring to a base 58 representation.
encodeBase58 :: Alphabet -> ByteString -> ByteString
encodeBase58 alpha bs =
    let (z, b) = BS.span (== 0) bs
        l = BS.pack
            $ replicate (BS.length z)
            $ b58 alpha 0
        r | BS.null b = BS.empty
          | otherwise = encodeBase58I alpha
                        $ bsToInteger b
    in BS.append l r

-- | Decode a base 58 encoded bytestring. This can fail if the input bytestring
-- contains invalid base 58 characters such as 0,O,l,I
decodeBase58 :: Alphabet -> ByteString -> Maybe ByteString
decodeBase58 alpha bs =
    let (z, b) = BS.span (== (b58 alpha 0)) bs
        prefix = BS.pack
                 $ replicate (BS.length z) 0
        r | BS.null b = Just BS.empty
          | otherwise = integerToBS <$> decodeBase58I alpha b
    in BS.append prefix <$> r

-- | Decode a big endian Integer from a bytestring
bsToInteger :: ByteString -> Integer
bsToInteger = (L.foldl' f 0) . BS.unpack
  where
    f n w = (toInteger w) .|. shiftL n 8

-- | Encode an Integer to a bytestring as big endian
integerToBS :: Integer -> ByteString
integerToBS 0 = BS.pack [0]
integerToBS i
    | i > 0     = BS.pack $ reverse $ L.unfoldr f i
    | otherwise = error "integerToBS not defined for negative values"
  where
    f 0 = Nothing
    f x = Just $ (fromInteger x :: Word8, x `shiftR` 8)