// Copyright 2021 The Erigon Authors
// This file is part of Erigon.
//
// Erigon is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Erigon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with Erigon. If not, see <http://www.gnu.org/licenses/>.

package kvcache

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"sync"
	"testing"
	"time"

	keccak "github.com/erigontech/fastkeccak"
	"github.com/holiman/uint256"
	"github.com/stretchr/testify/require"

	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/common/log/v3"
	"github.com/erigontech/erigon/db/datadir"
	"github.com/erigontech/erigon/db/kv"
	"github.com/erigontech/erigon/db/kv/temporal/temporaltest"
	"github.com/erigontech/erigon/db/state/execctx"
	"github.com/erigontech/erigon/execution/types/accounts"
	"github.com/erigontech/erigon/node/gointerfaces"
	"github.com/erigontech/erigon/node/gointerfaces/remoteproto"
)

func TestEvictionInUnexpectedOrder(t *testing.T) {
	// Order: View - 2, OnNewBlock - 2, View - 5, View - 6, OnNewBlock - 3, OnNewBlock - 4, View - 5, OnNewBlock - 5, OnNewBlock - 100
	require := require.New(t)
	cfg := DefaultCoherentConfig
	cfg.CacheSize = 3
	cfg.NewBlockWait = 0
	c := New(cfg)
	c.selectOrCreateRoot(2)
	require.Len(c.roots, 1)
	require.Zero(int(c.latestStateVersionID))
	require.False(c.roots[2].isCanonical)

	c.add([]byte{1}, nil, c.roots[2], 2)
	require.Zero(c.stateEvict.Len())

	c.advanceRoot(2)
	require.Len(c.roots, 1)
	require.Equal(2, int(c.latestStateVersionID))
	require.True(c.roots[2].isCanonical)

	c.add([]byte{1}, nil, c.roots[2], 2)
	require.Equal(1, c.stateEvict.Len())

	c.selectOrCreateRoot(5)
	require.Len(c.roots, 2)
	require.Equal(2, int(c.latestStateVersionID))
	require.False(c.roots[5].isCanonical)

	c.add([]byte{2}, nil, c.roots[5], 5) // not added to evict list
	require.Equal(1, c.stateEvict.Len())
	c.add([]byte{2}, nil, c.roots[2], 2) // added to evict list, because it's latest view
	require.Equal(2, c.stateEvict.Len())

	c.selectOrCreateRoot(6)
	require.Len(c.roots, 3)
	require.Equal(2, int(c.latestStateVersionID))
	require.False(c.roots[6].isCanonical) // parrent exists, but parent has isCanonical=false

	c.advanceRoot(3)
	require.Len(c.roots, 4)
	require.Equal(3, int(c.latestStateVersionID))
	require.True(c.roots[3].isCanonical)

	c.advanceRoot(4)
	require.Len(c.roots, 5)
	require.Equal(4, int(c.latestStateVersionID))
	require.True(c.roots[4].isCanonical)

	c.selectOrCreateRoot(5)
	require.Len(c.roots, 5)
	require.Equal(4, int(c.latestStateVersionID))
	require.False(c.roots[5].isCanonical)

	c.advanceRoot(5)
	require.Len(c.roots, 5)
	require.Equal(5, int(c.latestStateVersionID))
	require.True(c.roots[5].isCanonical)

	c.advanceRoot(100)
	require.Len(c.roots, 6)
	require.Equal(100, int(c.latestStateVersionID))
	require.True(c.roots[100].isCanonical)

	//c.add([]byte{1}, nil, c.roots[2], 2)
	require.Equal(0, c.latestStateView.cache.Len())
	require.Equal(0, c.stateEvict.Len())
}

func TestEviction(t *testing.T) {
	require, ctx := require.New(t), t.Context()
	cfg := DefaultCoherentConfig
	cfg.CacheSize = 21
	cfg.NewBlockWait = 0
	c := New(cfg)

	dirs := datadir.New(t.TempDir())
	db := temporaltest.NewTestDB(t, dirs)
	k1, k2 := [20]byte{1}, [20]byte{2}

	var id uint64
	_ = db.UpdateTemporal(ctx, func(tx kv.TemporalRwTx) error {
		id = tx.ViewID()
		var versionID [8]byte
		binary.BigEndian.PutUint64(versionID[:], id)
		_ = tx.Put(kv.Sequence, kv.PlainStateVersion, versionID[:])
		cacheView, _ := c.View(ctx, tx)
		view := cacheView.(*CoherentView)
		_, _ = c.Get(k1[:], tx, view.stateVersionID)
		_, _ = c.Get([]byte{1}, tx, view.stateVersionID)
		_, _ = c.Get([]byte{2}, tx, view.stateVersionID)
		_, _ = c.Get([]byte{3}, tx, view.stateVersionID)
		//require.Equal(c.roots[c.latestViewID].cache.Len(), c.stateEvict.Len())
		return nil
	})
	require.Equal(0, c.stateEvict.Len())
	//require.Equal(c.roots[c.latestViewID].cache.Len(), c.stateEvict.Len())
	c.OnNewBlock(&remoteproto.StateChangeBatch{
		StateVersionId: id + 1,
		ChangeBatch: []*remoteproto.StateChange{
			{
				Direction: remoteproto.Direction_FORWARD,
				Changes: []*remoteproto.AccountChange{{
					Action:  remoteproto.Action_UPSERT,
					Address: gointerfaces.ConvertAddressToH160(k1),
					Data:    []byte{2},
				}},
			},
		},
	})
	require.Equal(21, c.stateEvict.Size())
	require.Equal(1, c.stateEvict.Len())
	require.Equal(c.roots[c.latestStateVersionID].cache.Len(), c.stateEvict.Len())
	_ = db.UpdateTemporal(ctx, func(tx kv.TemporalRwTx) error {
		id = tx.ViewID()
		cacheView, _ := c.View(ctx, tx)
		var versionID [8]byte
		binary.BigEndian.PutUint64(versionID[:], id)
		_ = tx.Put(kv.Sequence, kv.PlainStateVersion, versionID[:])
		view := cacheView.(*CoherentView)
		_, _ = c.Get(k1[:], tx, view.stateVersionID)
		_, _ = c.Get(k2[:], tx, view.stateVersionID)
		_, _ = c.Get([]byte{5}, tx, view.stateVersionID)
		_, _ = c.Get([]byte{6}, tx, view.stateVersionID)
		return nil
	})
	require.Equal(c.roots[c.latestStateVersionID].cache.Len(), c.stateEvict.Len())
	require.Equal(int(cfg.CacheSize.Bytes()), c.stateEvict.Size())
}

func TestAPI(t *testing.T) {
	require := require.New(t)

	// Create a context with timeout for the entire test
	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	defer cancel()

	c := New(DefaultCoherentConfig)
	k1, k2 := [20]byte{1}, [20]byte{2}
	db := temporaltest.NewTestDB(t, datadir.New(t.TempDir()))

	acc := accounts.Account{
		Nonce:       1,
		Balance:     *uint256.NewInt(11),
		CodeHash:    accounts.EmptyCodeHash,
		Incarnation: 2,
	}
	account1Enc := accounts.SerialiseV3(&acc)
	acc.Incarnation = 3
	account2Enc := accounts.SerialiseV3(&acc)
	acc.Incarnation = 5
	account4Enc := accounts.SerialiseV3(&acc)

	get := func(key [20]byte, expectTxnID uint64) (res [1]chan []byte) {
		wg := sync.WaitGroup{}
		for idx := range len(res) {
			res[idx] = make(chan []byte, 1) // Buffered channel to prevent deadlock
			wg.Go(func() {
				err := db.ViewTemporal(ctx, func(tx kv.TemporalTx) error {
					if expectTxnID != tx.ViewID() {
						panic(fmt.Sprintf("epxected: %d, got: %d", expectTxnID, tx.ViewID()))
					}
					cacheView, err := c.View(ctx, tx)
					if err != nil {
						panic(fmt.Sprintf("View error: %v", err))
					}
					view := cacheView.(*CoherentView)
					v, err := c.Get(key[:], tx, view.stateVersionID)
					if err != nil {
						panic(fmt.Sprintf("Get error: %v", err))
					}

					select {
					case res[idx] <- bytes.Clone(v):
					case <-ctx.Done():
						panic("Context done while sending result")
					}
					return nil
				})
				if err != nil {
					panic(fmt.Sprintf("Database error: %v", err))
				}
			})
		}
		wg.Wait() // ensure that all goroutines started their transactions
		return res
	}

	put := func(k, v []byte) uint64 {
		var txID uint64
		err := db.UpdateTemporal(ctx, func(tx kv.TemporalRwTx) error {
			txID = tx.ViewID()
			d, err := execctx.NewSharedDomains(ctx, tx, log.New())
			if err != nil {
				return err
			}
			defer d.Close()
			txNum := uint64(0)
			if err := d.DomainPut(kv.AccountsDomain, tx, k, v, txNum, nil); err != nil {
				return err
			}
			return d.Flush(ctx, tx)
		})
		require.NoError(err)
		return txID
	}
	// block 1 - represents existing state (no notifications about this data will come to client)
	txID1 := put(k2[:], account1Enc)

	wg := sync.WaitGroup{}

	res1, res2 := get(k1, txID1), get(k2, txID1) // will return immediately
	wg.Go(func() {
		for i := range res1 {
			select {
			case v := <-res1[i]:
				if v != nil {
					panic(fmt.Sprintf("expected nil, got: %x", v))
				}
			case <-ctx.Done():
				panic("Context done while checking res1")
			}
		}
		for i := range res2 {
			select {
			case v := <-res2[i]:
				if !bytes.Equal(account1Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account1Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res2")
			}
		}

		fmt.Printf("done1: \n")
	})

	txID2 := put(k1[:], account2Enc)
	fmt.Printf("-----1 %d, %d\n", txID1, txID2)
	res3, res4 := get(k1, txID2), get(k2, txID2) // will see View of transaction 2
	txID3 := put(k1[:], account2Enc)             // even if core already on block 3

	c.OnNewBlock(&remoteproto.StateChangeBatch{
		StateVersionId:      txID2,
		PendingBlockBaseFee: 1,
		ChangeBatch: []*remoteproto.StateChange{
			{
				Direction:   remoteproto.Direction_FORWARD,
				BlockHeight: 2,
				BlockHash:   gointerfaces.ConvertHashToH256([32]byte{}),
				Changes: []*remoteproto.AccountChange{{
					Action:  remoteproto.Action_UPSERT,
					Address: gointerfaces.ConvertAddressToH160(k1),
					Data:    account2Enc,
				}},
			},
		},
	})

	wg.Go(func() {
		for i := range res3 {
			select {
			case v := <-res3[i]:
				if !bytes.Equal(account2Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account2Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res3")
			}
		}
		for i := range res4 {
			select {
			case v := <-res4[i]:
				if !bytes.Equal(account1Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account1Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res4")
			}
		}
		fmt.Printf("done2: \n")
	})
	fmt.Printf("-----2\n")

	res5, res6 := get(k1, txID3), get(k2, txID3) // will see View of transaction 3, even if notification has not enough changes
	c.OnNewBlock(&remoteproto.StateChangeBatch{
		StateVersionId:      txID3,
		PendingBlockBaseFee: 1,
		ChangeBatch: []*remoteproto.StateChange{
			{
				Direction:   remoteproto.Direction_FORWARD,
				BlockHeight: 3,
				BlockHash:   gointerfaces.ConvertHashToH256([32]byte{}),
				Changes: []*remoteproto.AccountChange{{
					Action:  remoteproto.Action_UPSERT,
					Address: gointerfaces.ConvertAddressToH160(k1),
					Data:    account2Enc,
				}},
			},
		},
	})

	wg.Go(func() {
		for i := range res5 {
			select {
			case v := <-res5[i]:
				if !bytes.Equal(account2Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account2Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res5")
			}
		}
		fmt.Printf("-----21\n")

		for i := range res6 {
			select {
			case v := <-res6[i]:
				if !bytes.Equal(account1Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account1Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res6")
			}
		}
		fmt.Printf("done3: \n")
	})
	fmt.Printf("-----3\n")
	txID4 := put(k1[:], account2Enc)

	c.OnNewBlock(&remoteproto.StateChangeBatch{
		StateVersionId:      txID4,
		PendingBlockBaseFee: 1,
		ChangeBatch: []*remoteproto.StateChange{
			{
				Direction:   remoteproto.Direction_UNWIND,
				BlockHeight: 2,
				BlockHash:   gointerfaces.ConvertHashToH256([32]byte{}),
				Changes: []*remoteproto.AccountChange{{
					Action:  remoteproto.Action_UPSERT,
					Address: gointerfaces.ConvertAddressToH160(k1),
					Data:    account4Enc,
				}},
			},
		},
	})
	fmt.Printf("-----4\n")
	txID5 := put(k1[:], account4Enc) // reorg to new chain
	c.OnNewBlock(&remoteproto.StateChangeBatch{
		StateVersionId:      txID5,
		PendingBlockBaseFee: 1,
		ChangeBatch: []*remoteproto.StateChange{
			{
				Direction:   remoteproto.Direction_FORWARD,
				BlockHeight: 3,
				BlockHash:   gointerfaces.ConvertHashToH256([32]byte{2}),
				Changes: []*remoteproto.AccountChange{{
					Action:  remoteproto.Action_UPSERT,
					Address: gointerfaces.ConvertAddressToH160(k1),
					Data:    account4Enc,
				}},
			},
		},
	})
	fmt.Printf("-----5\n")

	res7, res8 := get(k1, txID5), get(k2, txID5) // will see View of transaction 3, even if notification has not enough changes

	wg.Go(func() {
		for i := range res7 {
			select {
			case v := <-res7[i]:
				if !bytes.Equal(account4Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account4Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res7")
			}
		}
		for i := range res8 {
			select {
			case v := <-res8[i]:
				if !bytes.Equal(account1Enc, v) {
					panic(fmt.Sprintf("expected: %x, got: %x", account1Enc, v))
				}
			case <-ctx.Done():
				panic("Context done while checking res8")
			}
		}
		fmt.Printf("done4: \n")
	})
	// TODO: Used in other places too cant modify this.
	// err := db.View(t.Context(), func(tx kv.Tx) error {
	// 	_, err := AssertCheckValues(t.Context(), tx, c)
	// 	require.NoError(err)
	// 	return nil
	// })
	// require.NoError(err)

	// Wait for all goroutines to complete or timeout
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		// All goroutines completed successfully
	case <-ctx.Done():
		t.Error("Test timed out waiting for goroutines to complete")
	}
}

func TestOnNewBlockCodeHashKey(t *testing.T) {
	require := require.New(t)
	cfg := DefaultCoherentConfig
	cfg.NewBlockWait = 0
	c := New(cfg)

	code := []byte{0x01, 0x02, 0x03, 0x04}
	addr := common.Address{0xAA}

	batch := &remoteproto.StateChangeBatch{
		StateVersionId: 1,
		ChangeBatch: []*remoteproto.StateChange{
			{
				Direction: remoteproto.Direction_FORWARD,
				Changes: []*remoteproto.AccountChange{
					{
						Action:  remoteproto.Action_CODE,
						Address: gointerfaces.ConvertAddressToH160(addr),
						Code:    code,
					},
				},
			},
		},
	}

	c.OnNewBlock(batch)

	c.lock.Lock()
	defer c.lock.Unlock()

	require.NotNil(c.latestStateView)
	require.Equal(uint64(1), c.latestStateVersionID)

	var elems []*Element
	c.latestStateView.codeCache.Walk(func(items []*Element) bool {
		if len(items) > 0 {
			elems = append(elems, items...)
		}
		return true
	})

	require.Len(elems, 1)

	h := keccak.NewFastKeccak()
	h.Write(code)
	expectedKey := h.Sum(nil)

	require.Equal(expectedKey, elems[0].K)
	require.Equal(code, elems[0].V)
}

func TestCode(t *testing.T) {
	require, ctx := require.New(t), t.Context()
	c := New(DefaultCoherentConfig)
	db := temporaltest.NewTestDB(t, datadir.New(t.TempDir()))
	k1 := [20]byte{1}
	code := []byte{0x60, 0x00, 0x60, 0x00}

	require.NoError(db.UpdateTemporal(ctx, func(tx kv.TemporalRwTx) error {
		d, err := execctx.NewSharedDomains(ctx, tx, log.New())
		if err != nil {
			return err
		}
		defer d.Close()
		if err := d.DomainPut(kv.CodeDomain, tx, k1[:], code, 0, nil); err != nil {
			return err
		}
		return d.Flush(ctx, tx)
	}))

	require.NoError(db.ViewTemporal(ctx, func(tx kv.TemporalTx) error {
		cacheView, err := c.View(ctx, tx)
		require.NoError(err)
		view := cacheView.(*CoherentView)

		v, err := c.GetCode(k1[:], tx, view.stateVersionID)
		require.NoError(err)
		require.Equal(code, v)

		// second read is served from the cache
		v, err = c.GetCode(k1[:], tx, view.stateVersionID)
		require.NoError(err)
		require.Equal(code, v)
		return nil
	}))
}
