{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-type-defaults #-}

module Main (main) where

import qualified Bitcoin.Keys as K
import qualified Data.ByteString as B
import qualified Data.ByteString.Base16 as B16
import qualified Data.ByteString.Base58 as B58
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BL
import Data.Foldable
import Data.Maybe (isJust)
import Data.Word
import qualified Test.Tasty as Tasty
import qualified Test.Tasty.Runners as Tasty
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase, (@?=), (@=?))
import Test.Tasty.Hedgehog (testProperty, HedgehogTestLimit(..))
import Hedgehog (MonadGen, property, forAll, (===), diff)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import qualified BIP32 as S

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

main :: IO ()
main = Tasty.defaultMainWithIngredients
    [ Tasty.consoleTestReporter
    , Tasty.listingTests
    ] $ Tasty.localOption (HedgehogTestLimit (Just 1000))
      $ tt

tt :: TestTree
tt = testGroup "BIP32"
  [ testCase "versions" $ do
      -- https://bitcoin.stackexchange.com/questions/38878/how-does-the-bip32-version-bytes-convert-to-base58
      let sv = B.take 4 . B58.encodeBase58 B58.bitcoinAlphabet
             . flip mappend (B.replicate 78 0x00)
             . BL.toStrict . BB.toLazyByteString . BB.word32BE . S.unVersion
      sv S.version_xprv @?= "xprv"
      sv S.version_xpub @?= "xpub"
      sv S.version_tprv @?= "tprv"
      sv S.version_tpub @?= "tpub"
      sv S.version_Ltpv @?= "Ltpv"
      sv S.version_Ltub @?= "Ltub"
      sv S.version_ttpv @?= "ttpv"
      sv S.version_ttub @?= "ttub"

  , testProperty "encodeXPubRaw length" $ property $ do
      xpub <- forAll genXPub
      B.length (S.encodeXPubRaw xpub) === 78

  , testProperty "encodeXPrvRaw length" $ property $ do
      xprv <- forAll genXPrv
      B.length (S.encodeXPrvRaw xprv) === 78

  , testProperty "encodeXPub length" $ property $ do
      xpub <- forAll genXPub
      let xpubB = S.encodeXPub xpub
      diff (B.length xpubB) (>=) 111
      diff (B.length xpubB) (<=) 112

  , testProperty "encodeXPrv length" $ property $ do
      xprv <- forAll genXPrv
      let xprvB = S.encodeXPrv xprv
      diff (B.length xprvB) (>=) 111
      diff (B.length xprvB) (<=) 112

  , testProperty "XPub raw roundtrip" $ property $ do
      xpub :: S.XPub <- forAll genXPub
      let xpubB = S.encodeXPubRaw xpub
      Just xpub === S.decodeXPubRaw xpubB

  , testProperty "XPrv raw roundtrip" $ property $ do
      xprv :: S.XPrv <- forAll genXPrv
      let xprvB = S.encodeXPrvRaw xprv
      Just xprv === S.decodeXPrvRaw xprvB

  , testProperty "XPub Base58 roundtrip" $ property $ do
      xpub :: S.XPub <- forAll genXPub
      let xpubB = S.encodeXPub xpub
      Just xpub === S.decodeXPub xpubB

  , testProperty "XPrv Base58 roundtrip" $ property $ do
      xprv :: S.XPrv <- forAll genXPrv
      let xprvB = S.encodeXPrv xprv
      Just xprv === S.decodeXPrv xprvB

  , testCase "Pub invalid" $ do
      Nothing @=? K.parsePub "\x03hxxqhlivzpapqhxguwpftplfduateosa"

  , testProperty "subPubPub with normal index" $ property $ do
      S.XPub v d f i c p <- forAll genXPub
      let d' = S.Depth $ case S.unDepth d of { 255 -> 254; w -> w }
          xpub' = S.XPub v d' f i c p
      i1n :: S.Index <- forAll genIndexNormal
      case S.subXPubXPub xpub' i1n of
         Just _ -> pure ()
         Nothing -> do -- This has chance of happening of 1 in 2^127
           i2n :: S.Index <- forAll genIndexNormal
           True === isJust (S.subXPubXPub xpub' i2n)

  , testProperty "subPubPub with hardended index" $ property $ do
      xpub :: S.XPub <- forAll genXPub
      ih :: S.Index <- forAll genIndexHardened
      Nothing === S.subXPubXPub xpub ih

  , testProperty "subsPubPub" $ property $ do
      pubP :: K.Pub <- forAll genPub
      i0 :: S.Index <- forAll genIndexNormal
      chP :: S.Chain <- forAll genChain
      let mindTheGap = do
            (c0, pub0) <- S.subPubPub chP pubP i0
            i1 <- S.indexNext i0
            (c1, pub1) <- S.subPubPub chP pubP i1
            i2 <- S.indexNext i1
            (c2, pub2) <- S.subPubPub chP pubP i2
            pure [(i0, c0, pub0), (i1, c1, pub1), (i2, c2, pub2)]
      case mindTheGap of
        Nothing -> pure () -- We found some gaps. Testing is hard. Skip.
        Just x  -> x === take 3 (S.subsPubPub chP pubP i0)

  , testCase "indexIsHardened" $ do
      False @=? S.indexIsHardened (S.Index 0x00000000)
      False @=? S.indexIsHardened (S.Index 0x00000001)
      False @=? S.indexIsHardened (S.Index 0x7ffffffe)
      False @=? S.indexIsHardened (S.Index 0x7fffffff)

      True  @=? S.indexIsHardened (S.Index 0x80000000)
      True  @=? S.indexIsHardened (S.Index 0x80000001)
      True  @=? S.indexIsHardened (S.Index 0xfffffffe)
      True  @=? S.indexIsHardened (S.Index 0xffffffff)

  , testCase "indexNext" $ do
      S.Index 0x00000000 @=? S.Index minBound
      S.Index 0xffffffff @=? S.Index maxBound
      -- Normal
      Just (S.Index 0x00000001) @=? S.indexNext (S.Index 0x00000000)
      Just (S.Index 0x00000002) @=? S.indexNext (S.Index 0x00000001)
      Just (S.Index 0x7ffffffe) @=? S.indexNext (S.Index 0x7ffffffd)
      Just (S.Index 0x7fffffff) @=? S.indexNext (S.Index 0x7ffffffe)
      -- Normal upper bound
      Nothing @=? S.indexNext (S.Index 0x7fffffff)
      -- Hardened in bounds
      Just (S.Index 0x80000001) @=? S.indexNext (S.Index 0x80000000)
      Just (S.Index 0x80000002) @=? S.indexNext (S.Index 0x80000001)
      Just (S.Index 0x7ffffffe) @=? S.indexNext (S.Index 0x7ffffffd)
      Just (S.Index 0xffffffff) @=? S.indexNext (S.Index 0xfffffffe)
      -- Hardened upper bound
      Nothing @=? S.indexNext (S.Index 0xffffffff)

  , tt_vectors
  ]

tt_vectors :: TestTree
tt_vectors = testGroup "Test vectors"
  [ tv $ TV { tv_name       = "m"
            , tv_depth      = 0
            , tv_fingerp    = 0x00000000
            , tv_index      = 0
            , tv_chain      = "873dff81c02f525623fd1fe5167eac3a55a049de3d314bb42ee227ffed37d508"
            , tv_xpub       = "xpub661MyMwAqRbcFtXgS5sYJABqqG9YLmC4Q1Rdap9gSE8NqtwybGhePY2gZ29ESFjqJoCu1Rupje8YtGqsefD265TMg7usUDFdp6W1EGMcet8"
            , tv_pub        = "0339a36013301597daef41fbe593a02cc513d0b55527ec2df1050e2e8ff49c85c2"
            , tv_xprv       = "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi"
            , tv_prv        = "e8f32e723decf4051aefac8e2c93c9c5b214313817cdb01a1494b917c8436b35"
            , tv_subPubPubs = [ (0, "xpub68Gmy5EVb2BdFbj2LpWrk1M7obNuaPTpT5oh9QCCo5sRfqSHVYWex97WpDZzszdzHzxXDAzPLVSwybe4uPYkSk4G3gnrPqqkV9RyNzAcNJ1")
                              , (1, "xpub68Gmy5EVb2BdHTYHpekwGdcbBWax19w9HwA2DaADYvuCSSgt4YAErxxSN1KWSnmyqkwRNbnTj3XiUBKmHeC8rTjLRPjSULcDKQQgfgJDppq") ]
            , tv_subPrvPrvs = [ (0, "xprv9uHRZZhbkedL37eZEnyrNsQPFZYRAvjy5rt6M1nbEkLSo378x1CQQLo2xxBvREwiK6kqf7GRNvsNEchwibzXaV6i5GcsgyjBeRguXhKsi4R")
                              , (1, "xprv9uHRZZhbkedL4yTpidDvuVfrdUkTbhDHviERRBkbzbNDZeMjWzqzKAdxWhzftGDSxDmBdakjqHiZJbkwiaTEXJdjZAaAjMZEE3PMbMrPJih")
                              , (0 + 0x80000000, "xprv9uHRZZhk6KAJC1avXpDAp4MDc3sQKNxDiPvvkX8Br5ngLNv1TxvUxt4cV1rGL5hj6KCesnDYUhd7oWgT11eZG7XnxHrnYeSvkzY7d2bhkJ7")
                              , (1 + 0x80000000, "xprv9uHRZZhk6KAJFszJGW6LoUFq92uL7FvkBhmYiMurCWPHLJZkX2aGvNdRUBNnJu7nv36WnwCN59uNy6sxLDZvvNSgFz3TCCcKo7iutQzpg78") ]
            }
  , tv $ TV { tv_name       = "m/0"
            , tv_depth      = 1
            , tv_fingerp    = 0x3442193e
            , tv_index      = 0
            , tv_chain      = "d323f1be5af39a2d2f08f5e8f664633849653dbe329802e9847cfc85f8d7b52a"
            , tv_xpub       = "xpub68Gmy5EVb2BdFbj2LpWrk1M7obNuaPTpT5oh9QCCo5sRfqSHVYWex97WpDZzszdzHzxXDAzPLVSwybe4uPYkSk4G3gnrPqqkV9RyNzAcNJ1"
            , tv_pub        = "027c4b09ffb985c298afe7e5813266cbfcb7780b480ac294b0b43dc21f2be3d13c"
            , tv_xprv       = "xprv9uHRZZhbkedL37eZEnyrNsQPFZYRAvjy5rt6M1nbEkLSo378x1CQQLo2xxBvREwiK6kqf7GRNvsNEchwibzXaV6i5GcsgyjBeRguXhKsi4R"
            , tv_prv        = "4e2cdcf2f14e802810e878cf9e6411fc4e712edf19a06bcfcc5d5572e489a3b7"
            , tv_subPubPubs = [ (0, "xpub6AvUGrnEpfvJ8L7GLRkBTByQ9uBvUHp9o5VxHrFxhvzV4dSWkySpNaBoLR9FpbnwRmTa69yLHF3QfcaxbWT7gWdwws5k4dpmJvqpEuMWwnj") ]
            , tv_subPrvPrvs = [ (0, "xprv9ww7sMFLzJMzur2oEQDB642fbsMS4q6JRraMVTrM9bTWBq7NDS8ZpmsKVB4YF3mZecqax1fjnsPF19xnsJNfRp4RSyexacULXMKowSACTRc") ]
            }

  , tv $ TV { tv_name       = "m/0'"
            , tv_depth      = 1
            , tv_fingerp    = 0x3442193e
            , tv_index      = 0 + 0x80000000
            , tv_chain      = "47fdacbd0f1097043b78c63c20c34ef4ed9a111d980047ad16282c7ae6236141"
            , tv_xpub       = "xpub68Gmy5EdvgibQVfPdqkBBCHxA5htiqg55crXYuXoQRKfDBFA1WEjWgP6LHhwBZeNK1VTsfTFUHCdrfp1bgwQ9xv5ski8PX9rL2dZXvgGDnw"
            , tv_pub        = "035a784662a4a20a65bf6aab9ae98a6c068a81c52e4b032c0fb5400c706cfccc56"
            , tv_xprv       = "xprv9uHRZZhk6KAJC1avXpDAp4MDc3sQKNxDiPvvkX8Br5ngLNv1TxvUxt4cV1rGL5hj6KCesnDYUhd7oWgT11eZG7XnxHrnYeSvkzY7d2bhkJ7"
            , tv_prv        = "edb2e14f9ee77d26dd93b4ecede8d16ed408ce149b6cd80b0715a2d911a0afea"
            , tv_subPubPubs = [ (0, "xpub6ASuArnXKPbfEVRpCesNx4P939HDXENHkksgxsVG1yNp9958A33qYoPiTN9QrJmWFa2jNLdK84bWmyqTSPGtApP8P7nHUYwxHPhqmzUyeFG") ]
            , tv_subPrvPrvs = [ (0, "xprv9wTYmMFdV23N21MM6dLNavSQV7Sj7meSPXx6AV5eTdqqGLjycVjb115Ec5LgRAXscPZgy5G4jQ9csyyZLN3PZLxoM1h3BoPuEJzsgeypdKj")
                              , (0 + 0x80000000, "xprv9wTYmMFmpgaLB5Hge4YtaGqCKpsYPTD9vXWSsmdZrNU3Y2i4WoBykm6ZteeCLCCZpGxdHQuqEhM6Gdo2X6CVrQiTw6AAneF9WSkA9ewaxtS") ]
            }
  , tv $ TV { tv_name       = "m/44'/0'/0'/0"
            , tv_depth      = 4
            , tv_fingerp    = 0x31bf9083
            , tv_index      = 0
            , tv_chain      = "6dd51cfd6f41ad9d1edf007727a683a2cd317524e43bf30b5cd1a2035c1041cf"
            , tv_xpub       = "xpub6Du7UxgPt9xyZsijCkstyy1MEuR6SZbAh3MaaE1yvKesdKSygQKqfYLzVxdhoPeRwwTLwpEnjzMqVMb5NYvazx56sxNCRoExGNY1VNMqsSD"
            , tv_pub        = "03073cdc669834931577d21319e4414523ce32a638b4ae42eb4ace8ed3dac683b5"
            , tv_xprv       = "xprv9zum5T9W3nQgMPeG6jLtcq4cgsac36sKKpRymqcNMz7tkX7q8s1b7k2WegNtWCo91gRjnSZANnsMGLhjsofsQfVtuSLQdfeH48gDNosBAHk"
            , tv_prv        = "65e96c73a30bcb815fc0f0a5694a0ead061034eba030d1b82991dff1f9b68519"
            , tv_subPubPubs = [ (0, "xpub6FqQEmsqnmQYh8vkY8NsHUygAjSbfXCCXFzWYe2GHG3mL4ng9Ned6kEKHN4LpNQW3CwZAsnfgoj7VcqrbXoXiAUXz1bZdsam9QzP9RWjR2u") ]
            , tv_subPrvPrvs = [ (0, "xprvA2r3qGLwxPrFUerHS6qrvM2wchc7G4UMA34ukFceivWnTGTXbqLNYwuqS8nkbm7SjvgNSiHumkYLur27kjdnGZexAtiC6Ha4SNr7Wpeoz2V") ]
            }
  ]

data TV = TV
 { tv_name       :: String
 , tv_depth      :: Word8
 , tv_fingerp    :: Word32
 , tv_index      :: Word32
 , tv_chain      :: B.ByteString -- ^ Expected chain code in lower-case hexadecimal
 , tv_xpub       :: B.ByteString -- ^ Base58-serialized XPub
 , tv_pub        :: B.ByteString -- ^ Expected public key in lower-case hexadecimal
 , tv_xprv       :: B.ByteString -- ^ Base58-serialized XPrv
 , tv_prv        :: B.ByteString -- ^ Expected private key in lower-case hexadecimal
 , tv_subPubPubs :: [(Word32, B.ByteString)] -- ^ Pub->Pub Base58 subkeys with index
 , tv_subPrvPrvs :: [(Word32, B.ByteString)] -- ^ Prv->Prv Base58 subkeys with index
 }


tv :: TV -> TestTree
tv x = testGroup (tv_name x)
  [ testCase "chain" $ do
      Just (tv_chain x) @=? fmap (toBase16 . S.unChain)
                                 (S.chain (fromBase16 (tv_chain x)))

  , testCase "prv" $ do
      Just (tv_prv x) @=? fmap (toBase16 . K.prvRaw)
                               (K.parsePrv (fromBase16 (tv_prv x)))

  , testCase "pub" $ do
      Just (tv_pub x) @=? fmap (toBase16 . K.pubCompressed)
                               (K.parsePub (fromBase16 (tv_pub x)))

  , testCase "xprvToXPub" $ do
      let Just xprv = S.decodeXPrv (tv_xprv x)
          Just xpub = S.decodeXPub (tv_xpub x)
      xpub @=? S.xprvToXPub S.version_xpub xprv

  , testCase "xprv" $ do
      Just xprv@(S.XPrv v d f i c k) <- pure $ S.decodeXPrv (tv_xprv x)
      S.encodeXPrv xprv @?= tv_xprv x
      v @?= S.version_xprv
      d @?= S.Depth (tv_depth x)
      f @?= S.Fingerprint (tv_fingerp x)
      i @?= S.Index (tv_index x)
      Just c @?= S.chain (fromBase16 ((tv_chain x)))
      Just k @?= K.parsePrv (fromBase16 (tv_prv x))

  , testCase "xpub" $ do
      Just xpub@(S.XPub v d f i c k) <- pure $ S.decodeXPub (tv_xpub x)
      S.encodeXPub xpub @?= tv_xpub x
      v @?= S.version_xpub
      d @?= S.Depth (tv_depth x)
      f @?= S.Fingerprint (tv_fingerp x)
      i @?= S.Index (tv_index x)
      Just c @?= S.chain (fromBase16 ((tv_chain x)))
      Just k @?= K.parsePub (fromBase16 (tv_pub x))

  , testCase "subPubPub" $ do
      let Just xpub0@(S.XPub v0 d0 _ _ _ _) = S.decodeXPub (tv_xpub x)
      for_ (tv_subPubPubs x) $ \(i1w, xpub1B) -> do
        let i1 = S.Index i1w
            Just xpub1@(S.XPub v1 d1 _ i1' _ _) = S.subXPubXPub xpub0 i1
        v1 @?= v0
        d1 @?= S.Depth (1 + S.unDepth d0)
        i1' @?= i1
        S.encodeXPub xpub1 @?= xpub1B

  , testCase "subsPubPub" $ do
      let Just xpub0@(S.XPub v0 d0 _ _ _ _) = S.decodeXPub (tv_xpub x)
      case filter (not . S.indexIsHardened . S.Index . fst) (tv_subPrvPrvs x) of
        [] -> pure ()
        xxx -> do
          let i0 = S.Index (minimum (fmap fst xxx))
              subs = take 10 $ S.subsXPubXPub xpub0 i0
          for_ (tv_subPubPubs x) $ \(i1w, xpub1B) -> do
            let i1 = S.Index i1w
                Just xpub1@(S.XPub v1 d1 _ i1' _ _) =
                  find ((i1 ==) . S.xpub_index) subs
            v1 @?= v0
            d1 @?= S.Depth (1 + S.unDepth d0)
            i1' @?= i1
            S.encodeXPub xpub1 @?= xpub1B

  , testCase "subsPrvPrv: hardened" $ do
      let Just xprv0@(S.XPrv v0 d0 _ _ _ _) = S.decodeXPrv (tv_xprv x)
      case filter (S.indexIsHardened . S.Index . fst) (tv_subPrvPrvs x) of
        [] -> pure ()
        xxx -> do
          let i0 = S.Index (minimum (fmap fst xxx))
              subs = take 10 $ S.subsXPrvXPrv xprv0 i0
          for_ xxx $ \(i1w, xprv1B) -> do
            let i1 = S.Index i1w
                Just xprv1@(S.XPrv v1 d1 _ i1' _ _) =
                  find ((i1 ==) . S.xprv_index) subs
            v1 @?= v0
            d1 @?= S.Depth (1 + S.unDepth d0)
            i1' @?= i1
            S.encodeXPrv xprv1 @?= xprv1B

  , testCase "subsPrvPrv: normal" $ do
      let Just xprv0@(S.XPrv v0 d0 _ _ _ _) = S.decodeXPrv (tv_xprv x)
      case filter (not . S.indexIsHardened . S.Index . fst) (tv_subPrvPrvs x) of
        [] -> pure ()
        xxx -> do
          let i0 = S.Index (minimum (fmap fst xxx))
              subs = take 10 $ S.subsXPrvXPrv xprv0 i0
          for_ xxx $ \(i1w, xprv1B) -> do
            let i1 = S.Index i1w
                Just xprv1@(S.XPrv v1 d1 _ i1' _ _) =
                  find ((i1 ==) . S.xprv_index) subs
            v1 @?= v0
            d1 @?= S.Depth (1 + S.unDepth d0)
            i1' @?= i1
            S.encodeXPrv xprv1 @?= xprv1B

  , testCase "subPrvPrv" $ do
      let Just xprv0@(S.XPrv v0 d0 _ _ _ _) = S.decodeXPrv (tv_xprv x)
      for_ (tv_subPrvPrvs x) $ \(i1w, xprv1B) -> do
        let i1 = S.Index i1w
            Just xprv1@(S.XPrv v1 d1 _ i1' _ _) = S.subXPrvXPrv xprv0 i1
        v1 @?= v0
        d1 @?= S.Depth (1 + S.unDepth d0)
        i1' @?= i1
        S.encodeXPrv xprv1 @?= xprv1B
  ]

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

genIndex :: MonadGen m => m S.Index
genIndex = S.Index <$> Gen.word32 Range.constantBounded

genIndexNormal :: MonadGen m => m S.Index
genIndexNormal = S.Index <$> Gen.word32 (Range.constant 0 (0x80000000 - 1))

genIndexHardened :: MonadGen m => m S.Index
genIndexHardened = S.Index <$> Gen.word32 (Range.constant 0x80000000 maxBound)

genXPub :: MonadGen m => m S.XPub
genXPub = S.XPub <$> genVersion
                 <*> genDepth <*> genFingerprint
                 <*> genIndex <*> genChain <*> genPub

genXPrv :: MonadGen m => m S.XPrv
genXPrv = S.XPrv <$> genVersion
                 <*> genDepth <*> genFingerprint
                 <*> genIndex <*> genChain <*> genPrv

genVersion :: MonadGen m => m S.Version
genVersion = S.Version <$> Gen.word32 Range.constantBounded

genFingerprint :: MonadGen m => m S.Fingerprint
genFingerprint = S.Fingerprint <$> Gen.word32 Range.constantBounded

genDepth :: MonadGen m => m S.Depth
genDepth = S.Depth <$> Gen.word8 Range.constantBounded

genChain :: MonadGen m => m S.Chain
genChain = do
  b <- Gen.bytes (Range.singleton 32)
  let Just c = S.chain b
  pure c

genPrv :: MonadGen m => m K.Prv
genPrv = do
  b <- Gen.bytes (Range.singleton 32)
  let Just k = K.parsePrv b
  pure k

genPub :: MonadGen m => m K.Pub
genPub = go 10000 where
  go 0 = error "genPub: too many attempts"
  go n = do
    h <- Gen.element [2, 3 :: Word8]
    b <- Gen.bytes (Range.singleton 32)
    case K.parsePub (B.cons h b) of
      Just k -> pure k
      Nothing -> go (n - 1)

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

toBase16 :: B.ByteString -> B.ByteString
toBase16 = B16.encode

fromBase16 :: B.ByteString -> B.ByteString
fromBase16 a = case B16.decode a of
  (b, "") -> b
  _ -> error ("Invalid base16 string: " <> show a)