// Copyright 2024 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 beacon_indicies

import (
	"context"
	"errors"
	"testing"

	"github.com/klauspost/compress/zstd"
	"github.com/stretchr/testify/require"

	"github.com/erigontech/erigon/cl/clparams"
	"github.com/erigontech/erigon/cl/cltypes"
	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/db/kv"
	"github.com/erigontech/erigon/db/kv/dbcfg"
	"github.com/erigontech/erigon/db/kv/dbutils"
	"github.com/erigontech/erigon/db/kv/memdb"
)

type errorCursor struct {
	err error
}

func (c errorCursor) First() ([]byte, []byte, error)           { return nil, nil, nil }
func (c errorCursor) Seek([]byte) ([]byte, []byte, error)      { return nil, nil, nil }
func (c errorCursor) SeekExact([]byte) ([]byte, []byte, error) { return nil, nil, nil }
func (c errorCursor) Next() ([]byte, []byte, error)            { return nil, nil, c.err }
func (c errorCursor) Prev() ([]byte, []byte, error)            { return nil, nil, nil }
func (c errorCursor) Last() ([]byte, []byte, error)            { return nil, nil, nil }
func (c errorCursor) Current() ([]byte, []byte, error)         { return nil, nil, nil }
func (c errorCursor) Close()                                   {}

func setupTestDB(t *testing.T) kv.RwDB {
	// Create an in-memory SQLite DB for testing purposes
	db := memdb.NewTestDB(t, dbcfg.ChainDB)
	return db
}

func TestWriteBlockRoot(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	// Mock a block
	block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
	block.Block.Slot = 56
	block.EncodingSizeSSZ()

	require.NoError(t, WriteBeaconBlockHeaderAndIndicies(context.Background(), tx, block.SignedBeaconBlockHeader(), false))

	// Try to retrieve the block's slot by its blockRoot and verify
	blockRoot, err := block.Block.HashSSZ()
	require.NoError(t, err)

	retrievedSlot, err := ReadBlockSlotByBlockRoot(tx, blockRoot)
	require.NoError(t, err)
	require.Equal(t, block.Block.Slot, *retrievedSlot)

	canonicalRoot, err := ReadCanonicalBlockRoot(tx, *retrievedSlot)
	require.NoError(t, err)
	require.Equal(t, common.Hash{}, canonicalRoot)

	err = MarkRootCanonical(context.Background(), tx, *retrievedSlot, blockRoot)
	require.NoError(t, err)

	canonicalRoot, err = ReadCanonicalBlockRoot(tx, *retrievedSlot)
	require.NoError(t, err)
	require.Equal(t, common.Hash(blockRoot), canonicalRoot)
}

func TestWriteBeaconBlockStoresCompleteZstdFrame(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
	block.Block.Slot = 42
	block.EncodingSizeSSZ()

	blockRoot, err := block.Block.HashSSZ()
	require.NoError(t, err)
	require.NoError(t, WriteBeaconBlock(context.Background(), tx, block))

	stored, err := tx.GetOne(kv.BeaconBlocks, dbutils.BlockBodyKey(block.Block.Slot, blockRoot))
	require.NoError(t, err)
	require.NotEmpty(t, stored)

	dec, err := zstd.NewReader(nil)
	require.NoError(t, err)
	defer dec.Close()
	// DecodeAll reads the whole frame to EOF, so it fails on the unterminated frame
	// a Flush()-only encoder leaves behind — pinning that WriteBeaconBlock closes it.
	_, err = dec.DecodeAll(stored, nil)
	require.NoError(t, err)
}

func TestReadParentBlockRoot(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	mockParentRoot := common.Hash{1}
	// Mock a block
	block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
	block.Block.Slot = 56
	block.Block.ParentRoot = mockParentRoot
	block.EncodingSizeSSZ()

	require.NoError(t, WriteBeaconBlockHeaderAndIndicies(context.Background(), tx, block.SignedBeaconBlockHeader(), false))

	// Try to retrieve the block's slot by its blockRoot and verify
	blockRoot, err := block.Block.HashSSZ()
	require.NoError(t, err)

	retrieveParentRoot, err := ReadParentBlockRoot(context.Background(), tx, blockRoot)
	require.NoError(t, err)
	require.Equal(t, mockParentRoot, retrieveParentRoot)
}

func TestTruncateCanonicalChain(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	mockParentRoot := common.Hash{1}
	// Mock a block
	block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
	block.Block.Slot = 56
	block.Block.ParentRoot = mockParentRoot
	block.EncodingSizeSSZ()

	require.NoError(t, WriteBeaconBlockHeaderAndIndicies(context.Background(), tx, block.SignedBeaconBlockHeader(), true))

	// Try to retrieve the block's slot by its blockRoot and verify
	blockRoot, err := block.Block.HashSSZ()
	require.NoError(t, err)

	canonicalRoot, err := ReadCanonicalBlockRoot(tx, block.Block.Slot)
	require.NoError(t, err)
	require.Equal(t, common.Hash(blockRoot), canonicalRoot)

	require.NoError(t, TruncateCanonicalChain(context.Background(), tx, 0))

	canonicalRoot, err = ReadCanonicalBlockRoot(tx, block.Block.Slot)
	require.NoError(t, err)
	require.Equal(t, common.Hash{}, canonicalRoot)
}

func TestReadCanonicalHead(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	slot, root, err := ReadCanonicalHead(tx)
	require.NoError(t, err)
	require.Zero(t, slot)
	require.Equal(t, common.Hash{}, root)

	root10 := common.Hash{0x10}
	root12 := common.Hash{0x12}
	root11 := common.Hash{0x11}
	require.NoError(t, MarkRootCanonical(context.Background(), tx, 10, root10))
	require.NoError(t, MarkRootCanonical(context.Background(), tx, 12, root12))
	require.NoError(t, MarkRootCanonical(context.Background(), tx, 11, root11))

	slot, root, err = ReadCanonicalHead(tx)
	require.NoError(t, err)
	require.Equal(t, uint64(12), slot)
	require.Equal(t, root12, root)
}

func TestReadBeaconBlockHeader(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	mockParentRoot := common.Hash{1}
	mockSignature := [96]byte{23}

	// Mock a block
	block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
	block.Block.Slot = 56
	block.Block.ParentRoot = mockParentRoot
	block.Signature = mockSignature

	canonical := true
	block.EncodingSizeSSZ()

	require.NoError(t, WriteBeaconBlockHeaderAndIndicies(context.Background(), tx, block.SignedBeaconBlockHeader(), canonical))

	// Try to retrieve the block's slot by its blockRoot and verify
	blockRoot, err := block.Block.HashSSZ()
	require.NoError(t, err)

	header, isCanonical, err := ReadSignedHeaderByBlockRoot(context.Background(), tx, blockRoot)
	require.NoError(t, err)
	require.Equal(t, isCanonical, canonical)
	require.NotNil(t, header)

	headerRoot, err := header.Header.HashSSZ()
	require.NoError(t, err)

	require.Equal(t, headerRoot, blockRoot)

}

func TestWriteExecutionBlockNumber(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	tHash := common.HexToHash("0x2")
	require.NoError(t, WriteExecutionBlockNumber(tx, tHash, 1))
	require.NoError(t, WriteExecutionBlockNumber(tx, tHash, 2))
	require.NoError(t, WriteExecutionBlockNumber(tx, tHash, 3))

	// Try to retrieve the block's slot by its blockRoot and verify
	blockNumber, err := ReadExecutionBlockNumber(tx, tHash)
	require.NoError(t, err)
	require.Equal(t, uint64(3), *blockNumber)
}

func TestWriteExecutionBlockHash(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()
	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	tHash := common.HexToHash("0x2")
	tHash2 := common.HexToHash("0x3")
	require.NoError(t, WriteExecutionBlockHash(tx, tHash, tHash2))
	// Try to retrieve the block's slot by its blockRoot and verify
	tHash3, err := ReadExecutionBlockHash(tx, tHash)
	require.NoError(t, err)
	require.Equal(t, tHash2, tHash3)
}

func TestPruneBlocksRemovesAllBlocksBeforeSlot(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()

	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	type storedBlock struct {
		slot uint64
		root common.Hash
	}

	var blocks []storedBlock
	for _, slot := range []uint64{1, 2, 3, 4} {
		block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
		block.Block.Slot = slot
		block.EncodingSizeSSZ()

		root, err := block.Block.HashSSZ()
		require.NoError(t, err)
		require.NoError(t, WriteBeaconBlockAndIndicies(context.Background(), tx, block, false))

		blocks = append(blocks, storedBlock{slot: slot, root: root})
	}

	require.NoError(t, PruneBlocks(context.Background(), tx, 3))

	for _, block := range blocks {
		body, err := tx.GetOne(kv.BeaconBlocks, dbutils.BlockBodyKey(block.slot, block.root))
		require.NoError(t, err)
		if block.slot < 3 {
			require.Empty(t, body)
			continue
		}
		require.NotEmpty(t, body)
	}
}

func TestPruneBlocksLimitRemovesBoundedNumberOfBlocks(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()

	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	type storedBlock struct {
		slot uint64
		root common.Hash
	}

	var blocks []storedBlock
	for _, slot := range []uint64{1, 2, 3, 4} {
		block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
		block.Block.Slot = slot
		block.EncodingSizeSSZ()

		root, err := block.Block.HashSSZ()
		require.NoError(t, err)
		require.NoError(t, WriteBeaconBlockAndIndicies(context.Background(), tx, block, false))

		blocks = append(blocks, storedBlock{slot: slot, root: root})
	}

	deleted, hasMore, err := PruneBlocksLimit(context.Background(), tx, 4, 2)
	require.NoError(t, err)
	require.Equal(t, 2, deleted)
	require.True(t, hasMore)

	for _, block := range blocks {
		body, err := tx.GetOne(kv.BeaconBlocks, dbutils.BlockBodyKey(block.slot, block.root))
		require.NoError(t, err)
		if block.slot < 3 {
			require.Empty(t, body)
			continue
		}
		require.NotEmpty(t, body)
	}

	deleted, hasMore, err = PruneBlocksLimit(context.Background(), tx, 4, 2)
	require.NoError(t, err)
	require.Equal(t, 1, deleted)
	require.False(t, hasMore)
}

func TestPruneBlocksLimitExactLimitHasNoMore(t *testing.T) {
	db := setupTestDB(t)
	defer db.Close()

	tx, err := db.BeginRw(context.Background())
	require.NoError(t, err)
	defer tx.Rollback()

	for _, slot := range []uint64{1, 2} {
		block := cltypes.NewSignedBeaconBlock(&clparams.MainnetBeaconConfig, clparams.Phase0Version)
		block.Block.Slot = slot
		block.EncodingSizeSSZ()
		require.NoError(t, WriteBeaconBlockAndIndicies(context.Background(), tx, block, false))
	}

	deleted, hasMore, err := PruneBlocksLimit(context.Background(), tx, 3, 2)
	require.NoError(t, err)
	require.Equal(t, 2, deleted)
	require.False(t, hasMore)
}

func TestHasMorePrunableBeaconBlocksReturnsCursorError(t *testing.T) {
	wantErr := errors.New("cursor failed")

	hasMore, err := hasMorePrunableBeaconBlocks(errorCursor{err: wantErr}, 10)

	require.False(t, hasMore)
	require.ErrorIs(t, err, wantErr)
}
