{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wno-type-defaults #-}

module Main (main) where

import qualified Data.ByteString as B
import qualified Data.ByteString.Base16 as B16
import qualified Data.Text.Encoding as T
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, Range, property, forAll, (===))
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import qualified SLIP32 as S

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

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

tt :: TestTree
tt = testGroup "SLIP32"
  [ tt_vectors

  , testProperty "roundtrip" $ property $ do
      path <- forAll $ genPath Range.constantBounded
      chain <- forAll genChain

      prv <- forAll genPrv
      let xprv = S.XPrv path chain prv
          xprvB = S.renderXPrv xprv
          xprvT = S.renderXPrvText xprv
      xprvB === T.encodeUtf8 xprvT
      xprvT === T.decodeUtf8 xprvB
      Just xprv === S.parseXPrv xprvB
      Just xprv === S.parseXPrvText xprvT
      Just (Right xprv) === S.parse xprvB
      Just (Right xprv) === S.parseText xprvT

      pub <- forAll genPub
      let xpub = S.XPub path chain pub
          xpubB = S.renderXPub xpub
          xpubT = S.renderXPubText xpub
      xpubB === T.encodeUtf8 xpubT
      xpubT === T.decodeUtf8 xpubB
      Just xpub === S.parseXPub xpubB
      Just xpub === S.parseXPubText xpubT
      Just (Left xpub) === S.parse xpubB
      Just (Left xpub) === S.parseText xpubT
  ]

tt_vectors :: TestTree
tt_vectors = testGroup "Test vectors"
  [ tv $ TV { tv_pathD = "m"
            , tv_path  = []
            , tv_pathR = ""
            , tv_chain = "7923408dadd3c7b56eed15567707ae5e5dca089de972e07f3b860450e2a3b70e"
            , tv_xprv  = "xprv1qpujxsyd4hfu0dtwa524vac84e09mjsgnh5h9crl8wrqg58z5wmsuqqcxlqmar3fjhkprndzkpnp2xlze76g4hu7g7c4r4r2m2e6y8xlvu566tn6"
            , tv_prv   = "001837c1be8e2995ec11cda2b066151be2cfb48adf9e47b151d46adab3a21cdf67"
            , tv_xpub  = "xpub1qpujxsyd4hfu0dtwa524vac84e09mjsgnh5h9crl8wrqg58z5wmsuq7eqte474swq3cvvvcncumfz6xe6l0j6jdl990an7mukyyuemsyjszuwypl"
            , tv_pub   = "03d902f35f560e0470c63313c7369168d9d7df2d49bf295fd9fb7cb109ccee0494"
            }
  , tv $ TV { tv_pathD = "m/0"
            , tv_path  = [0]
            , tv_pathR = ""
            , tv_chain = "e0e6503ac057cf5dc76e0735e56dd44d193b2e9e271cc2d46bc759c99b021e3c"
            , tv_xprv  = "xprv1qyqqqqqqurn9qwkq2l84m3mwqu672mw5f5vnkt57yuwv94rtcavunxczrc7qpw4gn29a6cw9ug4e7yrqrkrerj0cl39jlfkln45dxdhsavpmqm4krfqykk"
            , tv_prv   = "00baa89a8bdd61c5e22b9f10601d8791c9f8fc4b2fa6df9d68d336f0eb03b06eb6"
            , tv_xpub  = "xpub1qyqqqqqqurn9qwkq2l84m3mwqu672mw5f5vnkt57yuwv94rtcavunxczrc7qxa4l2v75k923p75lgyjtdeyxzmc8m6709mcvlvv9ehz22aj9pdr4m6lwmk"
            , tv_pub   = "0376bf533d4b15510fa9f4124b6e48616f07debcf2ef0cfb185cdc4a576450b475"
            }

  , tv $ TV { tv_pathD = "m/1"
            , tv_path  = [1]
            , tv_pathR = "00000001"
            , tv_chain = "5c48917d6838b666aeb11eac7c4f98f807779b57c7522e38509719eeb1e7a592"
            , tv_prv   = "00c1beaff0c4db984670a40c69c2947b9d33cd7f6e749c67e1fcb5c6118dda1282"
            , tv_pub   = "02ea2649b3512b9a859ab658a85e2989a7ae39b2518877b2dc0f2b44b785d5788d"
            , tv_xprv  = "xprv1qyqqqqqpt3yfzltg8zmxdt43r6k8cnuclqrh0x6hcafzuwzsjuv7av085kfqpsd74lcvfkucgec2grrfc228h8fne4lkuayuvlsledwxzxxa5y5zefalyg"
            , tv_xpub  = "xpub1qyqqqqqpt3yfzltg8zmxdt43r6k8cnuclqrh0x6hcafzuwzsjuv7av085kfq963xfxe4z2u6skdtvk9gtc5cnfaw8xe9rzrhktwq726yk7za27ydw88adn"
            }
  , tv $ TV { tv_pathD = "m/0'"
            , tv_path  = [0 + 2^31]
            , tv_pathR = "80000000"
            , tv_chain = "f1c03f5ff97108912fd56761d3fada8879e4173aba45f10da4bbd94b1c497160"
            , tv_prv   = "00c08cf331996482c06db3d259ff99be4bf7083824d53185e33191ee7ceb2bf96f"
            , tv_pub   = "027f1d87730e460e921b382242911565bf93daf2081ed685b2edd1d01176b2c13c"
            , tv_xprv  = "xprv1qxqqqqqq78qr7hlewyyfzt74vasa87k63pu7g9e6hfzlzrdyh0v5k8zfw9sqpsyv7vcejeyzcpkm85jel7vmujlhpquzf4f3sh3nry0w0n4jh7t0jhc039"
            , tv_xpub  = "xpub1qxqqqqqq78qr7hlewyyfzt74vasa87k63pu7g9e6hfzlzrdyh0v5k8zfw9sqylcasaesu3swjgdnsgjzjy2kt0unmteqs8kkskewm5wsz9mt9sfuvlxj6p"
            }
  , tv $ TV { tv_pathD = "m/1'"
            , tv_path  = [1 + 2^31]
            , tv_pathR = "80000001"
            , tv_chain = "43cc4bca59c666a5f79265148125802ed2cec46df1c5ca8e6a058dab525a73f1"
            , tv_prv   = "003ef02fc53000742891fc90458ba9edc8363d8f1f267e326b1078710c7db34de5"
            , tv_pub   = "03b5184a526dac6abda3d8d54a541471ce83e8c2260d56706053e2780922319f5e"
            , tv_xprv  = "xprv1qxqqqqqpg0xyhjjecen2taujv52gzfvq9mfva3rd78zu4rn2qkx6k5j6w0csq0hs9lznqqr59zgleyz93w57mjpk8k837fn7xf43q7r3p37mxn095hysnx"
            , tv_xpub  = "xpub1qxqqqqqpg0xyhjjecen2taujv52gzfvq9mfva3rd78zu4rn2qkx6k5j6w0cs8dgcfffxmtr2hk3a34222s28rn5rarpzvr2kwps98cncpy3rr867k5u83k"
            }
  , tv $ TV { tv_pathD = "m/44'/0'/0'"
            , tv_path  = [44 + 2^31, 0 + 2^31, 0 + 2^31]
            , tv_pathR = "8000002c8000000080000000"
            , tv_chain = "3da4bc190a2680111d31fadfdc905f2a7f6ce77c6f109919116f253d43445219"
            , tv_prv   = "00fe64af825b5b78554c33a28b23085fc082f691b3c712cc1d4e66e133297da87a"
            , tv_pub   = "03774c910fcf07fa96886ea794f0d5caed9afe30b44b83f7e213bb92930e7df4bd"
            , tv_xprv  = "xprv1qwqqqqpvsqqqqqyqqqqqq0dyhsvs5f5qzywnr7klmjg972nldnnhcmcsnyv3zme984p5g5seqrlxftuztddhs42vxw3gkgcgtlqg9a53k0r39nqafenwzvef0k585enml6g"
            , tv_xpub  = "xpub1qwqqqqpvsqqqqqyqqqqqq0dyhsvs5f5qzywnr7klmjg972nldnnhcmcsnyv3zme984p5g5seqdm5eyg0eurl495gd6nefux4etke4l3sk39c8alzzwae9ycw0h6t6ltmssr"
            }
  , tv $ TV { tv_pathD = "m/44'/0'/1'"
            , tv_path  = [44 + 2^31, 0 + 2^31, 1 + 2^31]
            , tv_pathR = "8000002c8000000080000001"
            , tv_chain = "2971fa2db0ff5d69e166a406813aa3d9ed09c4adac2e0ce33523da8c5609f4f4"
            , tv_prv   = "008855dfda37fe663bffc0136618504e3cbd7d992134609cef6191c729339d5c65"
            , tv_pub   = "025d0261853d4c3a379160fb51d2f262ac64e65219139982c4e2180bcef1a233d9"
            , tv_xprv  = "xprv1qwqqqqpvsqqqqqyqqqqqz2t3lgkmpl6ad8skdfqxsya28k0dp8z2mtpwpn3n2g7633tqna85qzy9th76xllxvwllcqfkvxzsfc7t6lveyy6xp880vxguw2fnn4wx2mhtjy8"
            , tv_xpub  = "xpub1qwqqqqpvsqqqqqyqqqqqz2t3lgkmpl6ad8skdfqxsya28k0dp8z2mtpwpn3n2g7633tqna85qfwsycv984xr5du3vra4r5hjv2kxfejjryfenqkyugvqhnh35geajlgxhp0"
            }
  , tv $ TV { tv_pathD = "m/44'/2'/0'"
            , tv_path  = [44 + 2^31, 2 + 2^31, 0 + 2^31]
            , tv_pathR = "8000002c8000000280000000"
            , tv_chain = "869c5045e5fc789646babcd1961b101bc31e75fe50df8a585c79b05dca0ac758"
            , tv_prv   = "00983cd10d8d14160b10b9a4bb63207e9585054a3133619d57b78ea9d5aa3046d2"
            , tv_pub   = "0340fe3b8e89165258bac0cb711613c618d1af63dc321a90b751d0697301441bcc"
            , tv_xprv  = "xprv1qwqqqqpvsqqqqq5qqqqqpp5u2pz7tlrcjert40x3jcd3qx7rre6lu5xl3fv9c7dsth9q436cqzvre5gd352pvzcshxjtkceq062c2p22xyekr82hk782n4d2xprdysp4gxc"
            , tv_xpub  = "xpub1qwqqqqpvsqqqqq5qqqqqpp5u2pz7tlrcjert40x3jcd3qx7rre6lu5xl3fv9c7dsth9q436cqdq0uwuw3yt9yk96cr9hz9snccvdrtmrmsep4y9h28gxjucpgsducuj4f9r"
            }
  , tv $ TV { tv_pathD = "m/49'/0'/0'"
            , tv_path  = [49 + 2^31, 0 + 2^31, 0 + 2^31]
            , tv_pathR = "800000318000000080000000"
            , tv_chain = "6eaae365ae0e0a0aab84325cfe7cd76c3b909035f889e7d3f1b847a9a0797ecb"
            , tv_prv   = "00880d51752bda4190607e079588d3f644d96bfa03446bce93cddfda3c4a99c7e6"
            , tv_pub   = "02f1f347891b20f7568eae3ec9869fbfb67bcab6f358326f10ecc42356bd55939d"
            , tv_xprv  = "xprv1qwqqqqp3sqqqqqyqqqqqqm42udj6urs2p24cgvjule7dwmpmjzgrt7yfulflrwz84xs8jlktqzyq65t490dyryrq0cretzxn7ezdj6l6qdzxhn5neh0a50z2n8r7vumvllf"
            , tv_xpub  = "xpub1qwqqqqp3sqqqqqyqqqqqqm42udj6urs2p24cgvjule7dwmpmjzgrt7yfulflrwz84xs8jlktqtclx3ufrvs0w45w4clvnp5lh7m8hj4k7dvrymcsanzzx44a2kfe6xynfgh"
            }
  , tv $ TV { tv_pathD = "m/49'/2'/0'"
            , tv_path  = [49 + 2^31, 2 + 2^31, 0 + 2^31]
            , tv_pathR = "800000318000000280000000"
            , tv_chain = "67b7e1dc5c70a93504218ccf40c47ad46d4a9c858196376ce0e853aca7be0498"
            , tv_prv   = "00cf222cc2e097049fe2ca76626c19c7e7a3ef971b1f64195758ab3c832463fcf4"
            , tv_pub   = "02b07388bd2edaba3c0a2c0856716fd7c9965d212fb2736f7b925f57d922b10ace"
            , tv_xprv  = "xprv1qwqqqqp3sqqqqq5qqqqqqeahu8w9cu9fx5zzrrx0grz844rdf2wgtqvkxakwp6zn4jnmupycqr8jytxzuztsf8lzefmxymqecln68muhrv0kgx2htz4neqeyv070gg6dcn7"
            , tv_xpub  = "xpub1qwqqqqp3sqqqqq5qqqqqqeahu8w9cu9fx5zzrrx0grz844rdf2wgtqvkxakwp6zn4jnmupycq2c88z9a9mdt50q29sy9vut06lyevhfp97e8xmmmjf040kfzky9vu2pu92u"
            }
  , tv $ TV { tv_pathD = "m/84'/0'/0'"
            , tv_path  = [84 + 2^31, 0 + 2^31, 0 + 2^31]
            , tv_pathR = "800000548000000080000000"
            , tv_chain = "4a53a0ab21b9dc95869c4e92a161194e03c0ef3ff5014ac692f433c4765490fc"
            , tv_prv   = "00e14f274d16ca0d91031b98b162618061d03930fa381af6d4caf44b01819ab6d4"
            , tv_pub   = "02707a62fdacc26ea9b63b1c197906f56ee0180d0bcf1966e1a2da34f5f3a09a9b"
            , tv_xprv  = "xprv1qwqqqqz5sqqqqqyqqqqqqjjn5z4jrwwujkrfcn5j59s3jnsrcrhnlagpftrf9apnc3m9fy8uqrs57f6dzm9qmygrrwvtzcnpspsaqwfslgup4ak5et6ykqvpn2mdggeaxrp"
            , tv_xpub  = "xpub1qwqqqqz5sqqqqqyqqqqqqjjn5z4jrwwujkrfcn5j59s3jnsrcrhnlagpftrf9apnc3m9fy8uqfc85cha4npxa2dk8vwpj7gx74hwqxqdp083jehp5tdrfa0n5zdfkg3lp00"
            }
  ]

data TV = TV
 { tv_pathD   :: String
 , tv_path    :: [Word32]     -- ^ Derivation path.
 , tv_pathR   :: B.ByteString -- ^ Derivation path in lower-case hexadecimal.
 , tv_chain   :: B.ByteString -- ^ Expected chain code in lower-case hexadecimal.
 , tv_xprv    :: B.ByteString -- ^ Raw SLIP-0032 extended private key.
 , tv_prv     :: B.ByteString -- ^ Expected private key in lower-case hexadecimal.
 , tv_xpub    :: B.ByteString -- ^ Raw SLIP-0032 extended public key.
 , tv_pub     :: B.ByteString -- ^ Expected public key in lower-case hexadecimal.
 }

tv :: TV -> TestTree
tv x = testGroup (tv_pathD x)
  [ testCase "path" $ do
      fmap S.unPath (S.path (tv_path x)) @?= Just (tv_path 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 . S.unPrv)
                 (S.prv (fromBase16 (tv_prv x)))

  , testCase "pub" $ do
      Just (tv_pub x)
        @=? fmap (toBase16 . S.unPub)
                 (S.pub (fromBase16 (tv_pub x)))

  , testCase "xprv" $ do
      Just xprv@(S.XPrv p c k) <- pure $ S.parseXPrv (tv_xprv x)
      S.parse (tv_xprv x) @?= Just (Right xprv)
      S.unPath p @?= tv_path x
      toBase16 (S.unChain c) @?= tv_chain x
      toBase16 (S.unPrv k) @?= tv_prv x
      S.renderXPrv xprv @?= tv_xprv x

  , testCase "xpub" $ do
      Just xpub@(S.XPub p c k) <- pure $ S.parseXPub (tv_xpub x)
      S.parse (tv_xpub x) @?= Just (Left xpub)
      S.unPath p @?= tv_path x
      toBase16 (S.unChain c) @?= tv_chain x
      toBase16 (S.unPub k) @?= tv_pub x
      S.renderXPub xpub @?= tv_xpub x
  ]

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

genPath :: MonadGen m => Range Word8 -> m S.Path
genPath r = do
  ws <- Gen.list (fmap fromIntegral r) (Gen.word32 (Range.constantBounded))
  let Just p = S.path ws
  pure p

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

genPub :: MonadGen m => m S.Pub
genPub = do
  h <- Gen.element [2, 3 :: Word8]
  b <- Gen.bytes (Range.singleton 32)
  let Just k = S.pub (B.cons h b)
  pure k

genPrv :: MonadGen m => m S.Prv
genPrv = do
  b <- Gen.bytes (Range.singleton 32)
  let Just k = S.prv (B.cons 0 b)
  pure k

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

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"