package pruningmanager

import (
	"github.com/golang/protobuf/proto"
	"github.com/kaspanet/kaspad/domain/consensus/model"
	"github.com/kaspanet/kaspad/domain/consensus/model/externalapi"
	"github.com/kaspanet/kaspad/domain/consensus/utils/utxoserialization"
)

// pruningManager resolves and manages the current pruning point
type pruningManager struct {
	databaseContext model.DBReader

	dagTraversalManager    model.DAGTraversalManager
	dagTopologyManager     model.DAGTopologyManager
	consensusStateManager  model.ConsensusStateManager
	consensusStateStore    model.ConsensusStateStore
	ghostdagDataStore      model.GHOSTDAGDataStore
	pruningStore           model.PruningStore
	blockStatusStore       model.BlockStatusStore
	headerSelectedTipStore model.HeaderSelectedTipStore

	multiSetStore       model.MultisetStore
	acceptanceDataStore model.AcceptanceDataStore
	blocksStore         model.BlockStore
	utxoDiffStore       model.UTXODiffStore

	genesisHash      *externalapi.DomainHash
	finalityInterval uint64
	pruningDepth     uint64
}

// New instantiates a new PruningManager
func New(
	databaseContext model.DBReader,

	dagTraversalManager model.DAGTraversalManager,
	dagTopologyManager model.DAGTopologyManager,
	consensusStateManager model.ConsensusStateManager,
	consensusStateStore model.ConsensusStateStore,
	ghostdagDataStore model.GHOSTDAGDataStore,
	pruningStore model.PruningStore,
	blockStatusStore model.BlockStatusStore,
	headerSelectedTipStore model.HeaderSelectedTipStore,

	multiSetStore model.MultisetStore,
	acceptanceDataStore model.AcceptanceDataStore,
	blocksStore model.BlockStore,
	utxoDiffStore model.UTXODiffStore,

	genesisHash *externalapi.DomainHash,
	finalityInterval uint64,
	pruningDepth uint64,
) model.PruningManager {

	return &pruningManager{
		databaseContext:        databaseContext,
		dagTraversalManager:    dagTraversalManager,
		dagTopologyManager:     dagTopologyManager,
		consensusStateManager:  consensusStateManager,
		consensusStateStore:    consensusStateStore,
		ghostdagDataStore:      ghostdagDataStore,
		pruningStore:           pruningStore,
		blockStatusStore:       blockStatusStore,
		multiSetStore:          multiSetStore,
		acceptanceDataStore:    acceptanceDataStore,
		blocksStore:            blocksStore,
		utxoDiffStore:          utxoDiffStore,
		headerSelectedTipStore: headerSelectedTipStore,
		genesisHash:            genesisHash,
		pruningDepth:           pruningDepth,
		finalityInterval:       finalityInterval,
	}
}

// FindNextPruningPoint finds the next pruning point from the
// given blockHash
func (pm *pruningManager) UpdatePruningPointByVirtual() error {
	hasPruningPoint, err := pm.pruningStore.HasPruningPoint(pm.databaseContext)
	if err != nil {
		return err
	}

	if !hasPruningPoint {
		err = pm.savePruningPoint(pm.genesisHash)
		if err != nil {
			return err
		}
	}

	currentP, err := pm.pruningStore.PruningPoint(pm.databaseContext)
	if err != nil {
		return err
	}

	virtual, err := pm.ghostdagDataStore.Get(pm.databaseContext, model.VirtualBlockHash)
	if err != nil {
		return err
	}

	virtualSelectedParent, err := pm.ghostdagDataStore.Get(pm.databaseContext, virtual.SelectedParent())
	if err != nil {
		return err
	}

	currentPGhost, err := pm.ghostdagDataStore.Get(pm.databaseContext, currentP)
	if err != nil {
		return err
	}
	currentPBlueScore := currentPGhost.BlueScore()
	// Because the pruning point changes only once per finality, then there's no need to even check for that if a finality interval hasn't passed.
	if virtualSelectedParent.BlueScore() <= currentPBlueScore+pm.finalityInterval {
		return nil
	}

	// This means the pruning point is still genesis.
	if virtualSelectedParent.BlueScore() <= pm.pruningDepth+pm.finalityInterval {
		return nil
	}

	// get Virtual(pruningDepth)
	newPruningPoint, err := pm.calculatePruningPointFromBlock(model.VirtualBlockHash)
	if err != nil {
		return err
	}

	if *newPruningPoint != *currentP {
		err = pm.savePruningPoint(newPruningPoint)
		if err != nil {
			return err
		}
		return pm.deletePastBlocks(newPruningPoint)
	}

	return nil
}

func (pm *pruningManager) deletePastBlocks(pruningPoint *externalapi.DomainHash) error {
	// Go over all P.Past and P.AC that's not in V.Past
	queue := pm.dagTraversalManager.NewDownHeap()

	// Find P.AC that's not in V.Past
	dagTips, err := pm.consensusStateStore.Tips(pm.databaseContext)
	if err != nil {
		return err
	}
	for _, tip := range dagTips {
		hasPruningPointInPast, err := pm.dagTopologyManager.IsAncestorOf(pruningPoint, tip)
		if err != nil {
			return err
		}
		if !hasPruningPointInPast {
			isInVirtualPast, err := pm.dagTopologyManager.IsAncestorOf(model.VirtualBlockHash, tip)
			if err != nil {
				return err
			}
			if !isInVirtualPast {
				// Add them to the queue so they and their past will be pruned
				err := queue.Push(tip)
				if err != nil {
					return err
				}
			}
		}
	}

	// Add P.Parents
	parents, err := pm.dagTopologyManager.Parents(pruningPoint)
	if err != nil {
		return err
	}
	for _, parent := range parents {
		err = queue.Push(parent)
		if err != nil {
			return err
		}
	}

	visited := map[externalapi.DomainHash]struct{}{}
	// Prune everything in the queue including its past
	for queue.Len() > 0 {
		current := queue.Pop()
		if _, ok := visited[*current]; ok {
			continue
		}
		visited[*current] = struct{}{}

		alreadyPruned, err := pm.deleteBlock(current)
		if err != nil {
			return err
		}
		if !alreadyPruned {
			parents, err := pm.dagTopologyManager.Parents(current)
			if err != nil {
				return err
			}
			for _, parent := range parents {
				err = queue.Push(parent)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func (pm *pruningManager) savePruningPoint(blockHash *externalapi.DomainHash) error {
	utxoIter, err := pm.consensusStateManager.RestorePastUTXOSetIterator(blockHash)
	if err != nil {
		return err
	}
	serializedUtxo, err := serializeUTXOSetIterator(utxoIter)
	if err != nil {
		return err
	}
	pm.pruningStore.Stage(blockHash, serializedUtxo)

	return nil
}

func (pm *pruningManager) deleteBlock(blockHash *externalapi.DomainHash) (alreadyPruned bool, err error) {
	status, err := pm.blockStatusStore.Get(pm.databaseContext, blockHash)
	if err != nil {
		return false, err
	}
	if status == externalapi.StatusHeaderOnly {
		return true, nil
	}

	pm.multiSetStore.Delete(blockHash)
	pm.acceptanceDataStore.Delete(blockHash)
	pm.blocksStore.Delete(blockHash)
	pm.utxoDiffStore.Delete(blockHash)

	pm.blockStatusStore.Stage(blockHash, externalapi.StatusHeaderOnly)
	return false, nil
}

func (pm *pruningManager) CalculatePruningPointByHeaderSelectedTip() (*externalapi.DomainHash, error) {
	headersSelectedTip, err := pm.headerSelectedTipStore.HeadersSelectedTip(pm.databaseContext)
	if err != nil {
		return nil, err
	}

	return pm.calculatePruningPointFromBlock(headersSelectedTip)
}

func (pm *pruningManager) calculatePruningPointFromBlock(blockHash *externalapi.DomainHash) (*externalapi.DomainHash, error) {
	ghostdagData, err := pm.ghostdagDataStore.Get(pm.databaseContext, blockHash)
	if err != nil {
		return nil, err
	}

	targetBlueScore := uint64(0)
	if ghostdagData.BlueScore() > pm.pruningDepth {
		// The target blue is calculated by calculating ghostdagData.BlueScore() - pm.pruningDepth and rounding
		// down with the precision of finality interval.
		targetBlueScore = ((ghostdagData.BlueScore() - pm.pruningDepth) / pm.finalityInterval) * pm.finalityInterval
	}
	return pm.dagTraversalManager.LowestChainBlockAboveOrEqualToBlueScore(blockHash, targetBlueScore)
}

func serializeUTXOSetIterator(iter model.ReadOnlyUTXOSetIterator) ([]byte, error) {
	serializedUtxo, err := utxoserialization.ReadOnlyUTXOSetToProtoUTXOSet(iter)
	if err != nil {
		return nil, err
	}
	return proto.Marshal(serializedUtxo)
}