{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}

-- | SLIP-0032 is an extended serialization format
-- for [BIP-0032](https://github.com/bitcoin/bips/blob/master/bip-0032.mediawiki)
-- wallets.
--
-- This implementation is based on
-- the [draft SLIP-0032 spec](https://github.com/satoshilabs/slips/blob/71a3549388022820e77aa1f44c80d0f412e5529f/slip-0032.md).
--
-- Please refer to the "BIP32" module from
-- the [bip32 library](https://hackage.haskell.org/package/bip32) to
-- find more about 'A.Index' and 'A.Chain'.
--
-- Please refer to the "Bitcoin.Keys" module from
-- the [bitcoin-keys library](https://hackage.haskell.org/package/bitcoin-keys) to
-- find more about 'K.Pub' and 'K.Prv'.
module SLIP32
  ( -- * Parsing
    parse
  , parseXPub
  , parseXPrv

    -- ** Text
  , parseText
  , parseXPubText
  , parseXPrvText

    -- * Rendering
  , renderXPub
  , renderXPrv
    -- ** Text
  , renderXPubText
  , renderXPrvText

    -- * Public key
  , XPub(..)

    -- * Private key
  , XPrv(..)

    -- * Path
  , Path
  , path
  , unPath
  ) where

import qualified Bitcoin.Keys as K
import qualified BIP32 as A
import Control.Applicative
import Control.Monad
import qualified Codec.Binary.Bech32 as Bech32
import qualified Data.Binary.Get as Bin
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import Data.Word

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

-- | Extended public key.
data XPub = XPub !Path !A.Chain !K.Pub
  deriving (Eq, Show)

-- | Extended private key.
data XPrv = XPrv !Path !A.Chain !K.Prv
  deriving (Eq, Show)

-- | Derivation path.
--
-- Construct with 'path'.
data Path = Path !Word8 ![A.Index]
  deriving (Eq, Show)

-- | Obtains the derivation path as a list of up to 255 elements.
unPath :: Path -> [A.Index]
unPath (Path _ x) = x

-- | Construct a derivation 'Path'.
--
-- Hardened keys start from \(2^{31}\).
--
-- @
-- m           = 'path' []
-- m\/0         = 'path' [0]
-- m\/0'        = 'path' [0 + 2^31]
-- m\/1         = 'path' [1]
-- m\/1'        = 'path' [1 + 2^31]
-- m\/0'/1/2'/2 = 'path' [0 + 2^31, 1, 2 + 2^31, 2]
-- @
--
-- See Bitcoin's [BIP-0032](https://github.com/bitcoin/bips/blob/master/bip-0032.mediawiki)
-- for details.
--
-- Returns 'Nothing' if the list length is more than 255.
path :: [A.Index] -> Maybe Path
{-# INLINE path #-}
path x | l < 256 = Just (Path (fromIntegral l) x)
       | otherwise = Nothing
       where l = length x

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

-- | Parse an 'XPub' from its SLIP-0032 representation.
parseXPub :: B.ByteString -> Maybe XPub
{-# INLINE parseXPub #-}
parseXPub = parseXPubText <=< hush . T.decodeUtf8'

-- | Parse an 'XPrv' from its SLIP-0032 representation.
parseXPrv :: B.ByteString -> Maybe XPrv
{-# INLINE parseXPrv #-}
parseXPrv = parseXPrvText <=< hush . T.decodeUtf8'

-- | Parse either an 'XPub' or an 'XPrv' from its SLIP-0032 representation.
parse :: B.ByteString -> Maybe (Either XPub XPrv)
{-# INLINE parse #-}
parse = parseText <=< hush . T.decodeUtf8'

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

-- | Parse an 'XPub' from its SLIP-0032 representation.
--
-- Like 'parseXPub', but takes 'T.Text'.
parseXPubText :: T.Text -> Maybe XPub
{-# INLINE parseXPubText #-}
parseXPubText = either Just (\_ -> Nothing) <=< parseText

-- | Parse an 'XPrv' from its SLIP-0032 representation.
--
-- Like 'parseXPrv', but takes 'T.Text'.
parseXPrvText :: T.Text -> Maybe XPrv
{-# INLINE parseXPrvText #-}
parseXPrvText = either (\_ -> Nothing) Just <=< parseText

-- | Parse either an 'XPub' or an 'XPrv' from its SLIP-0032 representation.
--
-- Like 'parse', but takes 'T.Text'.
parseText :: T.Text -> Maybe (Either XPub XPrv)
parseText = \t0 -> do
  (hrp, dp) <- hush $ Bech32.decodeLenient t0
  raw <- Bech32.dataPartToBytes dp
  case Bin.runGetOrFail getRawSLIP32 (BL.fromStrict raw) of
    Right (lo, _, out@(Left  _)) | BL.null lo && hrp == hrpXPub -> Just out
    Right (lo, _, out@(Right _)) | BL.null lo && hrp == hrpXPrv -> Just out
    _ -> Nothing

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

getRawSLIP32 :: Bin.Get (Either XPub XPrv)
getRawSLIP32 = do
  pa <- getPath
  cc <- getChain
  fmap (Left . XPub pa cc) getPub <|> fmap (Right . XPrv pa cc) getPrv

getPath :: Bin.Get Path
{-# INLINE getPath #-}
getPath = do
  depth <- Bin.getWord8
  Path depth <$> replicateM (fromIntegral depth) getIndex

getIndex :: Bin.Get A.Index
{-# INLINE getIndex #-}
getIndex = A.Index <$> Bin.getWord32be

getChain :: Bin.Get A.Chain
{-# INLINE getChain #-}
getChain = do
  a <- Bin.getByteString 32
  case A.chain a of
     Just b -> pure b
     Nothing -> fail "Bad chain code"

getPrv :: Bin.Get K.Prv
{-# INLINE getPrv #-}
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
{-# INLINE getPub #-}
getPub = do
  a <- Bin.getByteString 33
  case K.parsePub a of
    Just b -> pure b
    Nothing -> fail "Bad public key"

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

-- | Render an 'XPub' using the SLIP-0032 encoding.
renderXPub :: XPub -> B.ByteString
{-# INLINE renderXPub #-}
renderXPub = T.encodeUtf8 . renderXPubText

-- | Render an 'XPub' using the SLIP-0032 encoding.
renderXPrv :: XPrv -> B.ByteString
{-# INLINE renderXPrv #-}
renderXPrv = T.encodeUtf8 . renderXPrvText

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

-- | Render an 'XPub' using the SLIP-0032 encoding.
--
-- The rendered 'T.Text' is ASCII compatible.
renderXPubText :: XPub -> T.Text
{-# INLINE renderXPubText #-}
renderXPubText = \(XPub p c k) -> renderText hrpXPub p c (keyPub k)

-- | Render an 'XPub' using the SLIP-0032 encoding.
--
-- The rendered 'T.Text' is ASCII compatible.
renderXPrvText :: XPrv -> T.Text
{-# INLINE renderXPrvText #-}
renderXPrvText = \(XPrv p c k) -> renderText hrpXPrv p c (keyPrv k)

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

-- | 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

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

-- | Render either an 'XPub' or an 'XPrv' using the SLIP-0032 encoding.
--
-- The rendered 'T.Text' is ASCII compatible.
renderText :: Bech32.HumanReadablePart -> Path -> A.Chain -> Key -> T.Text
renderText hrp (Path pl p) c (Key k)
  = Bech32.encodeLenient hrp
  $ Bech32.dataPartFromBytes
  $ BL.toStrict
  $ BB.toLazyByteString
  $ mconcat [ BB.word8 pl
            , foldMap BB.word32BE (fmap (\(A.Index w) -> w) p)
            , BB.byteString (A.unChain c)
            , BB.byteString k ]


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

hrpXPub :: Bech32.HumanReadablePart
Right hrpXPub = Bech32.humanReadablePartFromText "xpub"

hrpXPrv :: Bech32.HumanReadablePart
Right hrpXPrv = Bech32.humanReadablePartFromText "xprv"

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

hush :: Either a b -> Maybe b
{-# INLINE hush #-}
hush = either (\_ -> Nothing) Just