from hashlib import blake2b
import struct
from typing import (List, Optional)

from .mininode import (CBlockHeader, block_work_from_compact, ser_compactsize, ser_uint256)
from .util import (
    NU5_BRANCH_ID,
    assert_equal,
)

def H(msg: bytes, consensusBranchId: int) -> bytes:
    digest = blake2b(
        digest_size=32,
        person=b'ZcashHistory' + struct.pack("<I", consensusBranchId))
    digest.update(msg)
    return digest.digest()

class ZcashMMRNode():
    # leaf nodes have no children
    left_child: Optional['ZcashMMRNode']
    right_child: Optional['ZcashMMRNode']

    # commitments
    hashSubtreeCommitment: bytes
    nEarliestTimestamp: int
    nLatestTimestamp: int
    nEarliestTargetBits: int
    nLatestTargetBits: int
    hashEarliestSaplingRoot: bytes # left child's sapling root
    hashLatestSaplingRoot: bytes # right child's sapling root
    nSubTreeTotalWork: int # total difficulty accumulated within each subtree
    nEarliestHeight: int
    nLatestHeight: int
    nSaplingTxCount: int # number of Sapling transactions in block
    # NU5 only.
    hashEarliestOrchardRoot: Optional[bytes] # left child's Orchard root
    hashLatestOrchardRoot: Optional[bytes] # right child's Orchard root
    nOrchardTxCount: Optional[int] # number of Orchard transactions in block

    consensusBranchId: bytes

    @classmethod
    def from_block(
        Z, block: CBlockHeader, height,
        sapling_root, sapling_tx_count,
        consensusBranchId,
        v2_data=None
    ) -> 'ZcashMMRNode':
        '''Create a leaf node from a block'''
        if v2_data is not None:
            assert_equal(consensusBranchId, NU5_BRANCH_ID)
            orchard_root = v2_data[0]
            orchard_tx_count = v2_data[1]
        else:
            orchard_root = None
            orchard_tx_count = None

        node = Z()
        node.left_child = None
        node.right_child = None
        node.hashSubtreeCommitment = ser_uint256(block.rehash())
        node.nEarliestTimestamp = block.nTime
        node.nLatestTimestamp = block.nTime
        node.nEarliestTargetBits = block.nBits
        node.nLatestTargetBits = block.nBits
        node.hashEarliestSaplingRoot = sapling_root
        node.hashLatestSaplingRoot = sapling_root
        node.nSubTreeTotalWork = block_work_from_compact(block.nBits)
        node.nEarliestHeight = height
        node.nLatestHeight = height
        node.nSaplingTxCount = sapling_tx_count
        node.hashEarliestOrchardRoot = orchard_root
        node.hashLatestOrchardRoot = orchard_root
        node.nOrchardTxCount = orchard_tx_count
        node.consensusBranchId = consensusBranchId
        return node

    def serialize(self) -> bytes:
        '''serializes a node'''
        buf = b''
        buf += self.hashSubtreeCommitment
        buf += struct.pack("<I", self.nEarliestTimestamp)
        buf += struct.pack("<I", self.nLatestTimestamp)
        buf += struct.pack("<I", self.nEarliestTargetBits)
        buf += struct.pack("<I", self.nLatestTargetBits)
        buf += self.hashEarliestSaplingRoot
        buf += self.hashLatestSaplingRoot
        buf += ser_uint256(self.nSubTreeTotalWork)
        buf += ser_compactsize(self.nEarliestHeight)
        buf += ser_compactsize(self.nLatestHeight)
        buf += ser_compactsize(self.nSaplingTxCount)
        if self.hashEarliestOrchardRoot is not None:
            buf += self.hashEarliestOrchardRoot
            buf += self.hashLatestOrchardRoot
            buf += ser_compactsize(self.nOrchardTxCount)
        return buf

def make_parent(
        left_child: ZcashMMRNode,
        right_child: ZcashMMRNode) -> ZcashMMRNode:
    parent = ZcashMMRNode()
    parent.left_child = left_child
    parent.right_child = right_child
    parent.hashSubtreeCommitment = H(
        left_child.serialize() + right_child.serialize(),
        left_child.consensusBranchId,
    )
    parent.nEarliestTimestamp = left_child.nEarliestTimestamp
    parent.nLatestTimestamp = right_child.nLatestTimestamp
    parent.nEarliestTargetBits = left_child.nEarliestTargetBits
    parent.nLatestTargetBits = right_child.nLatestTargetBits
    parent.hashEarliestSaplingRoot = left_child.hashEarliestSaplingRoot
    parent.hashLatestSaplingRoot = right_child.hashLatestSaplingRoot
    parent.nSubTreeTotalWork = left_child.nSubTreeTotalWork + right_child.nSubTreeTotalWork
    parent.nEarliestHeight = left_child.nEarliestHeight
    parent.nLatestHeight = right_child.nLatestHeight
    parent.nSaplingTxCount = left_child.nSaplingTxCount + right_child.nSaplingTxCount
    parent.hashEarliestOrchardRoot = left_child.hashEarliestOrchardRoot
    parent.hashLatestOrchardRoot = right_child.hashLatestOrchardRoot
    parent.nOrchardTxCount = (
        left_child.nOrchardTxCount + right_child.nOrchardTxCount
        if left_child.nOrchardTxCount is not None and right_child.nOrchardTxCount is not None
        else None)
    parent.consensusBranchId = left_child.consensusBranchId
    return parent

def make_root_commitment(root: ZcashMMRNode) -> bytes:
    '''Makes the root commitment for a blockheader'''
    return H(root.serialize(), root.consensusBranchId)

def get_peaks(node: ZcashMMRNode) -> List[ZcashMMRNode]:
    peaks: List[ZcashMMRNode] = []

    # Get number of leaves.
    leaves = node.nLatestHeight - (node.nEarliestHeight - 1)
    assert(leaves > 0)

    # Check if the number of leaves in this subtree is a power of two.
    if (leaves & (leaves - 1)) == 0:
        # This subtree is full, and therefore a single peak. This also covers
        # the case of a single isolated leaf.
        peaks.append(node)
    else:
        # This is one of the generated nodes; search within its children.
        peaks.extend(get_peaks(node.left_child))
        peaks.extend(get_peaks(node.right_child))

    return peaks


def bag_peaks(peaks: List[ZcashMMRNode]) -> ZcashMMRNode:
    '''
    "Bag" a list of peaks, and return the final root
    '''
    root = peaks[0]
    for i in range(1, len(peaks)):
        root = make_parent(root, peaks[i])
    return root


def append(root: ZcashMMRNode, leaf: ZcashMMRNode) -> ZcashMMRNode:
    '''Append a leaf to an existing tree, return the new tree root'''
    # recursively find a list of peaks in the current tree
    peaks: List[ZcashMMRNode] = get_peaks(root)
    merged: List[ZcashMMRNode] = []

    # Merge peaks from right to left.
    # This will produce a list of peaks in reverse order
    current = leaf
    for peak in peaks[::-1]:
        current_leaves = current.nLatestHeight - (current.nEarliestHeight - 1)
        peak_leaves = peak.nLatestHeight - (peak.nEarliestHeight - 1)

        if current_leaves == peak_leaves:
            current = make_parent(peak, current)
        else:
            merged.append(current)
            current = peak
    merged.append(current)

    # finally, bag the merged peaks
    return bag_peaks(merged[::-1])

def delete(root: ZcashMMRNode) -> ZcashMMRNode:
    '''
    Delete the rightmost leaf node from an existing MMR
    Return the new tree root
    '''

    n_leaves = root.nLatestHeight - (root.nEarliestHeight - 1)
    # if there were an odd number of leaves,
    # simply replace root with left_child
    if n_leaves & 1:
        return root.left_child

    # otherwise, we need to re-bag the peaks.
    else:
        # first peak
        peaks = [root.left_child]

        # we do this traversing the right (unbalanced) side of the tree
        # we keep the left side (balanced subtree or leaf) of each subtree
        # until we reach a leaf
        subtree_root = root.right_child
        while subtree_root.left_child:
            peaks.append(subtree_root.left_child)
            subtree_root = subtree_root.right_child

    new_root = bag_peaks(peaks)
    return new_root