package merkle_test

import (
	"fmt"
	"testing"

	"github.com/brianvoe/gofakeit/v6"
	. "github.com/stretchr/testify/assert"
	"github.com/synapsecns/sanguine/core/merkle"
)

const (
	leafsAmount uint32 = 100
)

func TestMerkleTree(t *testing.T) {
	tree := merkle.NewTree(merkle.MessageTreeHeight)
	// Generate test leafs
	leafs := make([][]byte, leafsAmount)
	for i := uint32(0); i < leafsAmount; i++ {
		leafs[i] = fakeLeaf()
		fmt.Printf("%d: %v\n", i, leafs[i])
	}
	// Insert test leafs
	for i := uint32(0); i < leafsAmount; i++ {
		tree.Insert(leafs[i])
		Equal(t, tree.NumOfItems(), i+1)
	}
	// Check Items()
	items := tree.Items()
	Equal(t, len(items), int(leafsAmount))
	for i := uint32(0); i < leafsAmount; i++ {
		Equal(t, items[i], leafs[i])
	}
	// Check Item()
	for i := uint32(0); i < leafsAmount; i++ {
		item, err := tree.Item(i)
		Nil(t, err)
		Equal(t, item, leafs[i])
	}
	// Check proofs
	for count := uint32(1); count <= leafsAmount; count++ {
		// Get root after `count` leafs were inserted
		root, err := tree.Root(count)
		Nil(t, err)
		// Root should be a bytes32 value.
		Equal(t, len(root), 32)
		for index := uint32(0); index < count; index++ {
			proof, err := tree.MerkleProof(index, count)
			Nil(t, err)
			// Proof should have length equal to `TreeDepth`.
			Equal(t, len(proof), int(merkle.MessageTreeHeight))
			for i := 0; i < len(proof); i++ {
				// Each element should be a bytes32 value.
				Equal(t, len(proof[i]), 32)
			}
			branchRoot, err := merkle.BranchRoot(leafs[index], index, proof, merkle.MessageTreeHeight)
			Nil(t, err)
			// Should match the correct bytes32 value for the root.
			Equal(t, root, branchRoot)
			True(t, merkle.VerifyMerkleProof(root, leafs[index], index, proof, merkle.MessageTreeHeight))
		}
	}
}

func TestIncorrectRequests(t *testing.T) {
	tree := merkle.NewTree(merkle.MessageTreeHeight)
	// Generate test leafs.
	leafs := make([][]byte, leafsAmount)
	for i := uint32(0); i < leafsAmount; i++ {
		leafs[i] = fakeLeaf()
	}
	// Insert test leafs.
	for i := uint32(0); i < leafsAmount; i++ {
		tree.Insert(leafs[i])
	}
	// Check Item() with index out of bound.
	item, err := tree.Item(leafsAmount)
	Nil(t, item)
	NotNil(t, err)
	// Check Root() with count out of bound.
	root, err := tree.Root(leafsAmount + 1)
	Nil(t, root)
	NotNil(t, err)
	// Check MerkleProof() with count out of bound.
	proof, err := tree.MerkleProof(0, leafsAmount+1)
	Nil(t, proof)
	NotNil(t, err)
	// Check MerkleProof() with index >= count.
	proof, err = tree.MerkleProof(10, 10)
	Nil(t, proof)
	NotNil(t, err)
	// Generate a valid proof.
	proof, err = tree.MerkleProof(0, leafsAmount)
	NotNil(t, proof)
	Nil(t, err)
	// Check BranchRoot() with incorrect proof length.
	proof = make([][]byte, merkle.MessageTreeHeight+1)
	branchRoot, err := merkle.BranchRoot(leafs[0], leafsAmount, proof, merkle.MessageTreeHeight)
	Nil(t, branchRoot)
	NotNil(t, err)
	// Get current and historical root.
	root, err = tree.Root(leafsAmount)
	NotNil(t, root)
	Nil(t, err)
	histRoot, err := tree.Root(leafsAmount - 1)
	NotNil(t, histRoot)
	Nil(t, err)
	// Check VerifyMerkleProof() with incorrect proof length.
	False(t, merkle.VerifyMerkleProof(root, leafs[0], 0, proof, merkle.MessageTreeHeight))
	// Generate proofs against current root, they should not work with the historical root.
	for i := uint32(0); i < leafsAmount; i++ {
		proof, err = tree.MerkleProof(i, leafsAmount)
		NotNil(t, proof)
		Nil(t, err)
		// Should not be able to prove against the historical root.
		False(t, merkle.VerifyMerkleProof(histRoot, leafs[i], i, proof, merkle.MessageTreeHeight))
		// Should be able to prove against the current root.
		True(t, merkle.VerifyMerkleProof(root, leafs[i], i, proof, merkle.MessageTreeHeight))
	}
}

func TestNewTreeFromItems(t *testing.T) {
	tree := merkle.NewTree(merkle.MessageTreeHeight)
	// Generate test leafs
	leafs := make([][]byte, leafsAmount)
	for i := uint32(0); i < leafsAmount; i++ {
		leafs[i] = fakeLeaf()
		fmt.Printf("%d: %v\n", i, leafs[i])
	}
	// Insert test leafs
	for i := uint32(0); i < leafsAmount; i++ {
		tree.Insert(leafs[i])
		Equal(t, tree.NumOfItems(), i+1)
	}
	// Get items and generate a new tree from them
	items := tree.Items()
	newTree := merkle.NewTreeFromItems(items, merkle.MessageTreeHeight)
	// Check that the number of items are the same
	Equal(t, tree.NumOfItems(), newTree.NumOfItems())
	// Check that the new tree has the same root
	root, err := tree.Root(leafsAmount)
	Nil(t, err)
	newRoot, err := newTree.Root(leafsAmount)
	Nil(t, err)
	Equal(t, root, newRoot)
}

func fakeLeaf() []byte {
	leaf := make([]byte, 32)
	for i := 0; i < 32; i++ {
		leaf[i] = gofakeit.Uint8()
	}
	return leaf
}