{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StrictData #-}

-- | [BIP-0032](https://github.com/bitcoin/bips/blob/master/bip-0032.mediawiki)
-- Hierarchical Deterministic Wallets, for bitcoin and other cryptocurrencies.
module BIP32
  ( -- * Base58 serialization
    encodeXPrv
  , encodeXPub
  , decodeXPrv
  , decodeXPub
  , decode

    -- * Raw serialization
  , encodeXPrvRaw
  , encodeXPubRaw
  , decodeXPrvRaw
  , decodeXPubRaw
  , decodeRaw

    -- * Private key
  , XPrv(..)
  , xprvToXPub
    -- ** Single private subkey
  , subXPrvXPrv
  , subPrvPrv
    -- ** Multiple private subkeys
  , subsXPrvXPrv
  , subsPrvPrv
    -- ** Single public subkey
  , subXPrvXPub
  , subPrvPub
    -- ** Multiple public subkeys
  , subsXPrvXPub
  , subsPrvPub

    -- * Public key
  , XPub(..)
    -- ** Single public subkey
  , subXPubXPub
  , subPubPub
    -- ** Multiple public subkeys
  , subsXPubXPub
  , subsPubPub

    -- * Chain code
  , Chain
  , chain
  , unChain

    -- * Derivation path
  , Index(..)
  , indexIsHardened
  , indexNext

    -- * Depth
  , Depth(..)

    -- * Fingerprint
  , Fingerprint(..)
  , fingerprint

    -- * Version
  , Version(..)
    -- ** Example versions
  , version_xprv
  , version_xpub
  , version_tprv
  , version_tpub
  , version_Ltpv
  , version_Ltub
  , version_ttpv
  , version_ttub
  ) where

import Bitcoin.Hash (hash160, check32)
import Bitcoin.Hash.Prim (hmacSHA512)
import qualified Bitcoin.Keys as K
import Control.Applicative
import Control.Monad
import Data.Bits
import Data.Maybe
import qualified Data.Binary.Get as Bin
import qualified Data.ByteString as B
import qualified Data.ByteString.Base58 as B58
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Char8 as BL8
import Data.Word

--------------------------------------------------------------------------------

-- | Extended private key.
data XPrv = XPrv
  { xprv_version     :: Version
  , xprv_depth       :: Depth
  , xprv_fingerprint :: Fingerprint
  , xprv_index       :: Index
  , xprv_chain       :: Chain
  , xprv_prv         :: K.Prv
  } deriving (Eq, Show)

-- | Obtain the 'XPub' corresponding to a particular 'XPrv', at a particular
-- 'Version'.
xprvToXPub :: Version -> XPrv -> XPub
xprvToXPub v (XPrv _ d f i c prv) = XPub v d f i c (K.prvToPub prv)

--------------------------------------------------------------------------------

-- | Extended private key.
data XPub = XPub
  { xpub_version     :: Version
  , xpub_depth       :: Depth
  , xpub_fingerprint :: Fingerprint
  , xpub_index       :: Index
  , xpub_chain       :: Chain
  , xpub_pub         :: K.Pub
  } deriving (Eq, Show)

--------------------------------------------------------------------------------

-- | The 33-byte serialized contents of either 'K.Pub' or 'K.Prv'.
newtype Key = Key B.ByteString

keyPub :: K.Pub -> Key
{-# INLINE keyPub #-}
keyPub = Key . K.pubCompressed

keyPrv :: K.Prv -> Key
{-# INLINE keyPrv #-}
keyPrv = Key . B.cons 0 . K.prvRaw

--------------------------------------------------------------------------------

-- | Version bytes.
newtype Version = Version { unVersion :: Word32 }
  deriving (Eq, Show)

-- | @0x0488ade4@, “xprv”, Bitcoin mainnet private.
version_xprv :: Version
version_xprv = Version 0x0488ade4
{-# INLINE version_xprv #-}

-- | @0x0488b21e@, “xpub”, Bitcoin mainnet public.
version_xpub :: Version
version_xpub = Version 0x0488b21e
{-# INLINE version_xpub #-}

-- | @0x04358394@, “tprv”, Bitcoin testnet private.
version_tprv :: Version
version_tprv = Version 0x04358394
{-# INLINE version_tprv #-}

-- | @0x043587cf@, “tpub”, Bitcoin testnet public.
version_tpub :: Version
version_tpub = Version 0x043587cf
{-# INLINE version_tpub #-}

-- | @0x019d9cfe@, “Ltpv”, Litecoin mainnet private.
version_Ltpv :: Version
version_Ltpv = Version 0x019d9cfe
{-# INLINE version_Ltpv #-}

-- | @0x019da462@, “Ltub”, Litecoin mainnet public.
version_Ltub :: Version
version_Ltub = Version 0x019da462
{-# INLINE version_Ltub #-}

-- | @0x0436ef7d@, “ttpv”, Litecoin testnet private.
version_ttpv :: Version
version_ttpv = Version 0x0436ef7d
{-# INLINE version_ttpv #-}

-- | @0x0436f6e1@, “ttub”, Litecoin testnet public.
version_ttub :: Version
version_ttub = Version 0x0436f6e1
{-# INLINE version_ttub #-}

--------------------------------------------------------------------------------

-- | Derivation path depth.
--
-- * @0@ for master nodes — @m@
--
-- * @1@ for level-1 derived keys — @m\/0'@, or @m\/13@, or …
--
-- * @2@ for level-2 derived keys — @m\/0'\/28@, or @m\/44'/0'@, or …
--
-- * … up to @255@.
newtype Depth = Depth { unDepth :: Word8 }
  deriving (Eq, Show)

depthNext :: Depth -> Maybe Depth
{-# INLINE depthNext #-}
depthNext (Depth w) = do
  guard (w /= 0xFF)
  pure (Depth (w + 1))

--------------------------------------------------------------------------------

-- | 4-byte fingerprint of a 'K.Pub' key.
newtype Fingerprint = Fingerprint { unFingerprint :: Word32 }
  deriving (Eq, Show)

fingerprint :: K.Pub -> Fingerprint
{-# INLINE fingerprint #-}
fingerprint = fingerprint' . keyId

fingerprint' :: KeyId -> Fingerprint
{-# INLINE fingerprint' #-}
fingerprint' (KeyId x) =
  let b = B.take 4 x :: B.ByteString
  in Fingerprint $! unsafeShiftL (fromIntegral (B.index b 0)) 24 .|.
                    unsafeShiftL (fromIntegral (B.index b 1)) 16 .|.
                    unsafeShiftL (fromIntegral (B.index b 2))  8 .|.
                                  fromIntegral (B.index b 3)

--------------------------------------------------------------------------------

-- | 20-byte key identifier.
newtype KeyId = KeyId B.ByteString

keyId :: K.Pub -> KeyId
{-# INLINE keyId #-}
keyId = KeyId . hash160 . K.pubCompressed

--------------------------------------------------------------------------------

-- | Chain code.
--
-- Construct with 'chain'.
newtype Chain = Chain B.ByteString
  deriving (Eq)

instance Show Chain where
  showsPrec n (Chain b) = showParen (n > 10) $
    showString "Chain " .
    mappend (BL8.unpack (BB.toLazyByteString (BB.byteStringHex b)))

-- | Construct a 'Chain' code.
--
-- See Bitcoin's [BIP-0032](https://github.com/bitcoin/bips/blob/master/bip-0032.mediawiki)
-- for details.
--
-- 'Nothing' if the 'B.ByteString' length is not 32.
chain :: B.ByteString -> Maybe Chain
{-# INLINE chain #-}
chain x = do
  guard (B.length x == 32)
  Just (Chain x)

-- | Obtain the 32 raw bytes inside a 'Chain'.
unChain :: Chain -> B.ByteString
{-# INLINE unChain #-}
unChain (Chain x) = x

--------------------------------------------------------------------------------

-- | A derivation path 'Index'.
newtype Index = Index { unIndex :: Word32 }
  deriving (Eq, Ord, Show)

-- | Whether a derivation path 'Index' is hardened. That is, \(2^{31}\) or
-- larger.
indexIsHardened :: Index -> Bool
{-# INLINE indexIsHardened #-}
indexIsHardened (Index w) = w >= 0x80000000

-- | Increment the given 'Index' by one, but only if doing so would result in a
-- new 'Index' in the same index group (normal vs. hardened). Otherwise, returns
-- 'Nothing'.
indexNext :: Index -> Maybe Index
{-# INLINE indexNext #-}
indexNext (Index w) = case w of
  0xffffffff -> Nothing -- Would be a normal to hardened transition
  0x7fffffff -> Nothing -- Would be a hardened to normal transition
  _          -> Just (Index (w + 1))

--------------------------------------------------------------------------------

-- | Derive a 'XPrv' subkey from a parent 'XPrv' key.
--
-- Returns 'Nothing' if the given inputs result in an invalid key.
subXPrvXPrv :: XPrv -> Index -> Maybe XPrv
subXPrvXPrv (XPrv v d0 _ _ ch0 prv0) i1 = do
  d1 <- depthNext d0
  (ch1, prv1) <- subPrvPrv ch0 prv0 i1
  let f1 = fingerprint (K.prvToPub prv0)
  pure $ XPrv v d1 f1 i1 ch1 prv1

-- | Derive a 'XPub' subkey from a parent 'XPub' key.
--
-- The given 'Index' is expected to /not/ be hardened.
--
-- Returns 'Nothing' if the given inputs result in an invalid key.
subXPubXPub :: XPub -> Index -> Maybe XPub
subXPubXPub (XPub v d0 _ _ ch0 pub0) i1 = do
  d1 <- depthNext d0
  (ch1, pub1) <- subPubPub ch0 pub0 i1
  let f1 = fingerprint pub0
  pure $ XPub v d1 f1 i1 ch1 pub1

-- | Derive 'XPub' subkey from a parent 'XPrv' key.
--
-- Notice that while @'subXPubXPub' ('xprvToXPub' v) xprv i@ will fail with a
-- hardened 'Index', @'subXPrvXPub' v xprv i@ may succeed.
--
-- Returns 'Nothing' if the given inputs result in an invalid key.
subXPrvXPub :: Version -> XPrv -> Index -> Maybe XPub
subXPrvXPub v xprv i1 = xprvToXPub v <$> subXPrvXPrv xprv i1

--------------------------------------------------------------------------------

-- | Behaves exactly like 'subXPrvXPrv', but deals with 'Chain' and
-- 'K.Prv', rather than an 'XPrv'.
subPrvPrv :: Chain -> K.Prv -> Index -> Maybe (Chain, K.Prv)
subPrvPrv (Chain ch0) prv0 i1@(Index i1w) = do
  let pub0 = K.prvToPub prv0
      x = hmacSHA512 ch0 $ mconcat $ if indexIsHardened i1
            then [B.singleton 0, K.prvRaw prv0, word32BE i1w]
            else [K.pubCompressed pub0, word32BE i1w]
      (xl, xr) = B.splitAt 32 x
  prv1 <- flip K.prvAddTweak prv0 =<< K.parseTweak xl
  ch1 <- chain xr
  pure (ch1, prv1)

-- | Behaves exactly like 'subXPubXPub', but deals with 'Chain' and
-- 'K.Pub', rather than an 'XPub'.
subPubPub :: Chain -> K.Pub -> Index -> Maybe (Chain, K.Pub)
subPubPub (Chain ch0) pub0 i1@(Index i1w) = do
  guard (not (indexIsHardened i1))
  let x = hmacSHA512 ch0 (K.pubCompressed pub0 <> word32BE i1w)
      (xl, xr) = B.splitAt 32 x
  pub1 <- flip K.pubAddTweak pub0 =<< K.parseTweak xl
  ch1 <- chain xr
  pure (ch1, pub1)

-- | Behaves exactly like 'subXPrvXPub', but deals with 'Chain', 'K.Prv' and
-- 'K.Pub', rather than an 'XPrv' and 'XPub'.
subPrvPub :: Chain -> K.Prv -> Index -> Maybe (Chain, K.Pub)
subPrvPub ch0 prv0 i1 = fmap K.prvToPub <$> subPrvPrv ch0 prv0 i1

--------------------------------------------------------------------------------

-- | Lazily derive all the child 'XPrv' subkeys from a parent's 'XPrv'.
--
-- The first child returned is the one at the given 'Index', if any, and the
-- subsequent childs are those immediately after it (see 'indexNext').
--
-- The returned 'Index'es are not necessarily consecutive, since not all
-- 'XPrv' and 'Index' combinations are able to derive valid keys.
-- However, the produced 'Index'es do increase monotonically.
--
-- If the given 'Index' is hardened, then all produced children will have
-- hardened 'Index'es too. If the given 'Index' is normal, then all children
-- will have normal 'Index'es too. That is, this function will never produce
-- children with 'Index'es belonging to another group (normal vs. hardened).
subsXPrvXPrv :: XPrv -> Index -> [XPrv]
subsXPrvXPrv (XPrv v d0 _ _ ch0 prv0) i = do
  d1 <- maybeToList $ depthNext d0
  let f1 = fingerprint (K.prvToPub prv0)
  (i1, ch1, prv1) <- subsPrvPrv ch0 prv0 i
  pure $ XPrv v d1 f1 i1 ch1 prv1

-- | Lazily derive all the 'XPub' subkeys from a parent's 'XPub'.
--
-- The first child returned is the one at the given 'Index', if any, and the
-- subsequent childs are those immediately after it (see 'indexNext').
--
-- The returned 'Index'es are not necessarily consecutive, since not all
-- 'XPub' and 'Index' combinations are able to derive valid keys.
-- However, the produced 'Index'es do increase monotonically.
--
-- The produced list will never include a hardened 'Index'. If a hardened
-- 'Index' is provided as input, the produced list will be empty.
subsXPubXPub :: XPub -> Index -> [XPub]
subsXPubXPub (XPub v d0 _ _ ch0 pub0) i = do
  d1 <- maybeToList $ depthNext d0
  let f1 = fingerprint pub0
  (i1, ch1, pub1) <- subsPubPub ch0 pub0 i
  pure $ XPub v d1 f1 i1 ch1 pub1

-- | Lazily derive all the 'XPub' subkeys from a parent 'XPrv' key.
--
-- The first child returned is the one at the given 'Index', if any, and the
-- subsequent childs are those immediately after it (see 'indexNext').
--
-- The returned 'Index'es are not necessarily consecutive, since not all
-- 'XPrv' and 'Index' combinations are able to derive valid keys.
-- However, the produced 'Index'es do increase monotonically.
--
-- If the given 'Index' is hardened, then all produced children will have
-- hardened 'Index'es too. If the given 'Index' is normal, then all children
-- will have normal 'Index'es too. That is, this function will never produce
-- children with 'Index'es belonging to another group (normal vs. hardened).
--
-- /Notice/ that while @'subsXPubXPub' ('xprvToXPub' v) xprv i@ will fail with a
-- hardened 'Index', @'subsXPrvXPub' v xprv i@ may succeed.
subsXPrvXPub :: Version -> XPrv -> Index -> [XPub]
subsXPrvXPub v xprv i1 = xprvToXPub v <$> subsXPrvXPrv xprv i1

----------------------------------------------------------------------------------

-- | Behaves exactly like 'subsXPrvXPrv', but deals with 'Chain' and
-- 'K.Prv', rather than an 'XPrv'.
subsPrvPrv :: Chain -> K.Prv -> Index -> [(Index, Chain, K.Prv)]
subsPrvPrv ch0 prv0 = go
  where
    go :: Index -> [(Index, Chain, K.Prv)]
    go i = let rest = maybe [] go (indexNext i)
           in case subPrvPrv ch0 prv0 i of
                Just (ch1, prv1) -> (i, ch1, prv1) : rest
                Nothing          -> rest

-- | Behaves exactly like 'subsXPubXPub', but deals with 'Chain' and
-- 'K.Pub', rather than an 'XPub'.
subsPubPub :: Chain -> K.Pub -> Index -> [(Index, Chain, K.Pub)]
subsPubPub ch0 pub0 i0
  | indexIsHardened i0 = []
  | otherwise          = go i0
  where
    go :: Index -> [(Index, Chain, K.Pub)]
    go i = let rest = maybe [] go (indexNext i)
           in case subPubPub ch0 pub0 i of
                Just (ch1, pub1) -> (i, ch1, pub1) : rest
                Nothing          -> rest

-- | Behaves exactly like 'subsXPrvXPub', but deals with 'Chain', 'K.Prv' and
-- 'K.Pub', rather than an 'XPrv' and 'XPub'.
subsPrvPub :: Chain -> K.Prv-> Index -> [(Index, Chain, K.Pub)]
subsPrvPub ch0 prv0 i0 = do
  (i1, ch1, prv1) <- subsPrvPrv ch0 prv0 i0
  pure (i1, ch1, K.prvToPub prv1)

--------------------------------------------------------------------------------

-- | Obtain the Base58 representation for an 'XPub'.
--
-- It can be either 111 or 112 bytes in length, and when rendered as ASCII it
-- looks something like this:
--
-- @
-- xpub661MyMwAqRbcFtXgS5sYJABqqG9YLmC4Q1Rdap9gSE8NqtwybGhePY2gZ29ESFjqJoCu1Rupje8YtGqsefD265TMg7usUDFdp6W1EGMcet8
-- @
encodeXPub :: XPub -> B.ByteString
{-# INLINE encodeXPub #-}
encodeXPub a = base58EncodeWithChecksum (encodeXPubRaw a)

-- | Obtain the Base58 representation for an 'XPrv'.
--
-- It can be either 111 or 112 bytes in length, and when rendered as ASCII it
-- looks something like this:
--
-- @
-- xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi
-- @
encodeXPrv :: XPrv -> B.ByteString
{-# INLINE encodeXPrv #-}
encodeXPrv a = base58EncodeWithChecksum (encodeXPrvRaw a)

base58EncodeWithChecksum :: B.ByteString -> B.ByteString
{-# INLINE base58EncodeWithChecksum #-}
base58EncodeWithChecksum a =
  B58.encodeBase58 B58.bitcoinAlphabet (a <> check32 a)

--------------------------------------------------------------------------------

-- | Encode the 78 raw 'XPub' bytes.
encodeXPubRaw :: XPub -> B.ByteString
{-# INLINE encodeXPubRaw #-}
encodeXPubRaw (XPub v d f i c p) = encodeRaw v d f i c (keyPub p)

-- | Encode the 78 raw 'XPrv' bytes.
encodeXPrvRaw :: XPrv -> B.ByteString
{-# INLINE encodeXPrvRaw #-}
encodeXPrvRaw (XPrv v d f i c p) = encodeRaw v d f i c (keyPrv p)

-- | Encode the 78 raw extended key bytes.
encodeRaw
  :: Version -> Depth -> Fingerprint -> Index -> Chain -> Key -> B.ByteString
encodeRaw (Version v) (Depth d) (Fingerprint f) (Index i) (Chain c) (Key k) =
  BL.toStrict $ BB.toLazyByteString $
    BB.word32BE v <>
    BB.word8 d <>
    BB.word32BE f <>
    BB.word32BE i <>
    BB.byteString c <>
    BB.byteString k

--------------------------------------------------------------------------------

-- | Decode the Base58-encoded 'XPub' representation. See 'encodeXPub'.
decodeXPub :: B.ByteString -> Maybe XPub
{-# INLINE decodeXPub #-}
decodeXPub = decodeXPubRaw <=< base58DecodeWithChecksum

-- | Decode the Base58-encoded 'XPrv' representation. See 'encodeXPrv'.
decodeXPrv :: B.ByteString -> Maybe XPrv
{-# INLINE decodeXPrv #-}
decodeXPrv = decodeXPrvRaw <=< base58DecodeWithChecksum

-- | Decode the Base58-encoded representation of either and 'XPub' or an 'XPub'.
decode :: B.ByteString -> Maybe (Either XPub XPrv)
{-# INLINE decode #-}
decode = decodeRaw <=< base58DecodeWithChecksum

base58DecodeWithChecksum :: B.ByteString -> Maybe B.ByteString
{-# INLINE base58DecodeWithChecksum #-}
base58DecodeWithChecksum = \a -> do
  b <- B58.decodeBase58 B58.bitcoinAlphabet a
  guard (B.length b == 78 + 4)
  let (raw, ch) = B.splitAt 78 b
  guard (ch == check32 raw)
  pure raw

--------------------------------------------------------------------------------

-- | Decode the 78 raw 'XPub' bytes.
decodeXPubRaw :: B.ByteString -> Maybe XPub
{-# INLINE decodeXPubRaw #-}
decodeXPubRaw = either Just (\_ -> Nothing) <=< decodeRaw

-- | Decode the 78 raw 'XPrv' bytes.
decodeXPrvRaw :: B.ByteString -> Maybe XPrv
{-# INLINE decodeXPrvRaw #-}
decodeXPrvRaw = either (\_ -> Nothing) Just <=< decodeRaw

-- | Encode the 78 raw 'XPub' or 'XPrv' bytes.
decodeRaw :: B.ByteString -> Maybe (Either XPub XPrv)
decodeRaw = \b -> do
  guard (B.length b == 78)
  case Bin.runGetOrFail getX (BL.fromStrict b) of
    Right (lo, 78, ex) | BL.null lo -> Just ex
    _ -> Nothing

getX :: Bin.Get (Either XPub XPrv)
getX = do
  v <- Version <$> Bin.getWord32be
  d <- Depth <$> Bin.getWord8
  f <- Fingerprint <$> Bin.getWord32be
  i <- Index <$> Bin.getWord32be
  c <- Chain <$> Bin.getByteString 32
  ek <- fmap Left getPub <|> fmap Right getPrv
  pure $ case ek of
    Left  k -> Left  (XPub v d f i c k)
    Right k -> Right (XPrv v d f i c k)

getPrv :: Bin.Get K.Prv
getPrv = do
  0 <- Bin.getWord8
  a <- Bin.getByteString 32
  case K.parsePrv a of
    Just b -> pure b
    Nothing -> fail "Bad private key"

getPub :: Bin.Get K.Pub
getPub = do
  a <- Bin.getByteString 33
  case K.parsePub a of
    Just b -> pure b
    Nothing -> fail "Bad public key"

--------------------------------------------------------------------------------

word32BE :: Word32 -> B.ByteString -- ^ 4 bytes.
{-# INLINE word32BE #-}
word32BE = BL.toStrict . BB.toLazyByteString . BB.word32BE