// 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 mock_services

import (
	"context"
	"testing"

	"go.uber.org/mock/gomock"

	"github.com/erigontech/erigon/cl/cltypes"
	"github.com/erigontech/erigon/cl/cltypes/solid"
	"github.com/erigontech/erigon/cl/das"
	"github.com/erigontech/erigon/cl/das/mock_services"
	peerdasstatemock "github.com/erigontech/erigon/cl/das/state/mock_services"
	"github.com/erigontech/erigon/cl/phase1/core/state"
	"github.com/erigontech/erigon/cl/phase1/execution_client"
	"github.com/erigontech/erigon/cl/phase1/forkchoice"
	"github.com/erigontech/erigon/cl/pool"
	"github.com/erigontech/erigon/cl/transition/impl/eth2"
	"github.com/erigontech/erigon/cl/validator/sync_contribution_pool"
	syncpoolmock "github.com/erigontech/erigon/cl/validator/sync_contribution_pool/mock_services"
	"github.com/erigontech/erigon/common"
)

// Make mocks with maps and simple setters and getters, panic on methods from ForkChoiceStorageWriter

type ForkChoiceStorageMock struct {
	Ancestors              map[uint64]forkchoice.ForkChoiceNode
	AnchorSlotVal          uint64
	AnchorRootVal          common.Hash
	FinalizedCheckpointVal solid.Checkpoint
	FinalizedSlotVal       uint64
	LowestAvailableSlotVal *uint64
	HeadVal                common.Hash
	HeadSlotVal            uint64
	HeadPayloadStatusVal   cltypes.PayloadStatus
	HighestSeenVal         uint64
	JustifiedCheckpointVal solid.Checkpoint
	JustifiedSlotVal       uint64
	ProposerBoostRootVal   common.Hash
	SlotVal                uint64
	TimeVal                uint64

	ParticipationVal map[uint64]*solid.ParticipationBitList

	IsRootOptimisticVal bool

	StateAtBlockRootVal          map[common.Hash]*state.CachingBeaconState
	StateAtSlotVal               map[uint64]*state.CachingBeaconState
	GetSyncCommitteesVal         map[uint64][2]*solid.SyncCommittee
	GetFinalityCheckpointsVal    map[common.Hash][3]solid.Checkpoint
	PendingConsolidationsVal     map[common.Hash]*solid.ListSSZ[*solid.PendingConsolidation]
	PendingDepositsVal           map[common.Hash]*solid.ListSSZ[*solid.PendingDeposit]
	PendingPartialWithdrawalsVal map[common.Hash]*solid.ListSSZ[*solid.PendingPartialWithdrawal]
	WeightsMock                  []forkchoice.ForkNode
	LightClientBootstraps        map[common.Hash]*cltypes.LightClientBootstrap
	NewestLCUpdate               *cltypes.LightClientUpdate
	LCUpdates                    map[uint64]*cltypes.LightClientUpdate
	SyncContributionPool         sync_contribution_pool.SyncContributionPool
	Headers                      map[common.Hash]*cltypes.BeaconBlockHeader
	Blocks                       map[common.Hash]*cltypes.SignedBeaconBlock
	Envelopes                    map[common.Hash]*cltypes.SignedExecutionPayloadEnvelope
	VerifiedPayloads             map[common.Hash]bool
	OnExecutionPayloadErr        error
	GetBeaconCommitteeMock       func(slot, committeeIndex uint64) ([]uint64, error)

	Pool pool.OperationsPool

	Eth1Hashes map[common.Hash]common.Hash

	// Mock for PeerDas
	MockPeerDas *mock_services.MockPeerDas

	ShouldExtendPayloadVal bool

	// [New in Gloas:EIP7732] Execution payload status by execution block hash
	ExecutionPayloadStatusMap map[common.Hash]execution_client.PayloadStatus
	PayloadStatusByRootMap    map[common.Hash]execution_client.PayloadStatus
	// [New in Gloas:EIP7732] Execution payload gas limit by execution block hash
	ExecutionPayloadGasLimitMap map[common.Hash]uint64
}

func makeSyncContributionPoolMock(t *testing.T) sync_contribution_pool.SyncContributionPool {
	ctrl := gomock.NewController(t)
	type syncContributionKey struct {
		slot              uint64
		subcommitteeIndex uint64
		beaconBlockRoot   common.Hash
	}
	u := map[syncContributionKey]*cltypes.Contribution{}
	pool := syncpoolmock.NewMockSyncContributionPool(ctrl)
	pool.EXPECT().
		AddSyncContribution(gomock.Any(), gomock.Any()).
		DoAndReturn(func(headState *state.CachingBeaconState, contribution *cltypes.Contribution) error {
			key := syncContributionKey{
				slot:              contribution.Slot,
				subcommitteeIndex: contribution.SubcommitteeIndex,
				beaconBlockRoot:   contribution.BeaconBlockRoot,
			}
			u[key] = contribution
			return nil
		}).
		AnyTimes()
	pool.EXPECT().
		GetSyncContribution(gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(slot uint64, subcommitteeIndex uint64, beaconBlockRoot common.Hash) *cltypes.Contribution {
			key := syncContributionKey{
				slot:              slot,
				subcommitteeIndex: subcommitteeIndex,
				beaconBlockRoot:   beaconBlockRoot,
			}
			return u[key]
		}).
		AnyTimes()
	pool.EXPECT().
		AddSyncCommitteeMessage(gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(headState *state.CachingBeaconState, subCommitee uint64, message *cltypes.SyncCommitteeMessage) error {
			key := syncContributionKey{
				slot:              message.Slot,
				subcommitteeIndex: subCommitee,
				beaconBlockRoot:   message.BeaconBlockRoot,
			}
			u[key] = &cltypes.Contribution{
				Slot:              message.Slot,
				SubcommitteeIndex: subCommitee,
				BeaconBlockRoot:   message.BeaconBlockRoot,
				AggregationBits:   make([]byte, cltypes.DefaultSyncCommitteeAggregationBitsSize),
			}
			return nil
		}).AnyTimes()
	pool.EXPECT().
		GetSyncAggregate(gomock.Any(), gomock.Any()).
		Return(cltypes.NewSyncAggregate(), nil).
		AnyTimes()
	return pool
}

func NewForkChoiceStorageMock(t *testing.T) *ForkChoiceStorageMock {
	ctrl := gomock.NewController(t)
	mockPeerDas := mock_services.NewMockPeerDas(ctrl)
	mockPeerDasStateReader := peerdasstatemock.NewMockPeerDasStateReader(ctrl)

	// Set up default expectations for the mock
	mockPeerDas.EXPECT().
		DownloadColumnsAndRecoverBlobs(gomock.Any(), gomock.Any()).
		Return(nil).
		AnyTimes()
	mockPeerDas.EXPECT().
		IsDataAvailable(gomock.Any(), gomock.Any()).
		Return(true, nil).
		AnyTimes()
	mockPeerDas.EXPECT().
		Prune(gomock.Any()).
		Return(nil).
		AnyTimes()
	mockPeerDas.EXPECT().
		UpdateValidatorsCustody(gomock.Any()).
		AnyTimes()
	mockPeerDas.EXPECT().
		TryScheduleRecover(gomock.Any(), gomock.Any()).
		Return(nil).
		AnyTimes()
	mockPeerDas.EXPECT().
		StateReader().
		Return(mockPeerDasStateReader).
		AnyTimes()

	// Set up default expectations for the PeerDasStateReader mock
	mockPeerDasStateReader.EXPECT().
		GetEarliestAvailableSlot().
		Return(uint64(0)).
		AnyTimes()
	mockPeerDasStateReader.EXPECT().
		GetRealCgc().
		Return(uint64(0)).
		AnyTimes()
	mockPeerDasStateReader.EXPECT().
		GetAdvertisedCgc().
		Return(uint64(0)).
		AnyTimes()

	return &ForkChoiceStorageMock{
		Ancestors:                   make(map[uint64]forkchoice.ForkChoiceNode),
		AnchorSlotVal:               0,
		FinalizedCheckpointVal:      solid.Checkpoint{},
		FinalizedSlotVal:            0,
		HeadVal:                     common.Hash{},
		HeadPayloadStatusVal:        cltypes.PayloadStatusFull,
		HighestSeenVal:              0,
		JustifiedCheckpointVal:      solid.Checkpoint{},
		JustifiedSlotVal:            0,
		ProposerBoostRootVal:        common.Hash{},
		SlotVal:                     0,
		TimeVal:                     0,
		StateAtBlockRootVal:         make(map[common.Hash]*state.CachingBeaconState),
		StateAtSlotVal:              make(map[uint64]*state.CachingBeaconState),
		GetSyncCommitteesVal:        make(map[uint64][2]*solid.SyncCommittee),
		GetFinalityCheckpointsVal:   make(map[common.Hash][3]solid.Checkpoint),
		LightClientBootstraps:       make(map[common.Hash]*cltypes.LightClientBootstrap),
		LCUpdates:                   make(map[uint64]*cltypes.LightClientUpdate),
		Headers:                     make(map[common.Hash]*cltypes.BeaconBlockHeader),
		Blocks:                      make(map[common.Hash]*cltypes.SignedBeaconBlock),
		Envelopes:                   make(map[common.Hash]*cltypes.SignedExecutionPayloadEnvelope),
		GetBeaconCommitteeMock:      nil,
		Eth1Hashes:                  make(map[common.Hash]common.Hash),
		ShouldExtendPayloadVal:      true,
		SyncContributionPool:        makeSyncContributionPoolMock(t),
		MockPeerDas:                 mockPeerDas,
		ExecutionPayloadStatusMap:   make(map[common.Hash]execution_client.PayloadStatus),
		PayloadStatusByRootMap:      make(map[common.Hash]execution_client.PayloadStatus),
		ExecutionPayloadGasLimitMap: make(map[common.Hash]uint64),
	}
}

func (f *ForkChoiceStorageMock) GetPeerDas() das.PeerDas {
	return f.MockPeerDas
}

func (f *ForkChoiceStorageMock) Ancestor(root common.Hash, slot uint64) forkchoice.ForkChoiceNode {
	return f.Ancestors[slot]
}

func (f *ForkChoiceStorageMock) AnchorSlot() uint64 {
	return f.AnchorSlotVal
}

func (f *ForkChoiceStorageMock) AnchorRoot() common.Hash {
	return f.AnchorRootVal
}

func (f *ForkChoiceStorageMock) Engine() execution_client.ExecutionEngine {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) FinalizedCheckpoint() solid.Checkpoint {
	return f.FinalizedCheckpointVal
}

func (f *ForkChoiceStorageMock) FinalizedSlot() uint64 {
	return f.FinalizedSlotVal
}

func (f *ForkChoiceStorageMock) GetEth1Hash(eth2Root common.Hash) common.Hash {
	if f.Eth1Hashes != nil {
		return f.Eth1Hashes[eth2Root]
	}
	return common.Hash{}
}

func (f *ForkChoiceStorageMock) GetFinalizedExecutionHash(eth2Root common.Hash) common.Hash {
	return f.GetEth1Hash(eth2Root)
}

func (f *ForkChoiceStorageMock) GetHead(_ *state.CachingBeaconState) (common.Hash, uint64, error) {
	return f.HeadVal, f.HeadSlotVal, nil
}

func (f *ForkChoiceStorageMock) HighestSeen() uint64 {
	return f.HighestSeenVal
}

func (f *ForkChoiceStorageMock) JustifiedCheckpoint() solid.Checkpoint {
	return f.JustifiedCheckpointVal
}

func (f *ForkChoiceStorageMock) JustifiedSlot() uint64 {
	return f.JustifiedSlotVal
}

func (f *ForkChoiceStorageMock) ProposerBoostRoot() common.Hash {
	return f.ProposerBoostRootVal
}

func (f *ForkChoiceStorageMock) GetStateAtBlockRoot(
	blockRoot common.Hash,
	alwaysCopy bool,
) (*state.CachingBeaconState, error) {
	st := f.StateAtBlockRootVal[blockRoot]
	if st == nil || !alwaysCopy {
		return st, nil
	}
	return st.Copy()
}

func (f *ForkChoiceStorageMock) GetFinalityCheckpoints(
	blockRoot common.Hash,
) (solid.Checkpoint, solid.Checkpoint, solid.Checkpoint, bool) {
	oneNil := f.GetFinalityCheckpointsVal[blockRoot][0] != solid.Checkpoint{} &&
		f.GetFinalityCheckpointsVal[blockRoot][1] != solid.Checkpoint{} &&
		f.GetFinalityCheckpointsVal[blockRoot][2] != solid.Checkpoint{}

	return f.GetFinalityCheckpointsVal[blockRoot][0], f.GetFinalityCheckpointsVal[blockRoot][1], f.GetFinalityCheckpointsVal[blockRoot][2], oneNil
}

func (f *ForkChoiceStorageMock) GetSyncCommittees(
	period uint64,
) (*solid.SyncCommittee, *solid.SyncCommittee, bool) {
	return f.GetSyncCommitteesVal[period][0], f.GetSyncCommitteesVal[period][1], f.GetSyncCommitteesVal[period][0] != nil &&
		f.GetSyncCommitteesVal[period][1] != nil
}

func (f *ForkChoiceStorageMock) Slot() uint64 {
	return f.SlotVal
}

func (f *ForkChoiceStorageMock) Time() uint64 {
	return f.TimeVal
}

func (f *ForkChoiceStorageMock) OnAttestation(
	attestation *solid.Attestation,
	fromBlock, insert bool,
) error {
	f.Pool.AttestationsPool.Insert(attestation.Signature, attestation)
	return nil
}

func (f *ForkChoiceStorageMock) OnAttesterSlashing(
	attesterSlashing *cltypes.AttesterSlashing,
	test bool,
) error {
	f.Pool.AttesterSlashingsPool.Insert(
		pool.ComputeKeyForAttesterSlashing(attesterSlashing),
		attesterSlashing,
	)
	return nil
}

func (f *ForkChoiceStorageMock) OnBlock(
	ctx context.Context,
	block *cltypes.SignedBeaconBlock,
	newPayload bool,
	fullValidation bool,
	checkDataAvaiability bool,
) error {
	return nil
}

func (f *ForkChoiceStorageMock) OnExecutionPayload(ctx context.Context, signedEnvelope *cltypes.SignedExecutionPayloadEnvelope, checkBlobData, validatePayload bool) error {
	return f.OnExecutionPayloadErr
}

func (f *ForkChoiceStorageMock) ApplyLocalSelfBuildEnvelope(ctx context.Context, signedEnvelope *cltypes.SignedExecutionPayloadEnvelope) error {
	return nil
}

func (f *ForkChoiceStorageMock) StoreAnchorEnvelope(blockRoot common.Hash, signedEnvelope *cltypes.SignedExecutionPayloadEnvelope) error {
	return nil
}

func (f *ForkChoiceStorageMock) OnPayloadAttestationMessage(msg *cltypes.PayloadAttestationMessage, isFromBlock bool) error {
	return nil
}

func (f *ForkChoiceStorageMock) OnTick(time uint64) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) BlockRewards(root common.Hash) (*eth2.BlockRewardsCollector, bool) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) TotalActiveBalance(root common.Hash) (uint64, bool) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) RandaoMixes(blockRoot common.Hash, out solid.HashListSSZ) bool {
	return false
}

func (f *ForkChoiceStorageMock) LowestAvailableSlot() uint64 {
	if f.LowestAvailableSlotVal != nil {
		return *f.LowestAvailableSlotVal
	}
	return f.FinalizedSlotVal
}

func (f *ForkChoiceStorageMock) Participation(epoch uint64) (*solid.ParticipationBitList, bool) {
	if f.ParticipationVal == nil {
		return nil, false
	}
	p, ok := f.ParticipationVal[epoch]
	return p, ok
}

func (f *ForkChoiceStorageMock) ForkNodes() []forkchoice.ForkNode {
	return f.WeightsMock
}

func (f *ForkChoiceStorageMock) Synced() bool {
	return true
}

func (f *ForkChoiceStorageMock) SetSynced(synced bool) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) GetLightClientBootstrap(
	blockRoot common.Hash,
) (*cltypes.LightClientBootstrap, bool) {
	return f.LightClientBootstraps[blockRoot], f.LightClientBootstraps[blockRoot] != nil
}

func (f *ForkChoiceStorageMock) NewestLightClientUpdate() *cltypes.LightClientUpdate {
	return f.NewestLCUpdate
}

func (f *ForkChoiceStorageMock) GetLightClientUpdate(
	period uint64,
) (*cltypes.LightClientUpdate, bool) {
	return f.LCUpdates[period], f.LCUpdates[period] != nil
}

func (f *ForkChoiceStorageMock) GetHeader(
	blockRoot common.Hash,
) (*cltypes.BeaconBlockHeader, bool) {
	return f.Headers[blockRoot], f.Headers[blockRoot] != nil
}

func (f *ForkChoiceStorageMock) GetBlock(
	blockRoot common.Hash,
) (*cltypes.SignedBeaconBlock, bool) {
	return f.Blocks[blockRoot], f.Blocks[blockRoot] != nil
}

func (f *ForkChoiceStorageMock) HasEnvelope(blockRoot common.Hash) bool {
	_, ok := f.Envelopes[blockRoot]
	return ok
}

func (f *ForkChoiceStorageMock) IsPayloadVerified(blockRoot common.Hash) bool {
	if f.VerifiedPayloads == nil {
		return false
	}
	return f.VerifiedPayloads[blockRoot]
}

func (f *ForkChoiceStorageMock) ReadEnvelopeFromDisk(blockRoot common.Hash) (*cltypes.SignedExecutionPayloadEnvelope, error) {
	return f.Envelopes[blockRoot], nil
}

func (f *ForkChoiceStorageMock) IsBlobDataAvailable(slot uint64, blockRoot common.Hash) bool {
	return true
}

func (f *ForkChoiceStorageMock) GetHeadPayloadStatus() cltypes.PayloadStatus {
	return f.HeadPayloadStatusVal
}

func (f *ForkChoiceStorageMock) ShouldExtendPayload(root common.Hash) bool {
	return f.ShouldExtendPayloadVal
}

func (f *ForkChoiceStorageMock) ShouldBuildOnFull(head forkchoice.ForkChoiceNode) bool {
	return true
}

func (f *ForkChoiceStorageMock) GetBalances(blockRoot common.Hash) (solid.Uint64ListSSZ, error) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) GetInactivitiesScores(
	blockRoot common.Hash,
) (solid.Uint64ListSSZ, error) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) GetPreviousParticipationIndicies(
	blockRoot common.Hash,
) (*solid.ParticipationBitList, error) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) GetValidatorSet(
	blockRoot common.Hash,
) (*solid.ValidatorSet, error) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) GetCurrentParticipationIndicies(
	blockRoot common.Hash,
) (*solid.ParticipationBitList, error) {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) GetPublicKeyForValidator(
	blockRoot common.Hash,
	idx uint64,
) (common.Bytes48, error) {
	panic("implement me")
}

// func (f *ForkChoiceStorageMock) OnSignedContributionAndProof(signedContribution *cltypes.SignedContributionAndProof, test bool) error {
// 	f.SyncContributionPool.AddSyncContribution(nil, signedContribution.Message.Contribution)
// 	return nil
// }

func (f *ForkChoiceStorageMock) AddPreverifiedBlobSidecar(msg *cltypes.BlobSidecar) error {
	return nil
}
func (f *ForkChoiceStorageMock) ValidateOnAttestation(attestation *solid.Attestation) error {
	panic("implement me")
}

func (f *ForkChoiceStorageMock) ProcessAttestingIndicies(
	attestation *solid.Attestation,
	attestionIndicies []uint64,
) {
}

func (f *ForkChoiceStorageMock) IsRootOptimistic(root common.Hash) bool {
	return f.IsRootOptimisticVal
}

func (f *ForkChoiceStorageMock) IsHeadOptimistic() bool {
	return false
}

func (f *ForkChoiceStorageMock) GetPendingConsolidations(blockRoot common.Hash) (*solid.ListSSZ[*solid.PendingConsolidation], bool) {
	v, ok := f.PendingConsolidationsVal[blockRoot]
	return v, ok
}

func (f *ForkChoiceStorageMock) GetPendingDeposits(blockRoot common.Hash) (*solid.ListSSZ[*solid.PendingDeposit], bool) {
	v, ok := f.PendingDepositsVal[blockRoot]
	return v, ok
}

func (f *ForkChoiceStorageMock) GetPendingPartialWithdrawals(blockRoot common.Hash) (*solid.ListSSZ[*solid.PendingPartialWithdrawal], bool) {
	v, ok := f.PendingPartialWithdrawalsVal[blockRoot]
	return v, ok
}

func (f *ForkChoiceStorageMock) GetProposerLookahead(slot uint64) (solid.Uint64VectorSSZ, bool) {
	return nil, false
}

// GetRecentExecutionPayloadStatus returns the validation status of a recently validated execution payload.
// [New in Gloas:EIP7732]
func (f *ForkChoiceStorageMock) GetRecentExecutionPayloadStatus(executionBlockHash common.Hash) (execution_client.PayloadStatus, bool) {
	status, ok := f.ExecutionPayloadStatusMap[executionBlockHash]
	return status, ok
}

func (f *ForkChoiceStorageMock) GetRecentExecutionPayloadStatusByRoot(blockRoot common.Hash) (execution_client.PayloadStatus, bool) {
	status, ok := f.PayloadStatusByRootMap[blockRoot]
	return status, ok
}

// GetExecutionPayloadGasLimit returns the gas_limit of a recently validated execution payload.
// [New in Gloas:EIP7732]
func (f *ForkChoiceStorageMock) GetExecutionPayloadGasLimit(executionBlockHash common.Hash) (uint64, bool) {
	gasLimit, ok := f.ExecutionPayloadGasLimitMap[executionBlockHash]
	return gasLimit, ok
}
