package utils_test

import (
	"testing"

	"github.com/cerc-io/plugeth-statediff/utils"
	"github.com/ethereum/go-ethereum/core/rawdb"
	"github.com/ethereum/go-ethereum/trie"

	"github.com/stretchr/testify/assert"
)

type kvs struct{ k, v string }

var (
	testdata1 = []kvs{
		{"bar", "b"},
		{"barb", "ba"},
		{"bard", "bc"},
		{"bars", "bb"},
		{"fab", "z"},
		{"foo", "a"},
		{"food", "ab"},
	}

	testdata2 = []kvs{
		{"aardvark", "c"},
		{"bar", "b"},
		{"barb", "bd"},
		{"bars", "be"},
		{"fab", "z"},
		{"foo", "a"},
		{"foos", "aa"},
		{"jars", "d"},
	}
)

func TestSymmetricDifferenceIterator(t *testing.T) {
	t.Run("empty case", func(t *testing.T) {
		db := trie.NewDatabase(rawdb.NewMemoryDatabase())
		tree := trie.NewEmpty(db)
		di, count := utils.NewSymmetricDifferenceIterator(tree.NodeIterator(nil), tree.NodeIterator(nil))
		for di.Next(true) {
			t.Errorf("iterator should not yield any elements")
		}
		if *count != 0 {
			t.Errorf("node count should be 0 for empty trie, got %d", *count)
		}
	})

	dba := trie.NewDatabase(rawdb.NewMemoryDatabase())
	triea := trie.NewEmpty(dba)
	for _, val := range testdata1 {
		triea.MustUpdate([]byte(val.k), []byte(val.v))
	}
	rootA, nodesA := triea.Commit(false)
	dba.Update(trie.NewWithNodeSet(nodesA))
	triea, _ = trie.New(trie.TrieID(rootA), dba)

	dbb := trie.NewDatabase(rawdb.NewMemoryDatabase())
	trieb := trie.NewEmpty(dbb)
	for _, val := range testdata2 {
		trieb.MustUpdate([]byte(val.k), []byte(val.v))
	}
	rootB, nodesB := trieb.Commit(false)
	dbb.Update(trie.NewWithNodeSet(nodesB))
	trieb, _ = trie.New(trie.TrieID(rootB), dbb)

	onlyA := make(map[string]string)
	onlyB := make(map[string]string)
	it, _ := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator(nil))
	for it.Next(true) {
		if !it.Leaf() {
			continue
		}
		key, value := string(it.LeafKey()), string(it.LeafBlob())
		if it.FromA() {
			onlyA[key] = value
		} else {
			onlyB[key] = value
		}
	}

	expectedOnlyA := map[string]string{
		"barb": "ba",
		"bard": "bc",
		"bars": "bb",
		"food": "ab",
	}
	expectedOnlyB := map[string]string{
		"aardvark": "c",
		"barb":     "bd",
		"bars":     "be",
		"foos":     "aa",
		"jars":     "d",
	}
	assert.Equal(t, expectedOnlyA, onlyA)
	assert.Equal(t, expectedOnlyB, onlyB)
}