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{
		{"barb", "ba"},
		{"bard", "bc"},
		{"bars", "bb"},
		{"bar", "b"},
		{"fab", "z"},
		{"food", "ab"},
		{"foo", "a"},
	}

	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("with no difference", func(t *testing.T) {
		db := trie.NewDatabase(rawdb.NewMemoryDatabase())
		triea := trie.NewEmpty(db)
		di, count := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), triea.NodeIterator(nil))
		for di.Next(true) {
			t.Errorf("iterator should not yield any elements")
		}
		assert.Equal(t, 0, *count)

		triea.MustUpdate([]byte("foo"), []byte("bar"))
		di, count = utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), triea.NodeIterator(nil))
		for di.Next(true) {
			t.Errorf("iterator should not yield any elements")
		}
		assert.Equal(t, 2, *count)

		trieb := trie.NewEmpty(db)
		di, count = utils.NewSymmetricDifferenceIterator(
			triea.NodeIterator([]byte("jars")),
			trieb.NodeIterator(nil))
		for di.Next(true) {
			t.Errorf("iterator should not yield any elements")
			t.Logf("%s", di.LeafKey())
		}
		assert.Equal(t, 0, *count)

		// TODO will fail until merged: https://github.com/ethereum/go-ethereum/pull/27838
		// di, count = utils.NewSymmetricDifferenceIterator(
		// 	triea.NodeIterator([]byte("food")),
		// 	trieb.NodeIterator(nil))
		// for di.Next(true) {
		// 	t.Errorf("iterator should not yield any elements")
		// 	t.Logf("%s", di.LeafKey())
		// }
		// assert.Equal(t, 0, *count)
	})

	t.Run("small difference", func(t *testing.T) {
		dba := trie.NewDatabase(rawdb.NewMemoryDatabase())
		triea := trie.NewEmpty(dba)

		dbb := trie.NewDatabase(rawdb.NewMemoryDatabase())
		trieb := trie.NewEmpty(dbb)
		trieb.MustUpdate([]byte("foo"), []byte("bar"))

		di, count := utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator(nil))
		leaves := 0
		for di.Next(true) {
			if di.Leaf() {
				assert.False(t, di.CommonPath())
				assert.Equal(t, "foo", string(di.LeafKey()))
				assert.Equal(t, "bar", string(di.LeafBlob()))
				leaves++
			}
		}
		assert.Equal(t, 1, leaves)
		assert.Equal(t, 2, *count)

		trieb.MustUpdate([]byte("quux"), []byte("bars"))
		di, count = utils.NewSymmetricDifferenceIterator(triea.NodeIterator(nil), trieb.NodeIterator([]byte("quux")))
		leaves = 0
		for di.Next(true) {
			if di.Leaf() {
				assert.False(t, di.CommonPath())
				assert.Equal(t, "quux", string(di.LeafKey()))
				assert.Equal(t, "bars", string(di.LeafBlob()))
				leaves++
			}
		}
		assert.Equal(t, 1, leaves)
		assert.Equal(t, 1, *count)
	})

	dba := trie.NewDatabase(rawdb.NewMemoryDatabase())
	triea := trie.NewEmpty(dba)
	for _, val := range testdata1 {
		triea.MustUpdate([]byte(val.k), []byte(val.v))
	}
	dbb := trie.NewDatabase(rawdb.NewMemoryDatabase())
	trieb := trie.NewEmpty(dbb)
	for _, val := range testdata2 {
		trieb.MustUpdate([]byte(val.k), []byte(val.v))
	}

	onlyA := make(map[string]string)
	onlyB := make(map[string]string)
	var deletions, creations []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
			if !it.CommonPath() {
				deletions = append(deletions, key)
			}
		} else {
			onlyB[key] = value
			if !it.CommonPath() {
				creations = append(creations, key)
			}
		}
	}

	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",
	}
	expectedDeletions := []string{
		"bard",
		"food",
	}
	expectedCreations := []string{
		"aardvark",
		"foos",
		"jars",
	}
	assert.Equal(t, expectedOnlyA, onlyA)
	assert.Equal(t, expectedOnlyB, onlyB)
	assert.Equal(t, expectedDeletions, deletions)
	assert.Equal(t, expectedCreations, creations)
}