// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package zkvm

import (
	"container/heap"
	"context"
	"errors"
	"sync"

	"github.com/luxfi/log"

	"github.com/luxfi/ids"
	vmcore "github.com/luxfi/vm"
)

// Mempool manages pending transactions
type Mempool struct {
	log     log.Logger
	maxSize int

	// Transaction storage
	txs    map[ids.ID]*MempoolTx
	txHeap TxHeap

	// Nullifier tracking to prevent conflicts
	nullifiers map[string]ids.ID // nullifier -> txID

	// work tells consensus a block can be built. The mempool is what learns
	// that first, whichever path the transaction arrived by.
	work vmcore.Latch

	mu sync.RWMutex
}

// MempoolTx represents a transaction in the mempool
type MempoolTx struct {
	tx  *Transaction
	fee uint64
}

// NewMempool creates a new mempool
func NewMempool(maxSize int, log log.Logger) *Mempool {
	return &Mempool{
		log:        log,
		maxSize:    maxSize,
		txs:        make(map[ids.ID]*MempoolTx),
		txHeap:     make(TxHeap, 0),
		nullifiers: make(map[string]ids.ID),
	}
}

// AddTransaction adds a transaction to the mempool.
//
// The id is derived here, not read from the caller. An HTTP client decoding
// straight into a Transaction supplies whatever id it likes, and a proposer
// that carried that id would compute a block id its own peers do not — the
// block it built would not be the block they see.
func (mp *Mempool) AddTransaction(tx *Transaction) error {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	tx.ID = tx.ComputeID()

	// Check if transaction already exists
	if _, exists := mp.txs[tx.ID]; exists {
		return nil // Already in mempool
	}

	// Check for nullifier conflicts
	for _, nullifier := range tx.Nullifiers {
		if existingTxID, exists := mp.nullifiers[string(nullifier)]; exists {
			mp.log.Debug("Nullifier conflict in mempool",
				log.String("newTx", tx.ID.String()),
				log.String("existingTx", existingTxID.String()),
			)
			return errors.New("nullifier already in mempool")
		}
	}

	// A full pool gives up whatever pays least, so what it holds is always the
	// best it has been offered. An arrival that is itself the cheapest is
	// refused: taking it would mean dropping a better transaction for a worse
	// one. Ordering is by fee — it used to divide every fee by the same fixed
	// 256, which orders identically and reads as a rate it is not.
	if len(mp.txs) >= mp.maxSize {
		cheapest := mp.cheapest()
		if cheapest == nil || cheapest.fee >= tx.Fee {
			return errors.New("mempool is full and the transaction pays less than what it would displace")
		}
		mp.removeTxNoLock(cheapest.tx.ID)
	}

	mempoolTx := &MempoolTx{tx: tx, fee: tx.Fee}

	// Add to storage
	mp.txs[tx.ID] = mempoolTx
	heap.Push(&mp.txHeap, mempoolTx)

	// Track nullifiers
	for _, nullifier := range tx.Nullifiers {
		mp.nullifiers[string(nullifier)] = tx.ID
	}

	mp.log.Debug("Added transaction to mempool",
		log.String("txID", tx.ID.String()),
		log.Uint64("fee", tx.Fee),
		log.Int("mempoolSize", len(mp.txs)),
	)

	// Consensus builds nothing until it is told there is something to build.
	mp.work.Signal()

	return nil
}

// WaitForEvent blocks until there is a transaction to build a block from, or
// the caller gives up.
func (mp *Mempool) WaitForEvent(ctx context.Context) (vmcore.Message, error) {
	return mp.work.WaitForEvent(ctx)
}

// RemoveTransaction removes a transaction from the mempool
func (mp *Mempool) RemoveTransaction(txID ids.ID) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	mp.removeTxNoLock(txID)
}

// cheapest returns the transaction paying least, which is the one a full pool
// gives up. The heap is ordered highest-first, so the cheapest sits among its
// leaves rather than at its root; the pool is bounded and this runs only when
// it is full, so walking it is the whole of the work.
func (mp *Mempool) cheapest() *MempoolTx {
	var out *MempoolTx
	for _, candidate := range mp.txHeap {
		if out == nil || candidate.fee < out.fee {
			out = candidate
		}
	}
	return out
}

// removeTxNoLock removes a transaction without locking (internal use)
func (mp *Mempool) removeTxNoLock(txID ids.ID) {
	mempoolTx, exists := mp.txs[txID]
	if !exists {
		return
	}

	// Remove from storage
	delete(mp.txs, txID)

	// Remove nullifiers
	for _, nullifier := range mempoolTx.tx.Nullifiers {
		delete(mp.nullifiers, string(nullifier))
	}

	// Remove from heap (expensive, but necessary)
	for i, tx := range mp.txHeap {
		if tx.tx.ID == txID {
			heap.Remove(&mp.txHeap, i)
			break
		}
	}
}

// GetPendingTransactions returns pending transactions sorted by priority
func (mp *Mempool) GetPendingTransactions(limit int) []*Transaction {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	// Create a copy of the heap to sort
	tempHeap := make(TxHeap, len(mp.txHeap))
	copy(tempHeap, mp.txHeap)
	heap.Init(&tempHeap)

	// Extract top transactions
	txs := make([]*Transaction, 0, limit)
	for i := 0; i < limit && tempHeap.Len() > 0; i++ {
		mempoolTx := heap.Pop(&tempHeap).(*MempoolTx)
		txs = append(txs, mempoolTx.tx)
	}

	return txs
}

// HasTransaction checks if a transaction is in the mempool
func (mp *Mempool) HasTransaction(txID ids.ID) bool {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	_, exists := mp.txs[txID]
	return exists
}

// Size returns the number of transactions in the mempool
func (mp *Mempool) Size() int {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	return len(mp.txs)
}

// PruneExpired drops transactions the chain has passed. Nothing else does:
// a transaction that can never enter a block occupies a slot forever, and a
// pool full of those refuses every honest arrival that pays the same floor.
func (mp *Mempool) PruneExpired(currentHeight uint64) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	var toRemove []ids.ID

	for txID, mempoolTx := range mp.txs {
		if mempoolTx.tx.Expiry > 0 && mempoolTx.tx.Expiry < currentHeight {
			toRemove = append(toRemove, txID)
		}
	}

	for _, txID := range toRemove {
		mp.removeTxNoLock(txID)
	}

	if len(toRemove) > 0 {
		mp.log.Info("Pruned expired transactions",
			log.Int("count", len(toRemove)),
			log.Uint64("currentHeight", currentHeight),
		)
	}
}

// TxHeap implements heap.Interface for priority ordering
type TxHeap []*MempoolTx

func (h TxHeap) Len() int { return len(h) }

func (h TxHeap) Less(i, j int) bool {
	// Higher fee = higher priority
	return h[i].fee > h[j].fee
}

func (h TxHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *TxHeap) Push(x interface{}) {
	*h = append(*h, x.(*MempoolTx))
}

func (h *TxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	tx := old[n-1]
	*h = old[0 : n-1]
	return tx
}