module Network.Haskoin.Crypto.Base58
( Address(..)
, addrToBase58
, base58ToAddr
, encodeBase58
, decodeBase58
, encodeBase58Check
, decodeBase58Check
) where

import Control.DeepSeq (NFData, rnf)
import Control.Monad (guard, mzero)

import Data.Maybe (fromMaybe, isJust, listToMaybe)
import Numeric (showIntAtBase, readInt)
import Data.Aeson
    ( Value (String)
    , FromJSON
    , ToJSON
    , parseJSON
    , toJSON
    , withText
    )

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C
import Data.String (IsString, fromString)
import Data.String.Conversions (cs)

import Text.Read (readPrec, parens, lexP, pfail)
import qualified Text.Read as Read (Lexeme(Ident, String))

import Network.Haskoin.Crypto.Hash
import Network.Haskoin.Constants
import Network.Haskoin.Util

b58Data :: ByteString
b58Data = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"

b58 :: Int -> Char
b58 = C.index b58Data

b58' :: Char -> Maybe Int
b58' = flip C.elemIndex b58Data

encodeBase58I :: Integer -> ByteString
encodeBase58I i = cs $ showIntAtBase 58 b58 i ""

decodeBase58I :: ByteString -> Maybe Integer
decodeBase58I s =
    case go of
        Just (r,[]) -> Just r
        _           -> Nothing
  where
    p = isJust . b58'
    f = fromMaybe e . b58'
    go = listToMaybe $ readInt 58 p f (cs s)
    e = error "Could not decode base58"

-- | Encode a 'ByteString' to a base 58 representation.
encodeBase58 :: ByteString -> ByteString
encodeBase58 bs =
    l `mappend` r
  where
    (z, b) = BS.span (== 0) bs
    l = BS.replicate (BS.length z) (BS.index b58Data 0) -- preserve leading 0's
    r | BS.null b = BS.empty
      | otherwise = encodeBase58I $ bsToInteger b

-- | Decode a base58-encoded 'ByteString'. This can fail if the input
-- 'ByteString' contains invalid base58 characters such as 0, O, l, I.
decodeBase58 :: ByteString -> Maybe ByteString
decodeBase58 t =
    BS.append prefix <$> r
  where
    (z, b)  = BS.span (== BS.index b58Data 0) t
    prefix = BS.replicate (BS.length z) 0 -- preserve leading 1's
    r | BS.null b = Just BS.empty
      | otherwise = integerToBS <$> decodeBase58I b

-- | Computes a checksum for the input 'ByteString' and encodes the input and
-- the checksum to a base58 representation.
encodeBase58Check :: ByteString -> ByteString
encodeBase58Check bs = encodeBase58 $ BS.append bs (encode' $ checkSum32 bs)

-- | Decode a base58-encoded string that contains a checksum. This function
-- returns 'Nothing' if the input string contains invalid base58 characters or
-- if the checksum fails.
decodeBase58Check :: ByteString -> Maybe ByteString
decodeBase58Check bs = do
    rs <- decodeBase58 bs
    let (res, chk) = BS.splitAt (BS.length rs - 4) rs
    guard $ chk == encode' (checkSum32 res)
    return res

-- | Data type representing a Bitcoin address
data Address
    -- | Public Key Hash Address
    = PubKeyAddress { getAddrHash :: !Hash160 }
    -- | Script Hash Address
    | ScriptAddress { getAddrHash :: !Hash160 }
       deriving (Eq, Ord)

-- TODO: Test
instance Show Address where
    showsPrec d a = showParen (d > 10) $
        showString "Address " . shows (addrToBase58 a)

-- TODO: Test
instance Read Address where
    readPrec = parens $ do
        Read.Ident "Address" <- lexP
        Read.String str <- lexP
        maybe pfail return $ base58ToAddr $ cs str

-- TODO: Test
instance IsString Address where
    fromString =
        fromMaybe e . base58ToAddr . cs
      where
        e = error "Could not decode bitcoin address"

instance NFData Address where
    rnf (PubKeyAddress h) = rnf h
    rnf (ScriptAddress h) = rnf h

instance FromJSON Address where
    parseJSON = withText "Address" $
        maybe mzero return . base58ToAddr . cs

instance ToJSON Address where
    toJSON = String . cs . addrToBase58

-- | Transforms an Address into a base58 encoded String
addrToBase58 :: Address -> ByteString
addrToBase58 addr = encodeBase58Check $ case addr of
    PubKeyAddress h -> BS.cons addrPrefix   $ getHash160 h
    ScriptAddress h -> BS.cons scriptPrefix $ getHash160 h

-- | Decodes an Address from a base58 encoded String. This function can fail
-- if the String is not properly encoded as base58 or the checksum fails.
base58ToAddr :: ByteString -> Maybe Address
base58ToAddr str = do
    val <- decodeBase58Check str
    guard $ BS.length val == 21
    let f | BS.head val == addrPrefix   = Just PubKeyAddress
          | BS.head val == scriptPrefix = Just ScriptAddress
          | otherwise = Nothing
    f <*> bsToHash160 (BS.tail val)