{-# LANGUAGE OverloadedStrings #-}

module SpecMerkle (
    tests,
) where

import Crypto.Hash (HashAlgorithm (hashDigestSize), SHA256 (SHA256))
import Data.Binary (decodeOrFail, encode)
import Data.ByteString.Base32 (
    encodeBase32Unpadded,
 )
import Data.List (
    sort,
 )
import Data.Maybe (
    isJust,
 )
import Data.Text (
    pack,
 )
import Data.Text.Encoding (
    encodeUtf8,
 )
import Hedgehog (
    Gen,
    Property,
    annotateShow,
    assert,
    diff,
    failure,
    forAll,
    property,
    tripping,
 )
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Tahoe.CHK.Crypto (
    sha256,
    taggedHash,
 )
import Tahoe.CHK.Merkle (
    Direction (..),
    MerkleTree (MerkleLeaf, MerkleNode),
    breadthFirstList,
    buildTreeOutOfAllTheNodes,
    emptyLeafHash,
    height,
    leafNumberToNodeNumber,
    makeTree,
    mapTree,
    merklePath,
    merkleProof,
    neededHashes,
    pairHash,
    rootHash,
    size,
    treeFromRows,
 )
import Test.Tasty (
    TestTree,
    testGroup,
 )
import Test.Tasty.HUnit (
    assertBool,
    assertEqual,
    testCase,
 )
import Test.Tasty.Hedgehog (testProperty)

tests :: TestTree
tests =
    testGroup
        "Merkle"
        [ testCase "pairHash" $
            assertEqual
                "simple test vector"
                "MNP3F5B64GHVUPQ3U7ZT76D7ZP6NVHHV5KMFLT2IPORIGI5EL57Q"
                (encodeBase32Unpadded $ pairHash "abc" "xyz")
        , testCase "emptyLeafHash" $
            assertEqual
                "simple test vector"
                "T3KZA5VWX3TLOWDEMMGDYIGP62JU57QDUYFH7UULNFKC7MJ2NCRQ"
                (encodeBase32Unpadded $ emptyLeafHash 3)
        , testCase "two leaf tree" $
            assertEqual
                "root hash is leaf pair hash"
                (Just "NFOM5H52FQH5A4F3OL3JCPGAUECQJEW6FUWKW5HWVQDIFSKPM6DQ")
                (encodeBase32Unpadded . rootHash <$> makeTree [sha256 "abc", sha256 "xyz"])
        , testCase "three leaf tree" $
            assertEqual
                "root hash of three leaf tree includes empty node hash"
                (Just $ encodeBase32Unpadded $ pairHash (pairHash (sha256 "abc") (sha256 "xyz")) (pairHash (sha256 "mno") $ emptyLeafHash 3))
                (encodeBase32Unpadded . rootHash <$> makeTree [sha256 "abc", sha256 "xyz", sha256 "mno"])
        , testCase "empty tree" $
            assertEqual
                "empty list results in no tree"
                Nothing
                (makeTree [])
        , testCase "tiny tree" $
            assertEqual
                "a two leaf tree can be constructed"
                (Just (MerkleNode "\162a\224\215\DEL_\204[u-\134\200\245\&8_\210\177=\DELM\217\203V\157\220\169\224tE\145F\145" (MerkleLeaf (sha256 "bar")) (MerkleLeaf (sha256 "baz"))))
                (makeTree [sha256 "bar", sha256 "baz"])
        , testCase "make 6 leaf tree" $
            assertBool "it can be made" $
                isJust (makeTestTree 6)
        , testCase "breadth first traversal (small)" $
            assertEqual
                "tree with one leaf"
                (Just 1)
                (length . breadthFirstList <$> makeTestTree 1)
        , testCase "breadth first traversal (big)" $
            assertEqual
                "tree with 1024 leaves"
                (Just (1024 * 2 - 1))
                (length . breadthFirstList <$> makeTestTree 1024)
        , testCase "show it" $ do
            print $ makeTestTree 2
            return ()
        , testCase "neededHashes test vectors" $
            let Just tree = makeTestTree 8
                needed = (sort . map fst <$>) . neededHashes tree
             in do
                    assertEqual "test vector 1" (Just [2 :: Int, 4, 8]) (needed 0)
                    assertEqual "test vector 2" (Just [2, 4, 7]) (needed 1)
                    assertEqual "test vector 3" (Just [1, 5, 13]) (needed 7)
        , testProperty "all paths same length for merkleProof" prop_merkleProof_length
        , testProperty "all internal nodes have the correct hash" prop_makeTree_hashes
        , testProperty "all merkleProofs prove what they ought" spec_merkleProof_nodeNumbers
        , testProperty "all merkleProofs prove what they ought" spec_merkleProof_hashes
        , testProperty "all merkle paths have a consistent length" spec_merklePath_length
        , testProperty "node numbering round trips through the converters" spec_numberConversion_tripping
        , testProperty "merkle tree block construction" spec_merkleFromRows
        , testProperty "invalid merkle trees fail" spec_invalidMerkle
        , testProperty "merkle trees round-trip through encode / decode" prop_binary_tripping
        ]

prop_binary_tripping :: Property
prop_binary_tripping = property $ do
    (Just someTree) <- forAll genMerkleTree
    let third (_, _, x) = x
    tripping someTree encode ((third <$>) . decodeOrFail)

prop_merkleProof_length :: Property
prop_merkleProof_length = property $ do
    (Just someTree) <- forAll genMerkleTree
    someLeaf <- forAll $ Gen.integral (Range.linear 0 $ height someTree - 1)
    diff (Just $ height someTree - 1) (==) (length <$> merkleProof someTree someLeaf)

prop_makeTree_hashes :: Property
prop_makeTree_hashes = property $ do
    (Just some_tree) <- forAll genMerkleTree
    assert (and $ mapTree checkMerkleProperty some_tree)
  where
    checkMerkleProperty (MerkleLeaf _) = True
    checkMerkleProperty (MerkleNode h l r) = h == pairHash (rootHash l) (rootHash r)

{- | Convert a set of directions to a node to that node's number.  The first
 argument is the node number of the root node from which to follow the
 directions.  For the "true" root of the tree, use 1.
-}
pathToNumber :: Int -> [Direction] -> Int
pathToNumber rootNum [] = rootNum
pathToNumber rootNum (d : ds) = pathToNumber childNum ds
  where
    childNum = case d of
        TurnLeft -> rootNum * 2
        TurnRight -> rootNum * 2 + 1

{- | Convert a set of directions to a node to the numbers of the nodes on the
 proof path to that node.  These are the numbers of the nodes that are
 _siblings_ to nodes on the given path.
-}
proofPathNodes :: Int -> [Direction] -> [Int]
proofPathNodes _ [] = []
proofPathNodes rootNum (d : ds) = siblingNum : proofPathNodes childNum ds
  where
    childNum = case d of
        TurnLeft -> rootNum * 2
        TurnRight -> rootNum * 2 + 1

    siblingNum = case d of
        TurnLeft -> rootNum * 2 + 1
        TurnRight -> rootNum * 2

{- | merkleProof returns a list of tuples where each tuple gives a node number
 and the hash belonging to that node.
-}
spec_merkleProof_hashes :: Property
spec_merkleProof_hashes = property $ do
    (Just someTree) <- forAll genMerkleTree
    someLeafNum <- forAll $ Gen.integral (Range.linear 0 $ height someTree - 1)

    let proof = merkleProof someTree someLeafNum
        -- Brute force search the tree for a matching node.
        getNode :: Int -> MerkleTree -> Int -> [MerkleTree]
        getNode thisNodeNum n@(MerkleLeaf _) targetNodeNum
            | thisNodeNum == targetNodeNum = [n]
            | otherwise = []
        getNode thisNodeNum n@(MerkleNode _ left right) targetNodeNum
            | thisNodeNum == targetNodeNum = [n]
            | otherwise =
                getNode (thisNodeNum * 2) left targetNodeNum
                    ++ getNode (thisNodeNum * 2 + 1) right targetNodeNum

    annotateShow proof

    case proof of
        Nothing -> failure
        Just proof' -> diff (map snd proof') (==) (map (rootHash . head . getNode 1 someTree . fst) proof')

{- | merkleProof returns a list of tuples where each tuple contains a node
 number which is a sibling of a node on the path to a given leaf.
-}
spec_merkleProof_nodeNumbers :: Property
spec_merkleProof_nodeNumbers = property $ do
    (Just someTree) <- forAll genMerkleTree

    -- Choose an arbitrary path through the tree.
    somePath <-
        forAll $
            Gen.list (Range.singleton $ height someTree - 1) $
                Gen.element [TurnLeft, TurnRight]

    let -- Identify the node at the end of the path
        nodeNum = pathToNumber 1 somePath
        leafNum = nodeNumberToLeafNumber someTree nodeNum

        -- Determine the proof path.  It consists of the node numbers of the
        -- siblings of the nodes on the merkle path.
        someProof = proofPathNodes 1 somePath

    annotateShow nodeNum
    annotateShow leafNum

    -- The computed proof path has node numbers which match the proof path node
    -- numbers we computed above.
    diff (map fst <$> merkleProof someTree leafNum) (==) (Just someProof)

spec_numberConversion_tripping :: Property
spec_numberConversion_tripping = property $ do
    (Just someTree) <- forAll genMerkleTree
    someNum <- forAll $ Gen.integral (Range.linear 1 $ size someTree - 1)
    tripping someNum (leafNumberToNodeNumber someTree) (pure . nodeNumberToLeafNumber someTree :: Int -> Maybe Int)

-- | We can build a Merkle tree from its flattened form
spec_merkleFromRows :: Property
spec_merkleFromRows = property $ do
    (Just validTree) <- forAll genMerkleTree
    let nodes = breadthFirstList validTree

    let (Just alleged) = buildTreeOutOfAllTheNodes nodes
    diff alleged (==) validTree

-- | Invalid flattened trees produce errors
spec_invalidMerkle :: Property
spec_invalidMerkle = property $ do
    (Just validTree) <- forAll genMerkleTree
    -- it's a valid list, missing one of the elements
    let nodes = tail (breadthFirstList validTree)

    let maybeTree = buildTreeOutOfAllTheNodes nodes
    diff maybeTree (==) Nothing

-- | The length of all merkle paths equals one less than the given height.
spec_merklePath_length :: Property
spec_merklePath_length = property $ do
    height' <- forAll $ Gen.integral (Range.linear 2 16)
    leafNum <- forAll $ Gen.integral (Range.linear 0 (height' - 1))
    let path = merklePath height' leafNum
    diff (length path) (==) (height' - 1)

genMerkleTree :: Gen (Maybe MerkleTree)
genMerkleTree = makeTestTree <$> Gen.integral (Range.linear 1 256)

makeTestTree :: Int -> Maybe MerkleTree
makeTestTree numleaves = makeTree [taggedHash (hashDigestSize SHA256) "tag" (encodeUtf8 . pack . show $ n) | n <- [0 .. numleaves - 1]]

nodeNumberToLeafNumber :: MerkleTree -> Int -> Int
nodeNumberToLeafNumber tree nodeNum = nodeNum - 1 - size tree `div` 2