{-# 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
--
-- Implementation based on
-- the [draft SLIP-0032 spec](https://github.com/satoshilabs/slips/blob/71a3549388022820e77aa1f44c80d0f412e5529f/slip-0032.md).
module SLIP32
  ( -- * Parsing
    parse
  , parseXPub
  , parseXPrv
    -- ** Text
  , parseText
  , parseXPubText
  , parseXPrvText

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

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

    -- * Private key
  , XPrv(..)
  , Prv
  , prv
  , unPrv

    -- * Path
  , Path
  , path
  , unPath

    -- * Chain
  , Chain
  , chain
  , unChain
  ) where

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 !Chain !Pub
  deriving (Eq, Show)

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

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

-- | Obtains the derivation path as a list of up to 255 elements.
unPath :: Path -> [Word32]
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 :: [Word32] -> Maybe Path
{-# INLINE path #-}
path x | l < 256 = Just (Path (fromIntegral l) x)
       | otherwise = Nothing
       where l = length x

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

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

-- | 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 | B.length x == 32 = Just (Chain x)
        | otherwise = Nothing

-- | Private key.
--
-- Construct with 'prv'.
newtype Prv = Prv B.ByteString
  deriving (Eq, Show)

-- | Obtain the 33 raw bytes inside a 'Prv'. See 'prv'.
unPrv :: Prv -> B.ByteString
{-# INLINE unPrv #-}
unPrv (Prv x) = x

-- | Construct a 'Prv' key from its raw bytes.
--
-- * 33 bytes in total.
--
-- * The leftmost byte is @0x00@.
--
-- * The remaining 32 bytes are \(ser_{256}(k)\).
--
-- See Bitcoin's [BIP-0032](https://github.com/bitcoin/bips/blob/master/bip-0032.mediawiki)
-- for details.
--
-- 'Nothing' if something is not satisfied.
prv :: B.ByteString -> Maybe Prv
{-# INLINE prv #-}
prv x | B.length x == 33 && B.head x == 0 = Just (Prv x)
      | otherwise = Nothing

-- | Public key.
--
-- Construct with 'pub'.
newtype Pub = Pub B.ByteString
  deriving (Eq, Show)

-- | Obtain the 33 raw bytes inside a 'Pub'. See 'pub'.
unPub :: Pub -> B.ByteString
{-# INLINE unPub #-}
unPub (Pub x) = x

-- | Construct a 'Pub' key from its raw bytes.
--
-- * 33 bytes in total, containing \(ser_{P}(P)\).
--
-- * The leftmost byte is either @0x02@ or @0x03@, depending on the parity
-- of the omitted @y@ coordinate.
--
-- * The remaining 32 bytes are \(ser_{256}(x)\).
--
-- See Bitcoin's [BIP-0032](https://github.com/bitcoin/bips/blob/master/bip-0032.mediawiki)
-- for details.
--
-- 'Nothing' if something is not satisfied.
pub :: B.ByteString -> Maybe Pub
{-# INLINE pub #-}
pub x | B.length x == 33 && (h == 2 || h == 3) = Just (Pub x)
      | otherwise = Nothing
      where h = B.head 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
  depth <- Bin.getWord8
  pa <- Path depth <$> replicateM (fromIntegral depth) Bin.getWord32be
  cc <- Chain <$> Bin.getByteString 32
  kd <- Bin.getByteString 33
  case pub kd of
    Just k -> pure (Left (XPub pa cc k))
    Nothing -> case prv kd of
      Just k -> pure (Right (XPrv pa cc k))
      Nothing -> fail "Bad key prefix"

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

-- | 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 (Pub k)) -> renderText hrpXPub p c (Key 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 (Prv k)) -> renderText hrpXPrv p c (Key k)

-- | The contents of either 'XPub' or 'XPrv'.
newtype Key = Key B.ByteString

-- | Render either an 'XPub' or an 'XPrv' using the SLIP-0032 encoding.
--
-- The rendered 'T.Text' is ASCII compatible.
renderText :: Bech32.HumanReadablePart -> Path -> Chain -> Key -> T.Text
renderText hrp (Path pl p) (Chain c) (Key k)
  = Bech32.encodeLenient hrp $ Bech32.dataPartFromBytes
  $ BL.toStrict $ BB.toLazyByteString
  $ BB.word8 pl <>
    foldMap BB.word32BE p <>
    BB.byteString 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