package main

import (
	"crypto/sha256"
	"log"
	"math"
)

type MerkleTree struct {
	RootNode *MerkleNode
}

type MerkleNode struct {
	Left  *MerkleNode
	Right *MerkleNode
	Data  []byte
}

func NewMerkleTree(data [][]byte) *MerkleTree {
	var nodes []MerkleNode

	for _, datum := range data {
		node := NewMerkleNode(nil, nil, datum)
		nodes = append(nodes, node)
	}

	for i := 0; math.Pow(2, i) < len(data); i++ {
		var newLevel []MerkleNode
		nodeNums := len(nodes)
		for j := 1; j < nodeNums; j += 2 {
			node := NewMerkleNode(&nodes[j-1], &nodes[j], nil)
			newLevel = append(newLevel, *node)
		}
		if nodeNums%2 != 0 {
			node := NewMerkleNode(nil, &nodes[nodeNums-1], nil)
			newLevel = append(newLevel, *node)
		}
		nodes = newLevel
	}
	mTree := MerkleTree{&nodes[0]}

	return mTree
}

func NewMerkleNode(left, right *MerkleNode, data []byte) *MerkleNode {
	mNode := MerkleNode{}
	//规定子节点为单数时，落单的节点为父节点的right分支。
	if left == nil && right == nil {
		hash := sha256.Sum256(data)
		mNode.Data = hash[:]
	} else if left != nil && right == nil {
		log.Panic("规定子节点为单数时，落单的节点为父节点的right分支。")
	} else if left == nil && right != nil {
		mNode.Data = right.Data
	} else {
		prevHashs := append(left.Data, right.Data)
		hash := sha256.Sum256(prevHashs)
		mNode.Data = hash[:]
	}
	mNode.Left = left
	mNode.Right = right

	return &mNode
}