{-# 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.Bitcoin.Script as S
import qualified Data.ByteString as B
import qualified Data.ByteString.Base16 as B16
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, (===))
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import qualified Bitcoin.Address          as A
import qualified Bitcoin.Address.Hash     as AH
import qualified Bitcoin.Address.Internal as AI
import qualified Bitcoin.Address.Script   as ASC
import qualified Bitcoin.Address.SegWit   as ASW
import qualified Bitcoin.Address.Settings as ASE

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

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

tt :: TestTree
tt = testGroup "bitcoin-address"
  [ tt_settings
  , tt_script
  , tt_hash
  , tt_address
  ]

tt_settings :: TestTree
tt_settings = testGroup "Settings"
  [ testCase "btc" $ do
      let ASE.Settings preP2PKH preP2SH preSegWit = A.btc
      preP2PKH @?= ASE.PrefixP2PKH { ASE.unPrefixP2PKH = 0x00 }
      preP2SH @?= ASE.PrefixP2SH { ASE.unPrefixP2SH = 0x05 }
      Just preSegWit @?= ASE.prefixSegWit "bc"

  , testCase "btcTestnet" $ do
      let ASE.Settings preP2PKH preP2SH preSegWit = A.btcTestnet
      preP2PKH @?= ASE.PrefixP2PKH { ASE.unPrefixP2PKH = 0x6f }
      preP2SH @?= ASE.PrefixP2SH { ASE.unPrefixP2SH = 0xc4 }
      Just preSegWit @?= ASE.prefixSegWit "tb"
  ]

tt_hash :: TestTree
tt_hash = testGroup "Hash"
  [ testGroup "pubHash160"
    [ testCase "pk1" $ do
        toBase16 (AH.unPubHash160 pk1hash160)
          @?= "5f04d673fdbd072b97558a8f830a5ddf655839d7"

    , testCase "pk2" $ do
        toBase16 (AH.unPubHash160 pk2hash160)
          @?= "980938e055943e47b73e5d16b2de7343a61a0a7b"
    ]
  , testGroup "pubUncompressedHash160"
    [ testCase "pk1" $ do
        toBase16 (AH.unPubHash160 (AH.pubUncompressedHash160 pk1))
          @?= "6baec369d3d75a0a27818c3560fe2c4db8f916d6"

    , testCase "pk2" $ do
        toBase16 (AH.unPubHash160 (AH.pubUncompressedHash160 pk2))
          @?= "85b679742f15e5b24849747c8eab44c0eb508fbe"
    ]
  , testCase "scriptHash160" $ do
      toBase16 (AH.unScriptHash160 (AH.scriptHash160 (ASC.p2pkh pk1hash160)))
        @?= "929f8c79af58e615ffe465a76f44c02865b5b77b"

  , testCase "scriptSHA256" $ do
      toBase16 (AH.unScriptSHA256 (AH.scriptSHA256 (ASC.p2pkh pk1hash160)))
        @?= "a39b82cbb33b2ee51d0fd5e415ce924f3f4a120ae962c13ab2f9b6433326e0bf"
  ]

tt_script :: TestTree
tt_script = testGroup "Script"
  [ tt_script_p2pkh
  , tt_script_multiSig
  , tt_script_p2sh
  , tt_script_segWit
  ]

tt_script_p2pkh :: TestTree
tt_script_p2pkh = testGroup "P2PKH"
  [ testCase "pk1" $ do
      let s = ASC.p2pkh pk1hash160
      toBase16 (AI.scriptBytes s) @?=
        "76a9145f04d673fdbd072b97558a8f830a5ddf655839d788ac"

  , testCase "pk2" $ do
      let s = ASC.p2pkh pk2hash160
      toBase16 (AI.scriptBytes s) @?=
        "76a914980938e055943e47b73e5d16b2de7343a61a0a7b88ac"

  , testProperty "prop" $ property $ do
      pk <- forAll genPub
      let pkh = AH.pubHash160 pk
          s = ASC.p2pkh pkh
      toBase16 (AI.scriptBytes s) === mconcat
        [ "76a914", toBase16 (AH.unPubHash160 pkh), "88ac" ]
  ]


tt_script_segWit :: TestTree
tt_script_segWit = testGroup "SegWit"
  [ tt_script_segWit_p2wpkh
  , tt_script_segWit_p2wsh
  ]

tt_script_segWit_p2wpkh :: TestTree
tt_script_segWit_p2wpkh = testGroup "P2WPKH"
  [ testCase "pk1" $ do
      let s = ASC.segWit (ASW.p2wpkh pk1hash160)
      toBase16 (AI.scriptBytes s) @?=
        "00145f04d673fdbd072b97558a8f830a5ddf655839d7"

  , testCase "pk2" $ do
      let s = ASC.segWit (ASW.p2wpkh pk2hash160)
      toBase16 (AI.scriptBytes s) @?=
        "0014980938e055943e47b73e5d16b2de7343a61a0a7b"

  , testProperty "prop" $ property $ do
      pk <- forAll genPub
      let pkh = AH.pubHash160 pk
          s = ASC.segWit (ASW.p2wpkh pkh)
      toBase16 (AI.scriptBytes s) === mconcat
        [ "0014", toBase16 (AH.unPubHash160 pkh) ]
  ]

tt_script_multiSig :: TestTree
tt_script_multiSig = testGroup "MultiSig"
  [ testCase "1 of [pk1]" $ do
      let Just s = ASC.multiSig [pk1] 1
      toBase16 (AI.scriptBytes s) @?= mconcat
         ["5121", toBase16 (K.pubCompressed pk1), "51ae"]

  ,  testCase "1 of [pk2]" $ do
      let Just s = ASC.multiSig [pk2] 1
      toBase16 (AI.scriptBytes s) @?= mconcat
         ["5121", toBase16 (K.pubCompressed pk2), "51ae"]

  ,  testCase "1 of [pk1, pk2]" $ do
      let Just s = ASC.multiSig [pk1, pk2] 1
      toBase16 (AI.scriptBytes s) @?= mconcat
         [ "5121", toBase16 (K.pubCompressed pk1)
         , "21", toBase16 (K.pubCompressed pk2)
         , "52ae" ]

  ,  testCase "1 of [pk2, pk1]" $ do
      let Just s = ASC.multiSig [pk2, pk1] 1
      toBase16 (AI.scriptBytes s) @?= mconcat
         [ "5121", toBase16 (K.pubCompressed pk2)
         , "21", toBase16 (K.pubCompressed pk1)
         , "52ae" ]

  ,  testCase "2 of [pk1, pk2]" $ do
      let Just s = ASC.multiSig [pk1, pk2] 2
      toBase16 (AI.scriptBytes s) @?= mconcat
         [ "5221", toBase16 (K.pubCompressed pk1)
         , "21", toBase16 (K.pubCompressed pk2)
         , "52ae" ]

  ,  testCase "2 of [pk2, pk1]" $ do
      let Just s = ASC.multiSig [pk2, pk1] 2
      toBase16 (AI.scriptBytes s) @?= mconcat
         [ "5221", toBase16 (K.pubCompressed pk2)
         , "21", toBase16 (K.pubCompressed pk1)
         , "52ae"]

  , testProperty "0 of n (invalid)" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 0 20) genPub
      Nothing === ASC.multiSig pks 0

  , testProperty "m of [] (invalid)" $ property $ do
      n <- forAll $ Gen.int (Range.constant (-20) 20)
      Nothing === ASC.multiSig [] n

  , testProperty "m of n, m > n (invalid)" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 0 20) genPub
      i <- forAll $ Gen.int (Range.constant 1 20)
      Nothing === ASC.multiSig pks (length pks + i)

  , testProperty "m of n, n > 16 (invalid)" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 17 30) genPub
      m <- forAll $ Gen.int (Range.constant 1 20)
      Nothing === ASC.multiSig pks m

  , testProperty "m of n (correct)" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 1 16) genPub
      m <- forAll $ Gen.int (Range.constant 1 (length pks))
      let Just mOp = AI.op0to16 m
          Just nOp = AI.op0to16 (length pks)
          pksOps = map (\pk -> S.OP_PUSHDATA (K.pubCompressed pk) S.OPCODE) pks
          s = S.Script $ (mOp : pksOps) <> [nOp, S.OP_CHECKMULTISIG]
      Just s === ASC.multiSig pks m
  ]

tt_script_p2sh :: TestTree
tt_script_p2sh = testGroup "P2SH"
  [ tt_script_p2sh_p2wpkh
  , tt_script_p2sh_multiSig
  , tt_script_p2sh_p2wsh
  ]

tt_script_p2sh_p2wpkh :: TestTree
tt_script_p2sh_p2wpkh = testGroup "P2WPKH-in-P2SH"
  [ testCase "pk1" $ do
      let sw = ASC.segWit (ASW.p2wpkh pk1hash160)
          swh = AH.scriptHash160 sw
          s = ASC.p2sh swh
      toBase16 (AI.scriptBytes s) @?= mconcat
        [ "a914", toBase16 (AH.unScriptHash160 swh), "87" ]
      toBase16 (AI.scriptBytes s) @?=
        "a914ba9102a5c6407308a600e114cb3474bb179b1eca87"

  , testCase "pk2" $ do
      let sw = ASC.segWit (ASW.p2wpkh pk2hash160)
          swh = AH.scriptHash160 sw
          s = ASC.p2sh swh
      toBase16 (AI.scriptBytes s) @?= mconcat
        [ "a914", toBase16 (AH.unScriptHash160 swh), "87" ]
      toBase16 (AI.scriptBytes s) @?=
        "a9140e8ee6d7778e0b708844a8c40b1f1c700de2489287"

  , testProperty "prop" $ property $ do
      pk <- forAll genPub
      let pkh = AH.pubHash160 pk
          sw = ASC.segWit (ASW.p2wpkh pkh)
          swh = AH.scriptHash160 sw
          s = ASC.p2sh swh
      toBase16 (AI.scriptBytes s) === mconcat
        [ "a914", toBase16 (AH.unScriptHash160 swh), "87" ]
  ]

tt_script_p2sh_multiSig :: TestTree
tt_script_p2sh_multiSig = testGroup "MultiSig"
  [ testCase "1 of [pk1, pk2]" $ do
      let Just sMultiSig = ASC.multiSig [pk1, pk2] 1
          shMultiSig = AH.scriptHash160 sMultiSig
          s = ASC.p2sh shMultiSig
      toBase16 (AI.scriptBytes s) @?= mconcat
        [ "a914", toBase16 (AH.unScriptHash160 shMultiSig), "87" ]
      toBase16 (AI.scriptBytes s) @?=
        "a914db9775b7424ddfd3a69d83c2a9a91bf7953197b187"

  , testProperty "prop" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 1 16) genPub
      m <- forAll $ Gen.int (Range.constant 1 (length pks))
      let Just sMultiSig = ASC.multiSig pks m
          shMultiSig = AH.scriptHash160 sMultiSig
          s = ASC.p2sh shMultiSig
      s === S.Script [ S.OP_HASH160
                     , S.OP_PUSHDATA (AH.unScriptHash160 shMultiSig) S.OPCODE
                     , S.OP_EQUAL ]
      toBase16 (AI.scriptBytes s) === mconcat
        [ "a914", toBase16 (AH.unScriptHash160 shMultiSig), "87"]
  ]

tt_script_p2sh_p2wsh :: TestTree
tt_script_p2sh_p2wsh = testGroup "P2WSH-in-P2SH"
  [ tt_script_p2sh_p2wsh_multiSig
  ]

tt_script_p2sh_p2wsh_multiSig :: TestTree
tt_script_p2sh_p2wsh_multiSig = testGroup "P2WSH-MultiSig-in-P2SH"
  [ testCase "1 of [pk1, pk2]" $ do
      let Just sMultiSig = ASC.multiSig [pk1, pk2] 1
          shMultiSig = AH.scriptSHA256 sMultiSig
          sw = ASC.segWit (ASW.p2wsh shMultiSig)
          swh = AH.scriptHash160 sw
          s = ASC.p2sh swh
      toBase16 (AI.scriptBytes s) @?= mconcat
        [ "a914", toBase16 (AH.unScriptHash160 swh), "87" ]
      toBase16 (AI.scriptBytes s) @?=
        "a9140b5fd2183b91a75ae09e70122bf89e79e71ff1c387"

  , testProperty "prop" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 1 16) genPub
      m <- forAll $ Gen.int (Range.constant 1 (length pks))
      let Just sMultiSig = ASC.multiSig pks m
          shMultiSig = AH.scriptSHA256 sMultiSig
          sw = ASC.segWit (ASW.p2wsh shMultiSig)
          swh = AH.scriptHash160 sw
          s = ASC.p2sh swh
      s === S.Script [ S.OP_HASH160
                     , S.OP_PUSHDATA (AH.unScriptHash160 swh) S.OPCODE
                     , S.OP_EQUAL ]
      toBase16 (AI.scriptBytes s) === mconcat
        [ "a914", toBase16 (AH.unScriptHash160 swh), "87" ]
      let Just s' = A.addressScript <$> A.p2sh_p2wsh_multiSig A.btc pks m
      s' === s
  ]

tt_script_segWit_p2wsh :: TestTree
tt_script_segWit_p2wsh = testGroup "P2WSH"
  [ tt_script_segWit_p2wsh_multiSig
  ]

tt_script_segWit_p2wsh_multiSig :: TestTree
tt_script_segWit_p2wsh_multiSig = testGroup "MultiSig"
  [ testCase "1 of [pk1, pk2]" $ do
      let Just sMultiSig = ASC.multiSig [pk1, pk2] 1
          shMultiSig = AH.scriptSHA256 sMultiSig
          s = ASC.segWit (ASW.p2wsh shMultiSig)
      toBase16 (AI.scriptBytes s) @?= mconcat
        [ "0020", toBase16 (AH.unScriptSHA256 shMultiSig) ]
      toBase16 (AI.scriptBytes s) @?=
        "00209d0d2b0482cfd102ae24a2720645b502fc9a56164ec54e87a51c8d4966cebdbe"

  ,  testProperty "prop" $ property $ do
      pks <- forAll $ Gen.list (Range.constant 1 16) genPub
      m <- forAll $ Gen.int (Range.constant 1 (length pks))
      let Just sMultiSig = ASC.multiSig pks m
          shMultiSig = AH.scriptSHA256 sMultiSig
          s = ASC.segWit (ASW.p2wsh shMultiSig)
      s === S.Script [ S.OP_0
                     , S.OP_PUSHDATA (AH.unScriptSHA256 shMultiSig) S.OPCODE ]
      toBase16 (AI.scriptBytes s) === mconcat
        [ "0020", toBase16 (AH.unScriptSHA256 shMultiSig) ]
  ]

tt_address :: TestTree
tt_address = testGroup "Address"
  [ tt_address_btc
  -- , tt_address_btcTestnet
  ]


tt_address_btc :: TestTree
tt_address_btc = testGroup "BTC"
  [ testGroup "p2pkh"
    [ testCase "pk1" $ do
        A.renderAddress (A.p2pkh A.btc pk1) @?=
          "19fQzMWpVo5esC9e2HGJRWfB2bkYLGEqAb"
    , testCase "pk2" $ do
        A.renderAddress (A.p2pkh A.btc pk2) @?=
          "1ErtkviCvV4H9a2cMCrHwb5LV1ir7aXJpF"
    ]
  , testGroup "p2wpkh"
    [ testCase "pk1" $ do
        A.renderAddress (A.p2wpkh A.btc pk1) @?=
           "bc1qtuzdvulah5rjh964328cxzjamaj4swwhj47xg6"
    , testCase "pk2" $ do
        A.renderAddress (A.p2wpkh A.btc pk2) @?=
          "bc1qnqyn3cz4jsly0de7t5tt9hnngwnp5znmk06th7"
    ]
  , testGroup "p2sh_p2wpkh"
    [ testCase "pk1" $ do
        A.renderAddress (A.p2sh_p2wpkh A.btc pk1) @?=
          "3JhVKigMH4hyZzq1MJPjwZmKY5tV5z69nC"
    , testCase "pk2" $ do
        A.renderAddress (A.p2sh_p2wpkh A.btc pk2) @?=
          "331zaDSamvQKX41nDATJPH1PKeMTQzfm22"
    ]
  , testGroup "p2sh_multiSig"
    [ testCase "1 of [pk1]" $ do
        fmap A.renderAddress (A.p2sh_multiSig A.btc [pk1] 1) @?=
          Just "3H12mXUPXvZPd4ardUEEtZnyTCyWtq5ozH"
    , testCase "1 of [pk2]" $ do
        fmap A.renderAddress (A.p2sh_multiSig A.btc [pk2] 1) @?=
          Just "3QUSu15rwtJXpLUUqGSCXG6i6j6Kv9RvQ5"
    , testCase "1 of [pk1, pk2]" $ do
        fmap A.renderAddress (A.p2sh_multiSig A.btc [pk1, pk2] 1) @?=
          Just "3Mi7MAhA2MeiLM77yYCmBYNjQMkug2fET9"
    , testCase "1 of [pk2, pk1]" $ do
        fmap A.renderAddress (A.p2sh_multiSig A.btc [pk2, pk1] 1) @?=
          Just "3PAA17FxCGdNFMjJMQWGQp1dyvxWcntDvs"
    , testCase "2 of [pk1, pk2]" $ do
        fmap A.renderAddress (A.p2sh_multiSig A.btc [pk1, pk2] 2) @?=
          Just "36i3sEMYTsECQPq6yzrhFq1mWmQRqgn6g3"
    , testCase "2 of [pk2, pk1]" $ do
        fmap A.renderAddress (A.p2sh_multiSig A.btc [pk2, pk1] 2) @?=
          Just "35CaUy8cA7xFSEFAS6mvCJbvugNS11nPBx"
    ]
  , testGroup "p2wsh_multiSig"
    [ testCase "1 of [pk1]" $ do
        fmap A.renderAddress (A.p2wsh_multiSig A.btc [pk1] 1) @?=
          Just "bc1qw2v3lcjvu0mxsqhwt673nt0t2sehhge7hamqzzqg45jyah294lsqw2j5yj"
    , testCase "1 of [pk2]" $ do
        fmap A.renderAddress (A.p2wsh_multiSig A.btc [pk2] 1) @?=
          Just "bc1q7x6l850snyxcxdlwjx3qty6rngkngxuhcjlr9a4dgk4975x6krnqzh99h6"
    , testCase "1 of [pk1, pk2]" $ do
        fmap A.renderAddress (A.p2wsh_multiSig A.btc [pk1, pk2] 1) @?=
          Just "bc1qn5xjkpyzelgs9t3y5feqv3d4qt7f54skfmz5apa9rjx5jekwhklqajj224"
    , testCase "1 of [pk2, pk1]" $ do
        fmap A.renderAddress (A.p2wsh_multiSig A.btc [pk2, pk1] 1) @?=
          Just "bc1qm29gsj6mr73kd7rg4ec848rcmujkgk05sthjcnhx2zeh09n5zcssw7j3pn"
    , testCase "2 of [pk1, pk2]" $ do
        fmap A.renderAddress (A.p2wsh_multiSig A.btc [pk2, pk1] 2) @?=
          Just "bc1qwy7gugcea2f2kckzfhvfu34raw6fevu7dhncdec5gz03ll7lxauq3t0r6u"
    , testCase "2 of [pk2, pk1]" $ do
        fmap A.renderAddress (A.p2wsh_multiSig A.btc [pk1, pk2] 2) @?=
          Just "bc1qxemuancumx68770699xahey9uqdc5qdc3572uxukdhg9wncztg3su4ecfy"
    ]
  , testGroup "p2sh_p2wsh_multiSig"
    [ testCase "pk1" $ do
        fmap A.renderAddress (A.p2sh_p2wsh_multiSig A.btc [pk1] 1) @?=
          Just "323Jncw5hBt6KL1WMSao8RrE8CtrooBt9j"
    , testCase "pk2" $ do
        fmap A.renderAddress (A.p2sh_p2wsh_multiSig A.btc [pk2] 1) @?=
          Just "39U3xFLp1KDcaixq6wyfHHsjiPvCCyVyDa"
    , testCase "1 of [pk1, pk2]" $ do
        fmap A.renderAddress (A.p2sh_p2wsh_multiSig A.btc [pk1, pk2] 1) @?=
          Just "32jA9PfdTQ74svaSQjTZPLeV98WiHEHJF3"
    , testCase "1 of [pk2, pk1]" $ do
        fmap A.renderAddress (A.p2sh_p2wsh_multiSig A.btc [pk2, pk1] 1) @?=
          Just "3GSfokgqW2G77VzTrojnPa1sNjSnYuwDuX"
    , testCase "2 of [pk1, pk2]" $ do
        fmap A.renderAddress (A.p2sh_p2wsh_multiSig A.btc [pk1, pk2] 2) @?=
          Just "3EeuEEn1uwdm2GwctKqSjBV4xPp1JiECgm"
    , testCase "2 of [pk2, pk1]" $ do
        fmap A.renderAddress (A.p2sh_p2wsh_multiSig A.btc [pk2, pk1] 2) @?=
          Just "35LhobEZN2QLpYfMfH35tTyxrxYpANBBkA"
    ]
  ]

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

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)

--------------------------------------------------------------------------------
-- BIP39 mnemonic: abandon abandon arm
--
-- BIP39 seed: 2f7b8b6eeeb7201d88ed45dcd4706ccd76e5feb61f1961990869581781f8944ae2bebff36156c6965054add29f2cf9ee8b7b488179bfba24773ce2b06ec0ce04
--
-- BIP32 Derivation path: m/44'/0'/0'/0/0
--
-- Private key WIF: L1CnCq2yTvfzngQTF4Fnh5daYUwAuysrqU5dNwZFcv33Cwu4Gf2F

pk1 :: K.Pub
Just pk1 = K.parsePub $ fromBase16
  "024ed59d3f6548c41172ec26b9144f547185cb2a9f28fccb6e3ae53dc4f1827154"

pk1hash160 :: AH.PubHash160
pk1hash160 = AH.pubHash160 pk1

--------------------------------------------------------------------------------
-- BIP39 mnemonic: abandon abandon arm
--
-- BIP39 seed: 2f7b8b6eeeb7201d88ed45dcd4706ccd76e5feb61f1961990869581781f8944ae2bebff36156c6965054add29f2cf9ee8b7b488179bfba24773ce2b06ec0ce04
--
-- BIP32 Derivation path: m/84'/0'/0'/0/0
--
-- Private key WIF: L3XRaAfqnZLCcUpcCi4g1GqMLTKvr9nbvCYBmrs1TAjx4acpZ2eY

pk2 :: K.Pub
Just pk2 = K.parsePub $ fromBase16
  "0315bd146d8b45383ad431a8204717dc9f46a7a46c1f44eeaab02d49f4511ccff0"

pk2hash160 :: AH.PubHash160
pk2hash160 = AH.pubHash160 pk2

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

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)