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

-- | [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
  , I.Prv
  , I.prv
  , I.unPrv
  , I.prvToPub

    -- * Public key
  , XPub(..)
  , I.Pub
  , I.pub
  , I.unPub

    -- * Chain code
  , Chain
  , chain
  , unChain

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

    -- * Subkeys
  , subXPubXPub
  , subXPrvXPrv
  , subXPrvXPub
    -- ** Bare
  , subPubPub
  , subPrvPrv
  , subPrvPub

    -- * 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 Control.Applicative
import Control.Monad
import Data.Bits
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 Data.Word

#ifdef ghcjs_HOST_OS
import qualified BIP32.GHCJS as I
#else
import qualified BIP32.GHC as I
#endif

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

-- | Extended private key.
data XPrv = XPrv !Version !Depth !Fingerprint !Index !Chain !I.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 (I.prvToPub prv)

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

-- | Extended private key.
data XPub = XPub !Version !Depth !Fingerprint !Index !Chain !I.Pub
  deriving (Eq, Show)

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

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

keyPub :: I.Pub -> Key
{-# INLINE keyPub #-}
keyPub = Key . I.unPub

keyPrv :: I.Prv -> Key
{-# INLINE keyPrv #-}
keyPrv = Key . B.cons 0 . I.unPrv

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

-- | 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 'I.Pub' key.
newtype Fingerprint = Fingerprint { unFingerprint :: Word32 }
  deriving (Eq, Show)

fingerprint :: I.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 :: I.Pub -> KeyId
{-# INLINE keyId #-}
keyId = KeyId . hash160 . I.unPub

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

-- | 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 " .
    showsPrec 0 (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

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

-- | Derive a child 'XPrv' subkey from a parent 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 (I.prvToPub prv0)
  pure $ XPrv v d1 f1 i1 ch1 prv1

-- | Derive a child 'XPub' subkey from a parent 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 a child 'XPub' subkey from a parent 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

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

-- | Derive a child 'I.Prv' subkey and 'Chain' code at a particular 'Index' from
-- a parent 'I.Prv' and 'Chain' code.
--
-- Returns 'Nothing' if the given inputs result in an invalid key.
subPrvPrv :: Chain -> I.Prv -> Index -> Maybe (Chain, I.Prv)
subPrvPrv (Chain ch0) prv0 i1@(Index i1w) = do
  let pub0 = I.prvToPub prv0
      x = hmacSHA512 ch0 $ mconcat $ if indexIsHardened i1
            then [B.singleton 0, I.unPrv prv0, word32BE i1w]
            else [I.unPub pub0, word32BE i1w]
      (xl, xr) = B.splitAt 32 x
  prv1 <- I.addPrvTweak prv0 =<< I.tweak xl
  ch1 <- chain xr
  pure (ch1, prv1)

-- | Derive a child 'I.Pub' subkey and 'Chain' code at a particular 'Index' from
-- a parent 'I.Pub' and 'Chain' code.
--
-- The given 'Index' is expected to /not/ be hardened.
--
-- Returns 'Nothing' if the given inputs result in an invalid key.
subPubPub :: Chain -> I.Pub -> Index -> Maybe (Chain, I.Pub)
subPubPub (Chain ch0) pub0 i1@(Index i1w) = do
  guard (not (indexIsHardened i1))
  let x = hmacSHA512 ch0 (I.unPub pub0 <> word32BE i1w)
      (xl, xr) = B.splitAt 32 x
  pub1 <- I.addPubTweak pub0 =<< I.tweak xl
  ch1 <- chain xr
  pure (ch1, pub1)

-- | Derive a child 'I.Pub' subkey and 'Chain' code at a particular 'Index' from
-- a parent 'I.Prv' and 'Chain' code.
--
-- Notice that while @'subPubPub' ('I.prvToPub' prv) i@ will fail with a
-- hardened 'Index', @'subPrvPub' prv i@ may succeed.
--
-- Returns 'Nothing' if the given inputs result in an invalid key.
subPrvPub :: Chain -> I.Prv -> Index -> Maybe (Chain, I.Pub)
subPrvPub ch0 prv0 i1 = fmap I.prvToPub <$> subPrvPrv ch0 prv0 i1

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

-- | 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 I.Prv
getPrv = do
  0 <- Bin.getWord8
  a <- Bin.getByteString 32
  case I.prv a of
    Just b -> pure b
    Nothing -> fail "Bad private key"

getPub :: Bin.Get I.Pub
getPub = do
  a <- Bin.getByteString 33
  case I.pub 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