package merkle

import (
	"bytes"
	"crypto/sha256"

	"github.com/tendermint/tendermint/binary"
	. "github.com/tendermint/tendermint/common"
)

type IAVLProof struct {
	LeafNode   IAVLProofLeafNode
	InnerNodes []IAVLProofInnerNode
	RootHash   []byte
}

func (proof *IAVLProof) Verify(keyBytes, valueBytes, rootHash []byte) bool {
	if !bytes.Equal(keyBytes, proof.LeafNode.KeyBytes) {
		return false
	}
	if !bytes.Equal(valueBytes, proof.LeafNode.ValueBytes) {
		return false
	}
	if !bytes.Equal(rootHash, proof.RootHash) {
		return false
	}
	hash := proof.LeafNode.Hash()
	// fmt.Printf("leaf hash: %X\n", hash)
	for _, branch := range proof.InnerNodes {
		hash = branch.Hash(hash)
		// fmt.Printf("branch hash: %X\n", hash)
	}
	// fmt.Printf("root: %X, computed: %X\n", proof.RootHash, hash)
	return bytes.Equal(proof.RootHash, hash)
}

type IAVLProofInnerNode struct {
	Height uint8
	Size   uint
	Left   []byte
	Right  []byte
}

func (branch IAVLProofInnerNode) Hash(childHash []byte) []byte {
	hasher := sha256.New()
	buf := new(bytes.Buffer)
	n, err := int64(0), error(nil)
	binary.WriteUint8(branch.Height, buf, &n, &err)
	binary.WriteUvarint(branch.Size, buf, &n, &err)
	if branch.Left == nil {
		binary.WriteByteSlice(childHash, buf, &n, &err)
		binary.WriteByteSlice(branch.Right, buf, &n, &err)
	} else {
		binary.WriteByteSlice(branch.Left, buf, &n, &err)
		binary.WriteByteSlice(childHash, buf, &n, &err)
	}
	if err != nil {
		panic(Fmt("Failed to hash IAVLProofInnerNode: %v", err))
	}
	// fmt.Printf("InnerNode hash bytes: %X\n", buf.Bytes())
	hasher.Write(buf.Bytes())
	return hasher.Sum(nil)
}

type IAVLProofLeafNode struct {
	KeyBytes   []byte
	ValueBytes []byte
}

func (leaf IAVLProofLeafNode) Hash() []byte {
	hasher := sha256.New()
	buf := new(bytes.Buffer)
	n, err := int64(0), error(nil)
	binary.WriteUint8(0, buf, &n, &err)
	binary.WriteUvarint(1, buf, &n, &err)
	binary.WriteByteSlice(leaf.KeyBytes, buf, &n, &err)
	binary.WriteByteSlice(leaf.ValueBytes, buf, &n, &err)
	if err != nil {
		panic(Fmt("Failed to hash IAVLProofLeafNode: %v", err))
	}
	// fmt.Printf("LeafNode hash bytes:   %X\n", buf.Bytes())
	hasher.Write(buf.Bytes())
	return hasher.Sum(nil)
}

func (node *IAVLNode) constructProof(t *IAVLTree, key interface{}, proof *IAVLProof) (exists bool) {
	if node.height == 0 {
		if t.keyCodec.Compare(node.key, key) == 0 {
			keyBuf, valueBuf := new(bytes.Buffer), new(bytes.Buffer)
			n, err := int64(0), error(nil)
			t.keyCodec.Encode(node.key, keyBuf, &n, &err)
			if err != nil {
				panic(Fmt("Failed to encode node.key: %v", err))
			}
			t.valueCodec.Encode(node.value, valueBuf, &n, &err)
			if err != nil {
				panic(Fmt("Failed to encode node.value: %v", err))
			}
			leaf := IAVLProofLeafNode{
				KeyBytes:   keyBuf.Bytes(),
				ValueBytes: valueBuf.Bytes(),
			}
			proof.LeafNode = leaf
			return true
		} else {
			return false
		}
	} else {
		if t.keyCodec.Compare(key, node.key) < 0 {
			exists := node.getLeftNode(t).constructProof(t, key, proof)
			if !exists {
				return false
			}
			branch := IAVLProofInnerNode{
				Height: node.height,
				Size:   node.size,
				Left:   nil,
				Right:  node.getRightNode(t).hash,
			}
			proof.InnerNodes = append(proof.InnerNodes, branch)
			return true
		} else {
			exists := node.getRightNode(t).constructProof(t, key, proof)
			if !exists {
				return false
			}
			branch := IAVLProofInnerNode{
				Height: node.height,
				Size:   node.size,
				Left:   node.getLeftNode(t).hash,
				Right:  nil,
			}
			proof.InnerNodes = append(proof.InnerNodes, branch)
			return true
		}
	}
}

// Returns nil if key is not in tree.
func (t *IAVLTree) ConstructProof(key interface{}) *IAVLProof {
	if t.root == nil {
		return nil
	}
	t.root.hashWithCount(t) // Ensure that all hashes are calculated.
	proof := &IAVLProof{
		RootHash: t.root.hash,
	}
	t.root.constructProof(t, key, proof)
	return proof
}