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

import (
	"bytes"
	"errors"
	"math/rand"
	"sort"
	"time"

	"github.com/erigontech/erigon/cl/clparams"
	"github.com/erigontech/erigon/cl/cltypes"
	"github.com/erigontech/erigon/cl/cltypes/solid"
	"github.com/erigontech/erigon/cl/phase1/core/state"
	"github.com/erigontech/erigon/common"
	log "github.com/erigontech/erigon/common/log/v3"
)

// accountWeights updates the weights of the validators, given the vote and given an head leaf.
func (f *ForkChoiceStore) accountWeights(votes, weights map[common.Hash]uint64, justifedRoot, leaf common.Hash) {
	curr := leaf
	accumulated := uint64(0)
	for curr != justifedRoot {
		accumulated += votes[curr]
		votes[curr] = 0 // make sure we don't double count
		weights[curr] += accumulated
		header, has := f.forkGraph.GetHeader(curr)
		if !has {
			return
		}
		curr = header.ParentRoot
	}
}

const (
	sampleFactor = 100
	sampleBasis  = 80
)

func (f *ForkChoiceStore) computeVotes(justifiedCheckpoint solid.Checkpoint, checkpointState *checkpointState, auxilliaryState *state.CachingBeaconState) map[common.Hash]uint64 {
	votes := make(map[common.Hash]uint64)
	// make an rng generator
	gen := rand.New(rand.NewSource(time.Now().UnixNano()))
	if auxilliaryState != nil {
		startIdx := 0
		step := 1
		if f.probabilisticHeadGetter {
			startIdx = gen.Intn(sampleBasis)
			step = sampleBasis + gen.Intn(sampleFactor)
		}
		for validatorIndex := startIdx; validatorIndex < f.latestMessages.latestMessagesCount(); validatorIndex += step {
			message, _ := f.latestMessages.get(validatorIndex)
			v := auxilliaryState.ValidatorSet().Get(validatorIndex)
			if !v.Active(justifiedCheckpoint.Epoch) || v.Slashed() {
				continue
			}
			if _, hasLatestMessage := f.getLatestMessage(uint64(validatorIndex)); !hasLatestMessage || f.isUnequivocating(uint64(validatorIndex)) {
				continue
			}
			votes[message.Root] += v.EffectiveBalance()
		}
		boostRoot := f.proposerBoostRoot.Load().(common.Hash)
		if boostRoot != (common.Hash{}) {
			boost := auxilliaryState.GetTotalActiveBalance() / auxilliaryState.BeaconConfig().SlotsPerEpoch
			votes[boostRoot] += (boost * auxilliaryState.BeaconConfig().ProposerScoreBoost) / 100
		}
	} else {
		for validatorIndex := 0; validatorIndex < f.latestMessages.latestMessagesCount(); validatorIndex++ {
			message, _ := f.latestMessages.get(validatorIndex)
			if message == (LatestMessage{}) {
				continue
			}
			if !readFromBitset(checkpointState.actives, validatorIndex) || readFromBitset(checkpointState.slasheds, validatorIndex) {
				continue
			}
			if _, hasLatestMessage := f.getLatestMessage(uint64(validatorIndex)); !hasLatestMessage {
				continue
			}
			if f.isUnequivocating(uint64(validatorIndex)) {
				continue
			}
			votes[message.Root] += checkpointState.balances[validatorIndex]
		}
		boostRoot := f.proposerBoostRoot.Load().(common.Hash)
		if boostRoot != (common.Hash{}) {
			boost := checkpointState.activeBalance / checkpointState.beaconConfig.SlotsPerEpoch
			votes[boostRoot] += (boost * checkpointState.beaconConfig.ProposerScoreBoost) / 100
		}
	}

	return votes
}

// GetHead returns the head of the fork choice store.
// Dispatches to GLOAS or pre-GLOAS implementation based on current epoch.
func (f *ForkChoiceStore) GetHead(auxilliaryState *state.CachingBeaconState) (common.Hash, uint64, error) {
	f.mu.RLock()
	if f.headHash != (common.Hash{}) {
		headHash, headSlot := f.headHash, f.headSlot
		f.mu.RUnlock()
		return headHash, headSlot, nil
	}
	f.mu.RUnlock()

	// After checkpoint sync the justified root may predate the anchor and
	// not exist in the fork graph. Both getHead and getHeadGloas depend on
	// the justified root being present (for getCheckpointState and
	// getFilteredBlockTree respectively). Return the anchor as head until
	// forward sync produces a resolvable justified checkpoint.
	justifiedCheckpoint := f.justifiedCheckpoint.Load().(solid.Checkpoint)
	if _, hasJustified := f.forkGraph.GetHeader(justifiedCheckpoint.Root); !hasJustified {
		log.Debug("GetHead: justified root not in fork graph, using anchor as head",
			"justifiedRoot", justifiedCheckpoint.Root)
		return f.forkGraph.AnchorRoot(), f.forkGraph.AnchorSlot(), nil
	}

	currentEpoch := f.computeEpochAtSlot(f.Slot())
	if f.beaconCfg.GetCurrentStateVersion(currentEpoch) >= clparams.GloasVersion {
		return f.getHeadGloas()
	}
	return f.getHead(auxilliaryState)
}

// GetHeadPayloadStatus returns the payload status of the current head node.
// Must be called after GetHead has been called (head is cached).
func (f *ForkChoiceStore) GetHeadPayloadStatus() cltypes.PayloadStatus {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.headPayloadStatus
}

// getHeadGloas returns the head using GLOAS fork choice rules.
// [New in Gloas:EIP7732]
func (f *ForkChoiceStore) getHeadGloas() (common.Hash, uint64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	justifiedCheckpoint := f.justifiedCheckpoint.Load().(solid.Checkpoint)

	// Get filtered block tree
	blocks := f.getFilteredBlockTree(justifiedCheckpoint.Root)

	// Start from justified checkpoint with PENDING status
	head := ForkChoiceNode{
		Root:          justifiedCheckpoint.Root,
		PayloadStatus: cltypes.PayloadStatusPending,
	}

	// Get weight store for weight calculation
	ws := f.GetWeightStore()

	for {
		children := f.getNodeChildren(head, blocks)
		if len(children) == 0 {
			// No children, head is the result
			header, hasHeader := f.forkGraph.GetHeader(head.Root)
			if !hasHeader {
				return common.Hash{}, 0, errors.New("no slot for head is stored")
			}
			f.headHash = head.Root
			f.headSlot = header.Slot
			f.headPayloadStatus = head.PayloadStatus
			return f.headHash, f.headSlot, nil
		}

		// Find best child: max(children, key=(weight, root, tiebreaker))
		bestChild := children[0]
		bestWeight := ws.GetWeight(bestChild)
		bestTiebreaker := f.getPayloadStatusTiebreaker(bestChild)

		for i := 1; i < len(children); i++ {
			child := children[i]
			weight := ws.GetWeight(child)
			tiebreaker := f.getPayloadStatusTiebreaker(child)

			// Compare: weight first, then root, then tiebreaker
			if weight > bestWeight {
				bestChild = child
				bestWeight = weight
				bestTiebreaker = tiebreaker
			} else if weight == bestWeight {
				// Compare by root (lexicographically greater wins)
				rootCmp := bytes.Compare(child.Root[:], bestChild.Root[:])
				if rootCmp > 0 {
					bestChild = child
					bestWeight = weight
					bestTiebreaker = tiebreaker
				} else if rootCmp == 0 {
					// Same root, compare by tiebreaker
					if tiebreaker > bestTiebreaker {
						bestChild = child
						bestWeight = weight
						bestTiebreaker = tiebreaker
					}
				}
			}
		}

		head = bestChild
	}
}

// getHead returns the head using pre-GLOAS fork choice rules.
func (f *ForkChoiceStore) getHead(auxilliaryState *state.CachingBeaconState) (common.Hash, uint64, error) {
	justifiedCheckpoint := f.justifiedCheckpoint.Load().(solid.Checkpoint)
	var justificationState *checkpointState
	var err error
	if auxilliaryState == nil {
		// Compute checkpoint state BEFORE acquiring f.mu: this operation
		// can read a large state from disk, which would block the OnTick
		// goroutine if done under the lock.
		justificationState, err = f.getCheckpointState(justifiedCheckpoint)
		if err != nil {
			return common.Hash{}, 0, err
		}
	}
	f.mu.Lock()
	defer f.mu.Unlock()

	// Retrieve att
	f.headHash = justifiedCheckpoint.Root
	blocks := f.getFilteredBlockTree(f.headHash)
	// Do a simple scan to determine the fork votes.
	votes := f.computeVotes(justifiedCheckpoint, justificationState, auxilliaryState)
	// Account for weights on each head fork
	f.weights = make(map[common.Hash]uint64)
	for head := range f.headSet {
		f.accountWeights(votes, f.weights, justifiedCheckpoint.Root, head)
	}

	for {
		// Filter out current head children.
		unfilteredChildren := f.children(f.headHash)
		children := []common.Hash{}
		for _, child := range unfilteredChildren {
			if _, ok := blocks[child]; ok {
				children = append(children, child)
			}
		}
		// Stop if we dont have any more children
		if len(children) == 0 {
			header, hasHeader := f.forkGraph.GetHeader(f.headHash)
			if !hasHeader {
				return common.Hash{}, 0, errors.New("no slot for head is stored")
			}
			f.headSlot = header.Slot
			return f.headHash, f.headSlot, nil
		}

		// Average case scenario.
		if len(children) == 1 {
			f.headHash = children[0]
			continue
		}
		// Sort children by lexigographical order
		sort.Slice(children, func(i, j int) bool {
			childA := children[i]
			childB := children[j]
			return bytes.Compare(childA[:], childB[:]) < 0
		})
		// After sorting is done determine best fit.
		f.headHash = children[0]
		maxWeight := f.weights[children[0]]
		for i := 1; i < len(children); i++ {
			weight := f.weights[children[i]]
			// Lexicographical order is king.
			if weight >= maxWeight {
				f.headHash = children[i]
				maxWeight = weight
			}
		}
	}
}

// getFilteredBlockTree filters out dumb blocks.
func (f *ForkChoiceStore) getFilteredBlockTree(base common.Hash) map[common.Hash]*cltypes.BeaconBlockHeader {
	blocks := make(map[common.Hash]*cltypes.BeaconBlockHeader)
	// Snapshot the store epoch once for the whole walk: OnTick updates f.time
	// without holding f.mu, so calling f.Slot() per-leaf can mix epochs across
	// leaves around a slot boundary.
	currentEpoch := f.computeEpochAtSlot(f.Slot())
	f.getFilterBlockTree(base, blocks, currentEpoch)
	return blocks
}

// getFilterBlockTree recursively traverses the block tree to identify viable blocks.
// It takes a block hash and a map of viable blocks as input parameters, and returns a boolean value indicating
// whether the current block is viable.
func (f *ForkChoiceStore) getFilterBlockTree(blockRoot common.Hash, blocks map[common.Hash]*cltypes.BeaconBlockHeader, currentEpoch uint64) bool {
	header, has := f.forkGraph.GetHeader(blockRoot)
	if !has {
		return false
	}
	finalizedCheckpoint := f.finalizedCheckpoint.Load().(solid.Checkpoint)
	justifiedCheckpoint := f.justifiedCheckpoint.Load().(solid.Checkpoint)
	children := f.children(blockRoot)
	// If there are children iterate down recursively and see which branches are viable.
	if len(children) > 0 {
		isAnyViable := false
		for _, child := range children {
			if f.getFilterBlockTree(child, blocks, currentEpoch) {
				isAnyViable = true
			}
		}
		if isAnyViable {
			blocks[blockRoot] = header
		}
		return isAnyViable
	}
	// Leaf node — viability per spec filter_block_tree.

	// Spec get_voting_source: pick unrealized only for prior-epoch blocks
	// (pull-up justification view); current/future-epoch blocks use the
	// block's realized state checkpoint.
	blockEpoch := f.computeEpochAtSlot(header.Slot)
	var votingSource solid.Checkpoint
	if currentEpoch > blockEpoch {
		votingSource, has = f.getUnrealizedJustification(blockRoot)
	} else {
		votingSource, has = f.forkGraph.GetCurrentJustifiedCheckpoint(blockRoot)
	}
	if !has {
		return false
	}

	genesisEpoch := f.beaconCfg.GenesisEpoch

	// Spec correct_justified:
	//   store.justified_checkpoint.epoch == GENESIS_EPOCH
	//   || voting_source.epoch == store.justified_checkpoint.epoch
	//   || voting_source.epoch + 2 >= current_epoch
	justifiedOk := justifiedCheckpoint.Epoch == genesisEpoch ||
		votingSource.Epoch == justifiedCheckpoint.Epoch ||
		votingSource.Epoch+2 >= currentEpoch

	// Spec correct_finalized:
	//   store.finalized_checkpoint.epoch == GENESIS_EPOCH
	//   || store.finalized_checkpoint.root == get_checkpoint_block(block_root, finalized.epoch)
	finalizedOk := finalizedCheckpoint.Epoch == genesisEpoch
	if !finalizedOk {
		finalizedSlot := f.computeStartSlotAtEpoch(finalizedCheckpoint.Epoch)
		finalizedOk = finalizedCheckpoint.Root == f.Ancestor(blockRoot, finalizedSlot).Root
	}

	if justifiedOk && finalizedOk {
		blocks[blockRoot] = header
		return true
	}
	return false
}
