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

import (
	"context"
	"errors"
	"fmt"
	"math"
	"runtime"
	"strings"
	"time"

	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/common/dbg"
	"github.com/erigontech/erigon/common/log/v3"
	"github.com/erigontech/erigon/db/datadir"
	"github.com/erigontech/erigon/db/dbservices"
	"github.com/erigontech/erigon/db/integrity"
	"github.com/erigontech/erigon/db/kv"
	"github.com/erigontech/erigon/db/kv/backup"
	"github.com/erigontech/erigon/db/kv/kvcfg"
	"github.com/erigontech/erigon/db/rawdb/rawtemporaldb"
	"github.com/erigontech/erigon/db/snapshotsync/blocksnapshots"
	dbstate "github.com/erigontech/erigon/db/state"
	"github.com/erigontech/erigon/db/state/execctx"
	"github.com/erigontech/erigon/execution/chain"
	"github.com/erigontech/erigon/execution/exec"
	"github.com/erigontech/erigon/execution/protocol/rules"
	"github.com/erigontech/erigon/execution/stagedsync/stages"
	"github.com/erigontech/erigon/execution/types"
	"github.com/erigontech/erigon/node/ethconfig"
)

type CustomTraceCfg struct {
	tmpdir   string
	db       kv.TemporalRwDB
	ExecArgs *exec.ExecArgs

	Produce Produce
}
type Produce struct {
	ReceiptDomain bool
	RCacheDomain  bool
	LogAddr       bool
	LogTopic      bool
	TraceFrom     bool
	TraceTo       bool
}

func NewProduce(produceList []string) Produce {
	var produce Produce
	for _, p := range produceList {
		p = strings.TrimSpace(p)
		switch p {
		case kv.ReceiptDomain.String():
			produce.ReceiptDomain = true
		case kv.RCacheDomain.String():
			produce.RCacheDomain = true
		case kv.LogAddrIdx.String():
			produce.LogAddr = true
		case kv.LogTopicIdx.String():
			produce.LogTopic = true
		case kv.TracesFromIdx.String():
			produce.TraceFrom = true
		case kv.TracesToIdx.String():
			produce.TraceTo = true
		default:
			panic(fmt.Errorf("assert: unknown Produce %#v", p))
		}
	}
	return produce
}

func StageCustomTraceCfg(produce []string, db kv.TemporalRwDB, dirs datadir.Dirs, br dbservices.FullBlockReader,
	cc *chain.Config, engine rules.Engine,
	genesis *types.Genesis, syncCfg ethconfig.Sync) CustomTraceCfg {
	execArgs := &exec.ExecArgs{
		ChainDB:     db,
		BlockReader: br,
		ChainConfig: cc,
		Dirs:        dirs,
		Engine:      engine,
		Genesis:     genesis,
		Workers:     syncCfg.ExecWorkerCount,
	}
	return CustomTraceCfg{
		db:       db,
		ExecArgs: execArgs,
		Produce:  NewProduce(produce),
	}
}

func unalignProduced(agg *dbstate.Aggregator, produce Produce) (realign func()) {
	var undo []func()
	if produce.ReceiptDomain {
		undo = append(undo, agg.Unalign(kv.ReceiptDomain))
	}
	if produce.RCacheDomain {
		undo = append(undo, agg.Unalign(kv.RCacheDomain))
	}
	if produce.LogAddr {
		undo = append(undo, agg.UnalignIdx(kv.LogAddrIdx))
	}
	if produce.LogTopic {
		undo = append(undo, agg.UnalignIdx(kv.LogTopicIdx))
	}
	if produce.TraceFrom {
		undo = append(undo, agg.UnalignIdx(kv.TracesFromIdx))
	}
	if produce.TraceTo {
		undo = append(undo, agg.UnalignIdx(kv.TracesToIdx))
	}
	return func() {
		for _, f := range undo {
			f()
		}
	}
}

func SpawnCustomTrace(cfg CustomTraceCfg, ctx context.Context, logger log.Logger) error {
	if cfg.Produce.RCacheDomain {
		if err := cfg.db.View(context.Background(), func(tx kv.Tx) error {
			return kvcfg.PersistReceipts.MustBeEnabled(tx, "you must enable `--prune.include-receipts` flag in db. remove chaindata and start erigon with this flag")
		}); err != nil {
			panic(err)
		}
	}

	log.Info("[stage_custom_trace] start params", "produce", cfg.Produce)
	txNumsReader := cfg.ExecArgs.BlockReader.TxnumReader()

	// what this re-derives lags the state until it finishes, and must not clamp the files it
	// re-executes from
	defer unalignProduced(cfg.db.(dbstate.HasAgg).Agg().(*dbstate.Aggregator), cfg.Produce)()

	//agg := cfg.db.(dbstate.HasAgg).Agg().(*dbstate.Aggregator)
	//stepSize := agg.StepSize()

	// 1. Require stage_exec > 0: means don't need handle "half-block execution case here"
	// 2. Require stage_exec > 0: means has enough state-history
	var execProgress uint64
	var startBlock, endBlock uint64
	if err := cfg.db.ViewTemporal(ctx, func(tx kv.TemporalTx) (err error) {
		execProgress, err = stages.GetStageProgress(tx, stages.Execution)
		if err != nil {
			return err
		}
		if execProgress == 0 {
			return errors.New("stage_exec progress is 0. please run `integration stage_exec --batchSize=1m` for couple minutes")
		}

		fromTxNum := progressOfDomains(tx, cfg.Produce)
		var ok bool
		startBlock, ok, err = txNumsReader.FindBlockNum(ctx, tx, fromTxNum)
		if err != nil {
			return fmt.Errorf("getting last executed block: %w", err)
		}
		if !ok {
			panic(ok)
		}
		return nil
	}); err != nil {
		return err
	}
	endBlock = execProgress

	defer cfg.ExecArgs.BlockReader.Snapshots().(*blocksnapshots.RoSnapshots).MadvNormal().DisableReadAhead()
	//defer tx.(dbstate.HasAggTx).AggTx().(*dbstate.AggregatorRoTx).MadvNormal().DisableReadAhead()

	log.Info("SpawnCustomTrace", "startBlock", startBlock, "endBlock", endBlock)
	batchSize := uint64(50_000)
	for startBlock < endBlock {
		to := min(endBlock+1, startBlock+batchSize)
		if err := customTraceBatchProduce(ctx, cfg.Produce, cfg.ExecArgs, cfg.db, startBlock, to, "custom_trace", logger); err != nil {
			return err
		}
		startBlock = to
	}

	logEvery := time.NewTicker(20 * time.Second)
	defer logEvery.Stop()
	chkEvery := time.NewTicker(3 * time.Second)
	defer chkEvery.Stop()

Loop:
	for {
		select {
		case <-ctx.Done():
			logger.Warn("[snapshots] user has interrupted process, waiting for build & merge files")
		case <-cfg.db.(dbstate.HasAgg).Agg().(*dbstate.Aggregator).WaitForBuildAndMerge(context.Background()):
			break Loop // Here we don't quit due to ctx because it's not safe for files.
		case <-logEvery.C:
			var m runtime.MemStats
			dbg.ReadMemStats(&m)
			//TODO: log progress and list of domains/files
			logger.Info("[snapshots] Building files", "alloc", common.ByteCount(m.Alloc), "sys", common.ByteCount(m.Sys))
		}
	}

	tx, err := cfg.db.BeginTemporalRo(context.Background())
	if err != nil {
		return err
	}
	defer tx.Rollback()
	log.Info("SpawnCustomTrace finish",
		"accounts", tx.Debug().DomainProgress(kv.AccountsDomain),
		"receipts", tx.Debug().DomainProgress(kv.ReceiptDomain),
		"rcache", tx.Debug().DomainProgress(kv.RCacheDomain))

	if cfg.Produce.ReceiptDomain {
		if err := integrity.ValidateDomainProgress(ctx, cfg.db, kv.ReceiptDomain, txNumsReader); err != nil {
			return err
		}
	}
	if cfg.Produce.RCacheDomain {
		if err := integrity.ValidateDomainProgress(ctx, cfg.db, kv.RCacheDomain, txNumsReader); err != nil {
			return err
		}
	}

	return nil
}

// stage_custom_trace input range is block based, and not txNum based.
// it doesn't need to account for "half-block execution" case, because it
// must have some stage_exec progress, which means it resumes from full blocks.
// also, it appends/puts to db blockResults and not "txResult".
func customTraceBatchProduce(ctx context.Context, produce Produce, cfg *exec.ExecArgs, db kv.TemporalRwDB, fromBlock, toBlock uint64, logPrefix string, logger log.Logger) error {
	if err := db.UpdateTemporal(ctx, func(tx kv.TemporalRwTx) error {
		if _, err := tx.PruneSmallBatches(ctx, 10*time.Hour); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}
	var lastTxNum uint64
	{
		tx, err := db.BeginTemporalRw(ctx)
		if err != nil {
			return err
		}
		defer tx.Rollback()

		doms, err := execctx.NewSharedDomains(ctx, tx, logger)
		if err != nil {
			return err
		}
		defer doms.Close()

		if err := customTraceBatch(ctx, produce, cfg, tx, doms, fromBlock, toBlock, logPrefix, logger); err != nil {
			return err
		}

		// Commit runs the validate hook after flush, before the commit, so the
		// asserts + SeekCommitment see the flushed state on tx exactly as the
		// prior Flush-then-commit shape did.
		if err := doms.Commit(ctx, tx, func(tx kv.RwTx) error {
			if produce.ReceiptDomain {
				return AssertReceipts(ctx, cfg, db, fromBlock, toBlock)
			}
			return nil
		}); err != nil {
			return err
		}

	}

	agg := db.(dbstate.HasAgg).Agg().(*dbstate.Aggregator)
	var fromStep, toStep kv.Step
	if err := db.ViewTemporal(ctx, func(tx kv.TemporalTx) error {
		fromStep = firstStepNotInFiles(tx, produce)
		// toStep must reflect what this batch actually re-derived, not the commitment
		// frontier: when rebuilding a subset of domains, commitment can be ahead, which
		// would build empty files past the domain's data and desync pruning.
		var err error
		lastTxNum, err = cfg.BlockReader.TxnumReader().Max(ctx, tx, toBlock-1)
		if err != nil {
			return err
		}
		if lastTxNum/agg.StepSize() > 0 {
			toStep = kv.Step(lastTxNum / agg.StepSize())
		}
		return nil
	}); err != nil {
		return err
	}
	if err := agg.BuildFiles2(ctx, fromStep, toStep, true); err != nil {
		return err
	}
	if err := db.Update(ctx, func(tx kv.RwTx) error {
		if _, err := tx.(kv.TemporalRwTx).PruneSmallBatches(ctx, 10*time.Hour); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}

	return nil
}

func AssertReceipts(ctx context.Context, cfg *exec.ExecArgs, db kv.TemporalRoDB, fromBlock, toBlock uint64) (err error) {
	if !dbg.AssertEnabled {
		return
	}
	if cfg.ChainConfig.Bor != nil { //TODO: enable me
		return nil
	}
	return integrity.ReceiptsNoDupsRange(ctx, fromBlock, toBlock, db, cfg.BlockReader, true)
}

func customTraceBatch(ctx context.Context, produce Produce, cfg *exec.ExecArgs, tx kv.TemporalRwTx, doms *execctx.SharedDomains, fromBlock, toBlock uint64, logPrefix string, logger log.Logger) error {
	const logPeriod = 5 * time.Second
	logEvery := time.NewTicker(logPeriod)
	defer logEvery.Stop()

	var cumulativeBlobGasUsedInBlock uint64
	// The reduce side runs on a single goroutine, so one writer serves the whole batch.
	var receipts rawtemporaldb.ReceiptWriter

	txNumsReader := cfg.BlockReader.TxnumReader()
	fromTxNum, _ := txNumsReader.Min(ctx, tx, fromBlock)
	prevTxNumLog := fromTxNum

	var m runtime.MemStats
	if err := exec.CustomTraceMapReduce(ctx, fromBlock, toBlock, exec.TraceConsumerFunc(
		func(blockResult *exec.BlockResult, result *exec.TxResult, tx kv.TemporalTx) error {
			if result.Err != nil {
				return result.Err
			}

			txTask := result.Task.(*exec.TxTask)
			if txTask.Tx() != nil {
				cumulativeBlobGasUsedInBlock += txTask.Tx().GetBlobGas()
			}

			putter := doms.AsPutDel(tx)

			if produce.ReceiptDomain {
				var logIndexAfterTx uint32
				var cumGasUsed uint64

				if txTask.IsBlockEnd() { // block changed
					if cfg.ChainConfig.Bor != nil && txTask.TxIndex >= 1 {
						// get last receipt and store the last log index + 1
						lastReceipt := blockResult.Receipts[txTask.TxIndex-1]
						if lastReceipt == nil {
							return fmt.Errorf("receipt is nil but should be populated, txIndex=%d, block=%d", txTask.TxIndex-1, txTask.BlockNumber())
						}
						if len(lastReceipt.Logs) > 0 {
							firstIndex := lastReceipt.Logs[len(lastReceipt.Logs)-1].Index + 1
							logIndexAfterTx = uint32(firstIndex) + uint32(len(result.Logs))
							cumGasUsed = lastReceipt.CumulativeGasUsed
						}
					}
				} else {
					if txTask.TxIndex >= 0 {
						receipt := blockResult.Receipts[txTask.TxIndex]
						if receipt != nil {
							logIndexAfterTx = receipt.FirstLogIndexWithinBlock + uint32(len(result.Logs))
							cumGasUsed = receipt.CumulativeGasUsed
						}
					}
				}

				if err := receipts.AppendMetadata(putter, logIndexAfterTx, cumGasUsed, cumulativeBlobGasUsedInBlock, txTask.TxNum); err != nil {
					return err
				}

				if txTask.IsBlockEnd() { // block changed
					cumulativeBlobGasUsedInBlock = 0
				}
			}

			if produce.RCacheDomain {
				var receipt *types.Receipt
				if !txTask.IsBlockEnd() {
					receipt = result.Receipt
				} else {
					if cfg.ChainConfig.Bor != nil && txTask.TxIndex >= 1 {
						// issue: https://github.com/erigontech/erigon/issues/16037
						receipt = blockResult.Receipts[txTask.TxIndex-1]
						if receipt == nil {
							return fmt.Errorf("receipt is nil but should be populated, txIndex=%d, block=%d", txTask.TxIndex-1, txTask.BlockNumber())
						}
					}
				}
				if err := receipts.Append(putter, receipt, txTask.TxNum); err != nil {
					return err
				}
			}

			if produce.LogAddr {
				for _, lg := range result.Logs {
					if err := doms.IndexAdd(kv.LogAddrIdx, lg.Address[:], txTask.TxNum); err != nil {
						return err
					}
				}
			}
			if produce.LogTopic {
				for _, lg := range result.Logs {
					for i := range lg.Topics {
						if err := doms.IndexAdd(kv.LogTopicIdx, lg.Topics[i][:], txTask.TxNum); err != nil {
							return err
						}
					}
				}
			}
			if produce.TraceFrom {
				for addr := range result.TraceFroms {
					addrValue := addr.Value()
					if err := doms.IndexAdd(kv.TracesFromIdx, addrValue[:], txTask.TxNum); err != nil {
						return err
					}
				}
			}
			if produce.TraceTo {
				for addr := range result.TraceTos {
					addrValue := addr.Value()
					if err := doms.IndexAdd(kv.TracesToIdx, addrValue[:], txTask.TxNum); err != nil {
						return err
					}
				}
			}

			select {
			case <-logEvery.C:
				if prevTxNumLog > 0 {
					dbg.ReadMemStats(&m)
					txsPerSec := (txTask.TxNum - prevTxNumLog) / uint64(logPeriod.Seconds())
					log.Info(fmt.Sprintf("[%s] Scanned", logPrefix), "block", fmt.Sprintf("%.3fm", float64(txTask.BlockNumber())/1_000_000), "tx/s", fmt.Sprintf("%.1fK", float64(txsPerSec)/1_000.0), "alloc", common.ByteCount(m.Alloc), "sys", common.ByteCount(m.Sys))
				}
				prevTxNumLog = txTask.TxNum
			default:
			}
			return nil
		}), tx, cfg, logger); err != nil {
		return err
	}

	return nil
}

func progressOfDomains(tx kv.TemporalTx, produce Produce) uint64 {
	//TODO: need better way to detect start point. What if domain/index is sparse (has rare events).
	dbg := tx.Debug()
	txNum := uint64(math.MaxUint64)
	if produce.ReceiptDomain {
		txNum = min(txNum, dbg.DomainProgress(kv.ReceiptDomain))
	}
	if produce.RCacheDomain {
		txNum = min(txNum, dbg.DomainProgress(kv.RCacheDomain))
	}
	if produce.LogAddr {
		txNum = min(txNum, dbg.IIProgress(kv.LogAddrIdx))
	}
	if produce.LogTopic {
		txNum = min(txNum, dbg.IIProgress(kv.LogTopicIdx))
	}
	if produce.TraceFrom {
		txNum = min(txNum, dbg.IIProgress(kv.TracesFromIdx))
	}
	if produce.TraceTo {
		txNum = min(txNum, dbg.IIProgress(kv.TracesToIdx))
	}
	return txNum
}

func firstStepNotInFiles(tx kv.TemporalTx, produce Produce) kv.Step {
	//TODO: need better way to detect start point. What if domain/index is sparse (has rare events).
	ac := dbstate.AggTx(tx)
	fromStep := kv.Step(math.MaxUint64)
	if produce.ReceiptDomain {
		fromStep = min(fromStep, kv.Step(ac.DbgDomain(kv.ReceiptDomain).FirstStepNotInFiles()))
	}
	if produce.RCacheDomain {
		fromStep = min(fromStep, kv.Step(ac.DbgDomain(kv.RCacheDomain).FirstStepNotInFiles()))
	}
	if produce.LogAddr {
		fromStep = min(fromStep, ac.DbgII(kv.LogAddrIdx).FirstStepNotInFiles())
	}
	if produce.LogTopic {
		fromStep = min(fromStep, ac.DbgII(kv.LogTopicIdx).FirstStepNotInFiles())
	}
	if produce.TraceFrom {
		fromStep = min(fromStep, ac.DbgII(kv.TracesFromIdx).FirstStepNotInFiles())
	}
	if produce.TraceTo {
		fromStep = min(fromStep, ac.DbgII(kv.TracesToIdx).FirstStepNotInFiles())
	}
	return fromStep
}

func StageCustomTraceReset(ctx context.Context, db kv.TemporalRwDB, produce Produce) error {
	tx, err := db.BeginTemporalRw(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	var tables []string
	if produce.ReceiptDomain {
		tables = append(tables, db.Debug().DomainTables(kv.ReceiptDomain)...)
	}
	if produce.RCacheDomain {
		tables = append(tables, db.Debug().DomainTables(kv.RCacheDomain)...)
	}
	if produce.LogAddr {
		tables = append(tables, db.Debug().InvertedIdxTables(kv.LogAddrIdx)...)
	}
	if produce.LogTopic {
		tables = append(tables, db.Debug().InvertedIdxTables(kv.LogTopicIdx)...)
	}
	if produce.TraceFrom {
		tables = append(tables, db.Debug().InvertedIdxTables(kv.TracesFromIdx)...)
	}
	if produce.TraceTo {
		tables = append(tables, db.Debug().InvertedIdxTables(kv.TracesToIdx)...)
	}
	if err := backup.ClearTables(ctx, db, tx, tables...); err != nil {
		return err
	}
	// Clearing the data tables alone is not enough: the prune-progress bookmark in
	// TblPruningValsProg survives and would make the post-rebuild prune short-circuit
	// (prs.TxTo >= txTo && Done), so the re-written history never gets pruned.
	for _, table := range tables {
		if err := dbstate.InvalidatePruneProgress(tx, table); err != nil {
			return err
		}
	}
	return tx.Commit()
}
