content: {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  {-# LANGUAGE OverloadedRecordDot #-}
  {-# LANGUAGE NoFieldSelectors #-}
  
  -- |
  -- Module      : Haskoin.Block.Merkle
  -- Copyright   : No rights reserved
  -- License     : MIT
  -- Maintainer  : jprupp@protonmail.ch
  -- Stability   : experimental
  -- Portability : POSIX
  --
  -- Function to deal with Merkle trees inside blocks.
  module Haskoin.Block.Merkle
    ( -- * Merkle Blocks
      MerkleBlock (..),
      MerkleRoot,
      FlagBits,
      PartialMerkleTree,
      buildMerkleRoot,
      merkleBlockTxs,
      testMerkleRoot,
      buildPartialMerkle,
      decodeMerkleFlags,
      encodeMerkleFlags,
      calcTreeHeight,
      calcTreeWidth,
      hash2,
      calcHash,
      traverseAndBuild,
      traverseAndExtract,
      extractMatches,
      splitIn,
      boolsToWord8,
    )
  where
  
  import Control.DeepSeq
  import Control.Monad (forM_, replicateM, when)
  import Data.Binary (Binary (..))
  import Data.Bits
  import qualified Data.ByteString as BS
  import Data.Bytes.Get
  import Data.Bytes.Put
  import Data.Bytes.Serial
  import Data.Either (isRight)
  import Data.Hashable
  import Data.Maybe
  import Data.Serialize (Serialize (..))
  import Data.Word (Word32, Word8)
  import GHC.Generics
  import Haskoin.Block.Common
  import Haskoin.Crypto.Hash
  import Haskoin.Network.Common
  import Haskoin.Network.Data
  import Haskoin.Transaction.Common
  
  -- | Hash of the block's Merkle root.
  type MerkleRoot = Hash256
  
  -- | Bits that are used to rebuild partial merkle tree transaction hash list.
  type FlagBits = [Bool]
  
  -- | Partial Merkle tree for a filtered block.
  type PartialMerkleTree = [Hash256]
  
  -- | Filtered block: a block with a partial Merkle tree that only includes the
  -- transactions that pass a bloom filter that was negotiated.
  data MerkleBlock = MerkleBlock
    { -- | block header
      header :: !BlockHeader,
      -- | total number of transactions in block
      txn :: !Word32,
      -- | hashes in depth-first order
      hashes :: !PartialMerkleTree,
      -- | bits to rebuild partial merkle tree
      flags :: !FlagBits
    }
    deriving (Eq, Show, Read, Generic, Hashable, NFData)
  
  instance Serial MerkleBlock where
    deserialize = do
      header <- deserialize
      ntx <- getWord32le
      (VarInt matchLen) <- deserialize
      hashes <- replicateM (fromIntegral matchLen) deserialize
      (VarInt flagLen) <- deserialize
      ws <- replicateM (fromIntegral flagLen) getWord8
      return $ MerkleBlock header ntx hashes (decodeMerkleFlags ws)
  
    serialize (MerkleBlock h ntx hashes flags) = do
      serialize h
      putWord32le ntx
      putVarInt $ length hashes
      forM_ hashes serialize
      let ws = encodeMerkleFlags flags
      putVarInt $ length ws
      forM_ ws putWord8
  
  instance Binary MerkleBlock where
    put = serialize
    get = deserialize
  
  instance Serialize MerkleBlock where
    put = serialize
    get = deserialize
  
  -- | Unpack Merkle flags into 'FlagBits' structure.
  decodeMerkleFlags :: [Word8] -> FlagBits
  decodeMerkleFlags ws =
    [ b | p <- [0 .. length ws * 8 - 1], b <- [testBit (ws !! (p `div` 8)) (p `mod` 8)]
    ]
  
  -- | Pack Merkle flags from 'FlagBits'.
  encodeMerkleFlags :: FlagBits -> [Word8]
  encodeMerkleFlags bs = map boolsToWord8 $ splitIn 8 bs
  
  -- | Computes the height of a Merkle tree.
  calcTreeHeight ::
    -- | number of transactions (leaf nodes)
    Int ->
    -- | height of the merkle tree
    Int
  calcTreeHeight ntx
    | ntx < 2 = 0
    | even ntx = 1 + calcTreeHeight (ntx `div` 2)
    | otherwise = calcTreeHeight $ ntx + 1
  
  -- | Computes the width of a Merkle tree at a specific height. The transactions
  -- are at height 0.
  calcTreeWidth ::
    -- | number of transactions (leaf nodes)
    Int ->
    -- | height at which we want to compute the width
    Int ->
    -- | width of the Merkle tree
    Int
  calcTreeWidth ntx h = (ntx + (1 `shiftL` h) - 1) `shiftR` h
  
  -- | Computes the root of a Merkle tree from a list of leaf node hashes.
  buildMerkleRoot ::
    -- | transaction hashes (leaf nodes)
    [TxHash] ->
    -- | root of the Merkle tree
    MerkleRoot
  buildMerkleRoot txs = calcHash (calcTreeHeight $ length txs) 0 txs
  
  -- | Concatenate and compute double SHA256.
  hash2 :: Hash256 -> Hash256 -> Hash256
  hash2 a b = doubleSHA256 $ runPutS (serialize a) <> runPutS (serialize b)
  
  -- | Computes the hash of a specific node in a Merkle tree.
  calcHash ::
    -- | height of the node
    Int ->
    -- | position of the node (0 for the leftmost node)
    Int ->
    -- | transaction hashes (leaf nodes)
    [TxHash] ->
    -- | hash of the node at the specified position
    Hash256
  calcHash height pos txs
    | height < 0 || pos < 0 = error "calcHash: Invalid parameters"
    | height == 0 = (txs !! pos).get
    | otherwise = hash2 left right
    where
      left = calcHash (height - 1) (pos * 2) txs
      right
        | pos * 2 + 1 < calcTreeWidth (length txs) (height - 1) =
            calcHash (height - 1) (pos * 2 + 1) txs
        | otherwise = left
  
  -- | Build a partial Merkle tree. Provide a list of tuples with all transaction
  -- hashes in the block, and whether the transaction is to be included in the
  -- partial tree. Returns a flag bits structure and the computed partial Merkle
  -- tree.
  buildPartialMerkle ::
    -- | transaction hash and whether to include
    [(TxHash, Bool)] ->
    -- | flag bits and partial Merkle tree
    (FlagBits, PartialMerkleTree)
  buildPartialMerkle hs = traverseAndBuild (calcTreeHeight $ length hs) 0 hs
  
  -- | Helper function to build partial Merkle tree. Used by 'buildPartialMerkle'
  -- above.
  traverseAndBuild ::
    Int -> Int -> [(TxHash, Bool)] -> (FlagBits, PartialMerkleTree)
  traverseAndBuild height pos txs
    | height < 0 || pos < 0 = error "traverseAndBuild: Invalid parameters"
    | height == 0 || not match = ([match], [calcHash height pos t])
    | otherwise = (match : lb ++ rb, lh ++ rh)
    where
      t = map fst txs
      s = pos `shiftL` height
      e = min (length txs) $ (pos + 1) `shiftL` height
      match = any snd $ take (e - s) $ drop s txs
      (lb, lh) = traverseAndBuild (height - 1) (pos * 2) txs
      (rb, rh)
        | (pos * 2 + 1) < calcTreeWidth (length txs) (height - 1) =
            traverseAndBuild (height - 1) (pos * 2 + 1) txs
        | otherwise = ([], [])
  
  -- | Helper function to extract transaction hashes from partial Merkle tree.
  traverseAndExtract ::
    Int ->
    Int ->
    Int ->
    FlagBits ->
    PartialMerkleTree ->
    Maybe (MerkleRoot, [TxHash], Int, Int)
  traverseAndExtract height pos ntx flags hashes
    | null flags = Nothing
    | height == 0 || not match = leafResult
    | isNothing leftM = Nothing
    | (pos * 2 + 1) >= calcTreeWidth ntx (height - 1) =
        Just (hash2 lh lh, lm, lcf + 1, lch)
    | isNothing rightM = Nothing
    | otherwise =
        Just (hash2 lh rh, lm ++ rm, lcf + rcf + 1, lch + rch)
    where
      leafResult
        | null hashes = Nothing
        | otherwise = Just (h, [TxHash h | height == 0 && match], 1, 1)
      (match : fs) = flags
      (h : _) = hashes
      leftM = traverseAndExtract (height - 1) (pos * 2) ntx fs hashes
      (lh, lm, lcf, lch) = fromMaybe e leftM
      rightM =
        traverseAndExtract
          (height - 1)
          (pos * 2 + 1)
          ntx
          (drop lcf fs)
          (drop lch hashes)
      (rh, rm, rcf, rch) = fromMaybe e rightM
      e = error "traverseAndExtract: unexpected error extracting a Maybe value"
  
  -- | Extracts the matching hashes from a partial merkle tree. This will return
  -- the list of transaction hashes that have been included (set to true) in
  -- a call to 'buildPartialMerkle'.
  extractMatches ::
    Network ->
    FlagBits ->
    PartialMerkleTree ->
    -- | number of transaction at height 0 (leaf nodes)
    Int ->
    -- | Merkle root and list of matching transaction hashes
    Either String (MerkleRoot, [TxHash])
  extractMatches net flags hashes ntx
    | ntx == 0 =
        Left
          "extractMatches: number of transactions can not be 0"
    | ntx > net.maxBlockSize `div` 60 =
        Left
          "extractMatches: number of transactions excessively high"
    | length hashes > ntx =
        Left
          "extractMatches: More hashes provided than the number of transactions"
    | length flags < length hashes =
        Left
          "extractMatches: At least one bit per node and one bit per hash"
    | isNothing resM =
        Left
          "extractMatches: traverseAndExtract failed"
    | (nBitsUsed + 7) `div` 8 /= (length flags + 7) `div` 8 =
        Left
          "extractMatches: All bits were not consumed"
    | nHashUsed /= length hashes =
        Left $
          "extractMatches: All hashes were not consumed: " ++ show nHashUsed
    | otherwise = return (merkRoot, matches)
    where
      resM = traverseAndExtract (calcTreeHeight ntx) 0 ntx flags hashes
      (merkRoot, matches, nBitsUsed, nHashUsed) = fromMaybe e resM
      e = error "extractMatches: unexpected error extracting a Maybe value"
  
  -- | Helper function to split a list in chunks 'Int' length. Last chunk may be
  -- smaller.
  splitIn :: Int -> [a] -> [[a]]
  splitIn _ [] = []
  splitIn c xs = xs1 : splitIn c xs2
    where
      (xs1, xs2) = splitAt c xs
  
  -- | Pack up to eight bools in a byte.
  boolsToWord8 :: [Bool] -> Word8
  boolsToWord8 [] = 0
  boolsToWord8 xs = foldl setBit 0 (map snd $ filter fst $ zip xs [0 .. 7])
  
  -- | Get matching transactions from Merkle block.
  merkleBlockTxs :: Network -> MerkleBlock -> Either String [TxHash]
  merkleBlockTxs net b = do
    (root, ths) <- extractMatches net b.flags b.hashes (fromIntegral b.txn)
    when (root /= b.header.merkle) $
      Left "merkleBlockTxs: Merkle root incorrect"
    return ths
  
  -- | Check if Merkle block root is valid against the block header.
  testMerkleRoot :: Network -> MerkleBlock -> Bool
  testMerkleRoot net = isRight . merkleBlockTxs net
  