package mpt

import (
	"encoding/binary"
	"testing"

	"github.com/nspcc-dev/neo-go/pkg/core/storage"
	"github.com/nspcc-dev/neo-go/pkg/io"
	"github.com/nspcc-dev/neo-go/pkg/util"
	"github.com/stretchr/testify/require"
)

func TestBillet_RestoreHashNode(t *testing.T) {
	check := func(t *testing.T, tr *Billet, expectedRoot Node, expectedNode Node, expectedRefCount uint32) {
		_ = expectedRoot.Hash()
		_ = tr.root.Hash()
		require.Equal(t, expectedRoot, tr.root)
		expectedBytes, err := tr.Store.Get(makeStorageKey(expectedNode.Hash()))
		if expectedRefCount != 0 {
			require.NoError(t, err)
			require.Equal(t, expectedRefCount, binary.LittleEndian.Uint32(expectedBytes[len(expectedBytes)-4:]))
		} else {
			require.ErrorIs(t, err, storage.ErrKeyNotFound)
		}
	}

	t.Run("parent is Extension", func(t *testing.T) {
		t.Run("restore Branch", func(t *testing.T) {
			b := NewBranchNode()
			b.Children[0] = NewExtensionNode([]byte{0x01}, NewLeafNode([]byte{0xAB, 0xCD}))
			b.Children[5] = NewExtensionNode([]byte{0x01}, NewLeafNode([]byte{0xAB, 0xDE}))
			path := toNibbles([]byte{0xAC})
			e := NewExtensionNode(path, NewHashNode(b.Hash()))
			tr := NewBillet(e.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
			tr.root = e

			// OK
			n := new(NodeObject)
			n.DecodeBinary(io.NewBinReaderFromBuf(b.Bytes()))
			require.NoError(t, tr.RestoreHashNode(path, n.Node))
			expected := NewExtensionNode(path, n.Node)
			check(t, tr, expected, n.Node, 1)

			// One more time (already restored) => panic expected, no refcount changes
			require.Panics(t, func() {
				_ = tr.RestoreHashNode(path, n.Node)
			})
			check(t, tr, expected, n.Node, 1)

			// Same path, but wrong hash => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode(path, NewBranchNode()), ErrRestoreFailed)
			check(t, tr, expected, n.Node, 1)

			// New path (changes in the MPT structure are not allowed) => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode(toNibbles([]byte{0xAB}), n.Node), ErrRestoreFailed)
			check(t, tr, expected, n.Node, 1)
		})

		t.Run("restore Leaf", func(t *testing.T) {
			l := NewLeafNode([]byte{0xAB, 0xCD})
			path := toNibbles([]byte{0xAC})
			e := NewExtensionNode(path, NewHashNode(l.Hash()))
			tr := NewBillet(e.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
			tr.root = e

			// OK
			require.NoError(t, tr.RestoreHashNode(path, l))
			expected := NewHashNode(e.Hash()) // leaf should be collapsed immediately => extension should also be collapsed
			expected.Collapsed = true
			check(t, tr, expected, l, 1)

			// One more time (already restored and collapsed) => error expected, no refcount changes
			require.Error(t, tr.RestoreHashNode(path, l))
			check(t, tr, expected, l, 1)

			// Same path, but wrong hash => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode(path, NewLeafNode([]byte{0xAB, 0xEF})), ErrRestoreFailed)
			check(t, tr, expected, l, 1)

			// New path (changes in the MPT structure are not allowed) => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode(toNibbles([]byte{0xAB}), l), ErrRestoreFailed)
			check(t, tr, expected, l, 1)
		})

		t.Run("restore Hash", func(t *testing.T) {
			h := NewHashNode(util.Uint256{1, 2, 3})
			path := toNibbles([]byte{0xAC})
			e := NewExtensionNode(path, h)
			tr := NewBillet(e.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
			tr.root = e

			// no-op
			require.ErrorIs(t, tr.RestoreHashNode(path, h), ErrRestoreFailed)
			check(t, tr, e, h, 0)
		})
	})

	t.Run("parent is Leaf", func(t *testing.T) {
		l := NewLeafNode([]byte{0xAB, 0xCD})
		path := []byte{}
		tr := NewBillet(l.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
		tr.root = l

		// Already restored => panic expected
		require.Panics(t, func() {
			_ = tr.RestoreHashNode(path, l)
		})

		// Same path, but wrong hash => error expected, no refcount changes
		require.ErrorIs(t, tr.RestoreHashNode(path, NewLeafNode([]byte{0xAB, 0xEF})), ErrRestoreFailed)

		// Non-nil path, but MPT structure can't be changed => error expected, no refcount changes
		require.ErrorIs(t, tr.RestoreHashNode(toNibbles([]byte{0xAC}), NewLeafNode([]byte{0xAB, 0xEF})), ErrRestoreFailed)
	})

	t.Run("parent is Branch", func(t *testing.T) {
		t.Run("middle child", func(t *testing.T) {
			l1 := NewLeafNode([]byte{0xAB, 0xCD})
			l2 := NewLeafNode([]byte{0xAB, 0xDE})
			b := NewBranchNode()
			b.Children[5] = NewHashNode(l1.Hash())
			b.Children[lastChild] = NewHashNode(l2.Hash())
			tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
			tr.root = b

			// OK
			path := []byte{0x05}
			require.NoError(t, tr.RestoreHashNode(path, l1))
			check(t, tr, b, l1, 1)

			// One more time (already restored) => panic expected.
			// It's an MPT pool duty to avoid such situations during real restore process.
			require.Panics(t, func() {
				_ = tr.RestoreHashNode(path, l1)
			})
			// No refcount changes expected.
			check(t, tr, b, l1, 1)

			// Same path, but wrong hash => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode(path, NewLeafNode([]byte{0xAD})), ErrRestoreFailed)
			check(t, tr, b, l1, 1)

			// New path pointing to the empty HashNode (changes in the MPT structure are not allowed) => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode([]byte{0x01}, l1), ErrRestoreFailed)
			check(t, tr, b, l1, 1)
		})

		t.Run("last child", func(t *testing.T) {
			l1 := NewLeafNode([]byte{0xAB, 0xCD})
			l2 := NewLeafNode([]byte{0xAB, 0xDE})
			b := NewBranchNode()
			b.Children[5] = NewHashNode(l1.Hash())
			b.Children[lastChild] = NewHashNode(l2.Hash())
			tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
			tr.root = b

			// OK
			path := []byte{}
			require.NoError(t, tr.RestoreHashNode(path, l2))
			check(t, tr, b, l2, 1)

			// One more time (already restored) => panic expected.
			// It's an MPT pool duty to avoid such situations during real restore process.
			require.Panics(t, func() {
				_ = tr.RestoreHashNode(path, l2)
			})
			// No refcount changes expected.
			check(t, tr, b, l2, 1)

			// Same path, but wrong hash => error expected, no refcount changes
			require.ErrorIs(t, tr.RestoreHashNode(path, NewLeafNode([]byte{0xAD})), ErrRestoreFailed)
			check(t, tr, b, l2, 1)
		})

		t.Run("two children with same hash", func(t *testing.T) {
			l := NewLeafNode([]byte{0xAB, 0xCD})
			b := NewBranchNode()
			// two same hashnodes => leaf's refcount expected to be 2 in the end.
			b.Children[3] = NewHashNode(l.Hash())
			b.Children[4] = NewHashNode(l.Hash())
			tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
			tr.root = b

			// OK
			require.NoError(t, tr.RestoreHashNode([]byte{0x03}, l))
			expected := b
			expected.Children[3].(*HashNode).Collapsed = true
			check(t, tr, b, l, 1)

			// Restore another node with the same hash => no error expected, refcount should be incremented.
			// Branch node should be collapsed.
			require.NoError(t, tr.RestoreHashNode([]byte{0x04}, l))
			res := NewHashNode(b.Hash())
			res.Collapsed = true
			check(t, tr, res, l, 2)
		})
	})

	t.Run("parent is Hash", func(t *testing.T) {
		l := NewLeafNode([]byte{0xAB, 0xCD})
		b := NewBranchNode()
		b.Children[3] = NewHashNode(l.Hash())
		b.Children[4] = NewHashNode(l.Hash())
		tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())

		// Should fail, because if it's a hash node with non-empty path, then the node
		// has already been collapsed.
		require.Error(t, tr.RestoreHashNode([]byte{0x03}, l))
	})
}