package exec

import (
	"bytes"
	"context"
	"sync"
	"time"

	lru "github.com/hashicorp/golang-lru/v2"
	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"

	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/common/dbg"
	"github.com/erigontech/erigon/common/length"
	"github.com/erigontech/erigon/common/log/v3"
	"github.com/erigontech/erigon/db/dbservices"
	"github.com/erigontech/erigon/db/kv"
	"github.com/erigontech/erigon/db/kv/dbutils"
	"github.com/erigontech/erigon/db/rawdb"
	"github.com/erigontech/erigon/execution/cache"
	"github.com/erigontech/erigon/execution/protocol/rules"
	"github.com/erigontech/erigon/execution/state"
	"github.com/erigontech/erigon/execution/types"
	"github.com/erigontech/erigon/execution/types/accounts"
)

type BlockReadAheader struct {
	// keeps some caches for block themselves
	headers *lru.Cache[common.Hash, *types.Header]
	bodies  *lru.Cache[common.Hash, *types.Body]
	senders *lru.Cache[common.Hash, []byte] // just do raw senders

	// The single permit belongs either to one warmup or to the code suspending
	// warmup across an unwind. Warmups never wait for it: read-ahead is
	// best-effort, and queued work would be stale by the time an unwind ends.
	warmupGate *semaphore.Weighted

	// stateCache is the process-global state cache that SharedDomains.GetLatest
	// consults on the EVM hot path. When set, warmBody routes its prefetches
	// through a cache-populating getter so the same hashmap the EVM probes is
	// pre-warmed. Without it, prefetches only warm OS page cache + RoTx
	// cursors — disconnected from the cache layer the EVM actually reads.
	// Mirrors reth's CachedReads / ExecutionCache "same hashmap" property.
	stateCache *cache.StateCache
}

func NewBlockReadAheader() *BlockReadAheader {
	headers, err := lru.New[common.Hash, *types.Header](4)
	if err != nil {
		panic(err)
	}
	bodies, err := lru.New[common.Hash, *types.Body](4)
	if err != nil {
		panic(err)
	}
	senders, err := lru.New[common.Hash, []byte](4)
	if err != nil {
		panic(err)
	}
	return &BlockReadAheader{
		headers:    headers,
		bodies:     bodies,
		senders:    senders,
		warmupGate: semaphore.NewWeighted(1),
	}
}

// SetStateCache wires the process-global state cache so warmBody's
// prefetches land in the same hashmap that SharedDomains.GetLatest probes
// on the EVM hot path. Without this, prefetches warm OS page cache only —
// the EVM still pays the file accessor stack on its first per-address read.
// Idempotent; safe to call before the first AddHeaderAndBody.
func (bra *BlockReadAheader) SetStateCache(sc *cache.StateCache) {
	bra.stateCache = sc
}

// cachePopulatingGetter wraps a kv.TemporalGetter and fills a StateCache
// ReadView as a side effect. Used by warmBody to make read-ahead prefetches
// populate the same in-process cache layer that SharedDomains.GetLatest
// consults — eliminating the file-accessor stack cost on the EVM's first
// touch of any prefetched address.
//
// Code reads also populate the content-addressed and size-cache layers.
type cachePopulatingGetter struct {
	kv.TemporalGetter
	view     cache.ReadView
	stepSize uint64 // for the read txNum upper bound (last txNum of the read's step)
}

func readAheadGetter(ttx kv.TemporalTx, sc *cache.StateCache) kv.TemporalGetter {
	if sc == nil {
		return ttx
	}
	debug := ttx.Debug()
	stateVersion, err := rawdb.GetStateVersion(ttx)
	if err != nil {
		return ttx
	}
	frontier := cache.FrontierWithStateVersion(debug, stateVersion)
	return &cachePopulatingGetter{TemporalGetter: ttx, view: sc.View(frontier), stepSize: debug.StepSize()}
}

func (cpg *cachePopulatingGetter) GetLatest(name kv.Domain, k []byte) ([]byte, kv.Step, error) {
	v, step, err := cpg.TemporalGetter.GetLatest(name, k)
	if err == nil {
		readTxNum := (uint64(step)+1)*cpg.stepSize - 1
		cpg.view.Fill(name, k, v, readTxNum)
	}
	return v, step, err
}

func (bra *BlockReadAheader) AddHeaderAndBody(ctx context.Context, db kv.RoDB, header *types.Header, body *types.Body) {
	blockHash := header.Hash()
	bra.headers.Add(blockHash, header)
	bra.bodies.Add(blockHash, body)
	if db != nil && ctx != nil {
		bra.startWarmup(func() {
			bra.warmBody(ctx, db, header, body, 8) // use 8 workers for warming
		})
	}
}

func (bra *BlockReadAheader) startWarmup(warm func()) bool {
	if !bra.warmupGate.TryAcquire(1) {
		return false
	}
	go func() {
		defer bra.warmupGate.Release(1)
		warm()
	}()
	return true
}

// SuspendWarmup waits for active state-cache warmup and prevents another
// warmup from starting until the returned function is called. Keep it
// suspended while staged unwind state is being read or published. If ctx is
// cancelled first, no suspension remains pending.
func (bra *BlockReadAheader) SuspendWarmup(ctx context.Context) (func(), error) {
	if err := bra.warmupGate.Acquire(ctx, 1); err != nil {
		return nil, err
	}
	return sync.OnceFunc(func() { bra.warmupGate.Release(1) }), nil
}

// WaitForWarmup waits until neither a warmup nor a suspension owns the permit,
// or until the context is cancelled. Call it before closing the database to
// avoid waitTxsAllDoneOnClose hangs.
func (bra *BlockReadAheader) WaitForWarmup(ctx context.Context) {
	if err := bra.warmupGate.Acquire(ctx, 1); err == nil {
		bra.warmupGate.Release(1)
	}
}

func (bra *BlockReadAheader) AddSenders(senders []byte, blockHash common.Hash) {
	if _, ok := bra.bodies.Get(blockHash); !ok {
		return
	}
	bra.senders.Add(blockHash, bytes.Clone(senders))
}

// warmBody warms state for all transactions in a body using multiple workers.
// It reads: To accounts, To account code, To account storage from access lists,
// and block-level access lists. Each worker creates its own transaction.
// AddHeaderAndBody permits only one warmBody at a time; concurrent requests
// skip warming.
func (bra *BlockReadAheader) warmBody(ctx context.Context, db kv.RoDB, header *types.Header, body *types.Body, workers int) {
	if !dbg.ReadAhead {
		return
	}

	if workers <= 0 {
		workers = 1
	}

	var wg errgroup.Group

	// If BAL exists in DB, use BAL warming (more complete)
	var bal types.BlockAccessList
	if header != nil && db != nil {
		tx, err := db.BeginRo(ctx)
		if err != nil {
			log.Warn("[warmBody] failed to open tx for BAL", "blockNum", header.Number.Uint64(), "blockHash", header.Hash(), "err", err)
		} else {
			data, err := tx.GetOne(kv.BlockAccessList, dbutils.BlockBodyKey(header.Number.Uint64(), header.Hash()))
			if err != nil {
				log.Warn("[warmBody] failed to read BAL", "blockNum", header.Number.Uint64(), "blockHash", header.Hash(), "err", err)
			} else if len(data) > 0 {
				bal, err = types.DecodeBlockAccessListBytes(data)
				if err != nil {
					log.Warn("[warmBody] failed to decode BAL", "blockNum", header.Number.Uint64(), "blockHash", header.Hash(), "err", err)
				}
			}
			tx.Rollback()
		}
	}

	balLen := len(bal)
	if balLen > 0 {
		balWorkers := min(workers, balLen)

		// Pre-divide work: each worker gets a dedicated range of BAL entries
		entriesPerWorker := (balLen + balWorkers - 1) / balWorkers

		for w := range balWorkers {
			start := w * entriesPerWorker
			end := min(start+entriesPerWorker, balLen)
			if start >= balLen {
				break
			}

			// Capture loop variables for closure
			workerStart, workerEnd, workerID := start, end, w
			wg.Go(func() error {
				startTime := time.Now()
				tx, err := db.BeginRo(ctx)
				if err != nil {
					return err
				}
				defer tx.Rollback()

				ttx, ok := tx.(kv.TemporalTx)
				if !ok {
					return nil
				}
				stateReader := state.NewReaderV3(readAheadGetter(ttx, bra.stateCache))

				for idx := workerStart; idx < workerEnd; idx++ {
					select {
					case <-ctx.Done():
						return ctx.Err()
					default:
					}

					acctChanges := bal[idx]
					acct, _ := stateReader.ReadAccountData(acctChanges.Address)
					// Warm code if account has code or if there are code changes.
					if acct != nil && !acct.CodeHash.IsEmpty() {
						stateReader.ReadAccountCode(acctChanges.Address)
					} else if len(acctChanges.CodeChanges) > 0 {
						stateReader.ReadAccountCode(acctChanges.Address)
					}
					for _, slotChanges := range acctChanges.StorageChanges {
						stateReader.ReadAccountStorage(acctChanges.Address, slotChanges.Slot)
					}
					for _, slot := range acctChanges.StorageReads {
						stateReader.ReadAccountStorage(acctChanges.Address, slot)
					}
				}
				log.Debug("[warmBody] BAL worker finished", "worker", workerID, "entries", workerEnd-workerStart, "elapsed", time.Since(startTime))
				return nil
			})
		}
		wg.Wait()
		return
	}
	// Fallback: per-transaction warming when no BAL
	txns := body.Transactions
	if len(txns) == 0 {
		return
	}

	txnLen := len(txns)
	if workers > txnLen {
		workers = txnLen
	}

	// Pre-divide work: each worker gets a dedicated range of transactions
	txnsPerWorker := (txnLen + workers - 1) / workers

	for w := 0; w < workers; w++ {
		start := w * txnsPerWorker
		end := min(start+txnsPerWorker, txnLen)
		if start >= txnLen {
			break
		}

		// Capture loop variables for closure
		workerStart, workerEnd, workerID := start, end, w
		wg.Go(func() error {
			startTime := time.Now()
			tx, err := db.BeginRo(ctx)
			if err != nil {
				return err
			}
			defer tx.Rollback()

			ttx, ok := tx.(kv.TemporalTx)
			if !ok {
				return nil
			}
			stateReader := state.NewReaderV3(readAheadGetter(ttx, bra.stateCache))

			for txIdx := workerStart; txIdx < workerEnd; txIdx++ {
				select {
				case <-ctx.Done():
					return ctx.Err()
				default:
				}

				txn := txns[txIdx]

				// Warm To account and its code if it has one
				if toAddr := txn.GetTo(); toAddr != nil {
					to := accounts.InternAddress(*toAddr)
					if acct, _ := stateReader.ReadAccountData(to); acct != nil && !acct.CodeHash.IsEmpty() {
						stateReader.ReadAccountCode(to)
					}
				}

				// Warm transaction access list accounts and their code
				for _, entry := range txn.GetAccessList() {
					addr := accounts.InternAddress(entry.Address)
					if acct, _ := stateReader.ReadAccountData(addr); acct != nil && !acct.CodeHash.IsEmpty() {
						stateReader.ReadAccountCode(addr)
					}
					for _, slot := range entry.StorageKeys {
						stateReader.ReadAccountStorage(addr, accounts.InternKey(slot))
					}
				}
			}
			log.Debug("[warmBody] TX worker finished", "worker", workerID, "txns", workerEnd-workerStart, "elapsed", time.Since(startTime))
			return nil
		})
	}

	wg.Wait()
}

func (bra *BlockReadAheader) ReadBodyWithTransactions(blockHash common.Hash) (*types.Body, bool) {
	return bra.bodies.Get(blockHash)
}

func (bra *BlockReadAheader) ReadBlockWithSenders(blockHash common.Hash) (*types.Block, bool) {
	header, ok := bra.headers.Get(blockHash)
	if header == nil || !ok {
		return nil, false
	}
	body, ok := bra.bodies.Get(blockHash)
	if body == nil || !ok {
		return nil, false
	}
	senders, ok := bra.senders.Get(blockHash)
	if len(senders) == 0 || !ok {
		return nil, false
	}
	sendersAddresses := make([]common.Address, 0, len(senders)/length.Addr)
	for i := 0; i < len(senders); i += length.Addr {
		sendersAddresses = append(sendersAddresses, common.BytesToAddress(senders[i:i+length.Addr]))
	}
	body.SendersToTxs(sendersAddresses)
	return types.NewBlockFromStorage(header.Hash(), header, body.Transactions, body.Uncles, body.Withdrawals), true
}

func BlocksReadAhead(ctx context.Context, workers int, db kv.RoDB, engine rules.Engine, blockReader dbservices.FullBlockReader) (chan uint64, context.CancelFunc) {
	const readAheadBlocks = 500
	readAhead := make(chan uint64, readAheadBlocks)
	g, gCtx := errgroup.WithContext(ctx)
	for range workers {
		g.Go(func() (err error) {
			var bn uint64
			var ok bool
			var tx kv.Tx
			defer func() {
				if tx != nil {
					tx.Rollback()
				}
			}()

			for i := 0; ; i++ {
				select {
				case bn, ok = <-readAhead:
					if !ok {
						return
					}
				case <-gCtx.Done():
					return gCtx.Err()
				}

				if i%100 == 0 {
					if tx != nil {
						tx.Rollback()
					}
					tx, err = db.BeginRo(ctx)
					if err != nil {
						return err
					}
				}

				if err := blocksReadAheadFunc(gCtx, tx, bn+readAheadBlocks, engine, blockReader); err != nil {
					return err
				}
			}
		})
	}
	return readAhead, func() {
		close(readAhead)
		_ = g.Wait()
	}
}
func blocksReadAheadFunc(ctx context.Context, tx kv.Tx, blockNum uint64, engine rules.Engine, blockReader dbservices.FullBlockReader) error {
	block, err := blockReader.BlockByNumber(ctx, tx, blockNum)
	if err != nil {
		return err
	}
	if block == nil {
		return nil
	}
	_, _ = engine.Author(block.HeaderNoCopy()) // Bor consensus: this calc is heavy and has cache

	ttx, ok := tx.(kv.TemporalTx)
	if !ok {
		return nil
	}

	stateReader := state.NewReaderV3(ttx)
	senders := block.Body().SendersFromTxs()

	for _, sender := range senders {
		a, _ := stateReader.ReadAccountData(accounts.InternAddress(sender))
		if a == nil {
			continue
		}

		//Code domain using .bt index - means no false-positives
		if code, _ := stateReader.ReadAccountCode(accounts.InternAddress(sender)); len(code) > 0 {
			_, _ = code[0], code[len(code)-1]
		}
	}

	for _, txn := range block.Transactions() {
		toaddr := txn.GetTo()
		if toaddr != nil {
			to := accounts.InternAddress(*toaddr)
			a, _ := stateReader.ReadAccountData(to)
			if a == nil {
				continue
			}
			//if account != nil && !bytes.Equal(account.CodeHash, types.EmptyCodeHash.Bytes()) {
			//	reader.Code(*tx.To(), common.BytesToHash(account.CodeHash))
			//}
			if code, _ := stateReader.ReadAccountCode(to); len(code) > 0 {
				_, _ = code[0], code[len(code)-1]
			}

			for _, list := range txn.GetAccessList() {
				stateReader.ReadAccountData(accounts.InternAddress(list.Address))
				if len(list.StorageKeys) > 0 {
					for _, slot := range list.StorageKeys {
						stateReader.ReadAccountStorage(accounts.InternAddress(list.Address), accounts.InternKey(slot))
					}
				}
			}
			//TODO: exec txn and pre-fetch commitment keys. see also: `func (p *statePrefetcher) Prefetch` in geth
		}

	}
	_, _ = stateReader.ReadAccountData(accounts.InternAddress(block.Coinbase()))

	return nil
}
