package messagepool

import (
	"context"
	"sort"
	"time"

	"github.com/filecoin-project/go-address"
	"github.com/filecoin-project/lotus/chain/types"
	"github.com/ipfs/go-cid"
	"golang.org/x/xerrors"
)

func (mp *MessagePool) pruneExcessMessages() error {
	mp.curTsLk.Lock()
	ts := mp.curTs
	mp.curTsLk.Unlock()

	mp.lk.Lock()
	defer mp.lk.Unlock()

	if mp.currentSize < mp.cfg.SizeLimitHigh {
		return nil
	}

	return mp.pruneMessages(context.TODO(), ts)
}

func (mp *MessagePool) pruneMessages(ctx context.Context, ts *types.TipSet) error {
	start := time.Now()
	defer func() {
		log.Infof("message pruning took %s", time.Since(start))
	}()

	baseFee, err := mp.api.ChainComputeBaseFee(ctx, ts)
	if err != nil {
		return xerrors.Errorf("computing basefee: %w", err)
	}

	pending, _ := mp.getPendingMessages(ts, ts)

	// priority actors -- not pruned
	priority := make(map[address.Address]struct{})
	for _, actor := range mp.cfg.PriorityAddrs {
		priority[actor] = struct{}{}
	}

	// Collect all messages to track which ones to remove and create chains for block inclusion
	pruneMsgs := make(map[cid.Cid]*types.SignedMessage, mp.currentSize)
	keepCount := 0

	var chains []*msgChain
	for actor, mset := range pending {
		// we never prune priority actors
		_, keep := priority[actor]
		if keep {
			keepCount += len(mset)
			continue
		}

		// not a priority actor, track the messages and create chains
		for _, m := range mset {
			pruneMsgs[m.Message.Cid()] = m
		}
		actorChains := mp.createMessageChains(actor, mset, baseFee, ts)
		chains = append(chains, actorChains...)
	}

	// Sort the chains
	sort.Slice(chains, func(i, j int) bool {
		return chains[i].Before(chains[j])
	})

	// Keep messages (remove them from pruneMsgs) from chains while we are under the low water mark
	loWaterMark := mp.cfg.SizeLimitLow
keepLoop:
	for _, chain := range chains {
		for _, m := range chain.msgs {
			if keepCount < loWaterMark {
				delete(pruneMsgs, m.Message.Cid())
				keepCount++
			} else {
				break keepLoop
			}
		}
	}

	// and remove all messages that are still in pruneMsgs after processing the chains
	log.Infof("Pruning %d messages", len(pruneMsgs))
	for _, m := range pruneMsgs {
		mp.remove(m.Message.From, m.Message.Nonce)
	}

	return nil
}