// Copyright 2023 The go-ethereum Authors
// This file is part of go-ethereum.
//
// go-ethereum is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// go-ethereum 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 General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with go-ethereum. If not, see <http://www.gnu.org/licenses/>.

package utils

import (
	"bytes"
	"crypto/sha256"
	"io"
	"math/big"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/consensus"
	"github.com/ethereum/go-ethereum/consensus/beacon"
	"github.com/ethereum/go-ethereum/consensus/ethash"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/rawdb"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/internal/era"
	"github.com/ethereum/go-ethereum/internal/era/execdb"
	"github.com/ethereum/go-ethereum/internal/era/onedb"
	"github.com/ethereum/go-ethereum/params"
	"github.com/ethereum/go-ethereum/trie"
	"github.com/ethereum/go-ethereum/triedb"
)

var (
	count uint64 = 128
	step  uint64 = 16
)

func TestHistoryImportAndExport(t *testing.T) {
	for _, tt := range []struct {
		name     string
		builder  func(io.Writer) era.Builder
		filename func(network string, epoch int, root common.Hash) string
		from     func(f era.ReadAtSeekCloser) (era.Era, error)
	}{
		{"era1", onedb.NewBuilder, onedb.Filename, onedb.From},
		{"erae", execdb.NewBuilder, execdb.Filename, execdb.From},
	} {
		t.Run(tt.name, func(t *testing.T) {
			var (
				key, _  = crypto.HexToECDSA("b71c71a67e1177ad4e901695e1b4b9ee17ae16c6668d313eac2f96dbcda3f291")
				address = crypto.PubkeyToAddress(key.PublicKey)
				genesis = &core.Genesis{
					Config: params.TestChainConfig,
					Alloc:  types.GenesisAlloc{address: {Balance: big.NewInt(1000000000000000000)}},
				}
				signer = types.LatestSigner(genesis.Config)
			)

			// Generate chain.
			db, blocks, _ := core.GenerateChainWithGenesis(genesis, ethash.NewFaker(), int(count), func(i int, g *core.BlockGen) {
				if i == 0 {
					return
				}
				tx, err := types.SignNewTx(key, signer, &types.DynamicFeeTx{
					ChainID:    genesis.Config.ChainID,
					Nonce:      uint64(i - 1),
					GasTipCap:  common.Big0,
					GasFeeCap:  g.PrevBlock(0).BaseFee(),
					Gas:        50000,
					To:         &common.Address{0xaa},
					Value:      big.NewInt(int64(i)),
					Data:       nil,
					AccessList: nil,
				})
				if err != nil {
					t.Fatalf("error creating tx: %v", err)
				}
				g.AddTx(tx)
			})

			// Initialize BlockChain.
			chain, err := core.NewBlockChain(db, genesis, ethash.NewFaker(), nil)
			if err != nil {
				t.Fatalf("unable to initialize chain: %v", err)
			}
			if _, err := chain.InsertChain(blocks); err != nil {
				t.Fatalf("error inserting chain: %v", err)
			}

			// Make temp directory for era files.
			dir := t.TempDir()

			// Export history to temp directory.
			if err := ExportHistory(chain, dir, 0, count, tt.builder, tt.filename); err != nil {
				t.Fatalf("error exporting history: %v", err)
			}

			// Read checksums.
			b, err := os.ReadFile(filepath.Join(dir, "checksums.txt"))
			if err != nil {
				t.Fatalf("failed to read checksums: %v", err)
			}
			checksums := strings.Split(string(b), "\n")

			// Verify each Era.
			entries, _ := era.ReadDir(dir, "mainnet")
			for i, filename := range entries {
				func() {
					f, err := os.Open(filepath.Join(dir, filename))
					if err != nil {
						t.Fatalf("error opening era file: %v", err)
					}
					var (
						h   = sha256.New()
						buf = bytes.NewBuffer(nil)
					)
					if _, err := io.Copy(h, f); err != nil {
						t.Fatalf("unable to recalculate checksum: %v", err)
					}
					if got, want := common.BytesToHash(h.Sum(buf.Bytes()[:])).Hex(), checksums[i]; got != want {
						t.Fatalf("checksum %d does not match: got %s, want %s", i, got, want)
					}
					e, err := tt.from(f)
					if err != nil {
						t.Fatalf("error opening era: %v", err)
					}
					defer e.Close()
					it, err := e.Iterator()
					if err != nil {
						t.Fatalf("error making era reader: %v", err)
					}
					for j := 0; it.Next(); j++ {
						n := i*int(step) + j
						if it.Error() != nil {
							t.Fatalf("error reading block entry %d: %v", n, it.Error())
						}
						block, receipts, err := it.BlockAndReceipts()
						if err != nil {
							t.Fatalf("error reading block entry %d: %v", n, err)
						}
						want := chain.GetBlockByNumber(uint64(n))
						if want, got := uint64(n), block.NumberU64(); want != got {
							t.Fatalf("blocks out of order: want %d, got %d", want, got)
						}
						if want.Hash() != block.Hash() {
							t.Fatalf("block hash mismatch %d: want %s, got %s", n, want.Hash().Hex(), block.Hash().Hex())
						}
						if got := types.DeriveSha(block.Transactions(), trie.NewStackTrie(nil)); got != want.TxHash() {
							t.Fatalf("tx hash %d mismatch: want %s, got %s", n, want.TxHash(), got)
						}
						if got := types.CalcUncleHash(block.Uncles()); got != want.UncleHash() {
							t.Fatalf("uncle hash %d mismatch: want %s, got %s", n, want.UncleHash(), got)
						}
						if got := types.DeriveSha(receipts, trie.NewStackTrie(nil)); got != want.ReceiptHash() {
							t.Fatalf("receipt root %d mismatch: want %s, got %s", n, want.ReceiptHash(), got)
						}
					}
				}()
			}

			// Now import Era.
			db2, err := rawdb.Open(rawdb.NewMemoryDatabase(), rawdb.OpenOptions{})
			if err != nil {
				panic(err)
			}
			t.Cleanup(func() {
				db2.Close()
			})

			genesis.MustCommit(db2, triedb.NewDatabase(db2, triedb.HashDefaults))
			imported, err := core.NewBlockChain(db2, genesis, ethash.NewFaker(), nil)
			if err != nil {
				t.Fatalf("unable to initialize chain: %v", err)
			}
			if err := ImportHistory(imported, dir, "mainnet", tt.from); err != nil {
				t.Fatalf("failed to import chain: %v", err)
			}
			if have, want := imported.CurrentHeader(), chain.CurrentHeader(); have.Hash() != want.Hash() {
				t.Fatalf("imported chain does not match expected, have (%d, %s) want (%d, %s)", have.Number, have.Hash(), want.Number, want.Hash())
			}
		})
	}
}

func TestImportHistoryRejectsForgedTxRootArchive(t *testing.T) {
	genesis, honest, honestReceipts, _ := makeHistorySingleTxFixture(t)
	header := honest.Header()
	header.TxHash = common.HexToHash("0x1234")
	if header.TxHash == honest.TxHash() {
		t.Fatal("fixture is not malformed: tx root unexpectedly matches original block")
	}
	poisoned := honest.WithSeal(header)

	dir := t.TempDir()
	if err := writeEra1(dir, poisoned, honestReceipts); err != nil {
		t.Fatalf("failed to build forged-tx-root era1 archive: %v", err)
	}
	imported := newHistoryImportChain(t, genesis, ethash.NewFaker())

	if err := ImportHistory(imported, dir, "mainnet", onedb.From); err == nil {
		t.Fatal("ImportHistory unexpectedly accepted forged tx root")
	} else if !strings.Contains(err.Error(), "transaction root hash mismatch") {
		t.Fatalf("unexpected ImportHistory error: %v", err)
	}
}

func TestImportHistoryRejectsForgedWithdrawalsArchive(t *testing.T) {
	genesis, honest, honestReceipts := makeMergedHistoryFixture(t)
	header := honest.Header()
	if header.WithdrawalsHash == nil {
		t.Fatal("fixture does not have a withdrawals root")
	}
	forgedWithdrawalsHash := common.HexToHash("0x5678")
	if forgedWithdrawalsHash == *header.WithdrawalsHash {
		t.Fatal("fixture is not malformed: withdrawals root unexpectedly matches original block")
	}
	header.WithdrawalsHash = &forgedWithdrawalsHash
	poisoned := honest.WithSeal(header)

	dir := t.TempDir()
	if err := writeEra1(dir, poisoned, honestReceipts); err != nil {
		t.Fatalf("failed to build forged-withdrawals era1 archive: %v", err)
	}
	imported := newHistoryImportChain(t, genesis, beacon.New(ethash.NewFaker()))

	if err := ImportHistory(imported, dir, "mainnet", onedb.From); err == nil {
		t.Fatal("ImportHistory unexpectedly accepted forged withdrawals root")
	} else if !strings.Contains(err.Error(), "withdrawals root hash mismatch") {
		t.Fatalf("unexpected ImportHistory error: %v", err)
	}
}

func TestImportHistoryRejectsForgedReceiptArchive(t *testing.T) {
	genesis, honest, honestReceipts, _ := makeHistorySingleTxFixture(t)
	forgedReceipt := new(types.Receipt)
	*forgedReceipt = *honestReceipts[0]
	forgedReceipt.Status = types.ReceiptStatusFailed
	if got := types.DeriveSha(types.Receipts{forgedReceipt}, trie.NewStackTrie(nil)); got == honest.ReceiptHash() {
		t.Fatalf("fixture is not malformed: receipt root unexpectedly matches header %s", got)
	}

	dir := t.TempDir()
	if err := writeEra1(dir, honest, types.Receipts{forgedReceipt}); err != nil {
		t.Fatalf("failed to build forged-receipt era1 archive: %v", err)
	}
	imported := newHistoryImportChain(t, genesis, ethash.NewFaker())

	if err := ImportHistory(imported, dir, "mainnet", onedb.From); err == nil {
		t.Fatal("ImportHistory unexpectedly accepted forged receipt root")
	} else if !strings.Contains(err.Error(), "receipt root hash mismatch") {
		t.Fatalf("unexpected ImportHistory error: %v", err)
	}
}

func writeEra1(dir string, block *types.Block, receipts types.Receipts) error {
	buf := new(bytes.Buffer)
	builder := onedb.NewBuilder(buf)

	td := new(big.Int)
	if diff := block.Difficulty(); diff != nil {
		td.Set(diff)
	}
	root, err := func() (common.Hash, error) {
		if err := builder.Add(block, receipts, td); err != nil {
			return common.Hash{}, err
		}
		return builder.Finalize()
	}()
	if err != nil {
		return err
	}

	filename := filepath.Join(dir, onedb.Filename("mainnet", 0, root))
	if err := os.WriteFile(filename, buf.Bytes(), 0o644); err != nil {
		return err
	}
	sum := sha256.Sum256(buf.Bytes())
	return os.WriteFile(filepath.Join(dir, "checksums.txt"), []byte(common.BytesToHash(sum[:]).Hex()), 0o644)
}

func makeHistorySingleTxFixture(t *testing.T) (*core.Genesis, *types.Block, types.Receipts, *types.Transaction) {
	t.Helper()

	key, err := crypto.GenerateKey()
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	address := crypto.PubkeyToAddress(key.PublicKey)
	genesis := &core.Genesis{
		Config: params.TestChainConfig,
		Alloc:  types.GenesisAlloc{address: {Balance: big.NewInt(1000000000000000000)}},
	}
	signer := types.LatestSigner(genesis.Config)
	_, blocks, receipts := core.GenerateChainWithGenesis(genesis, ethash.NewFaker(), 1, func(i int, g *core.BlockGen) {
		tx, err := types.SignNewTx(key, signer, &types.LegacyTx{
			Nonce:    uint64(i),
			To:       &common.Address{0xaa},
			Value:    big.NewInt(1),
			Gas:      50_000,
			GasPrice: big.NewInt(1_000_000_000),
		})
		if err != nil {
			t.Fatalf("sign tx: %v", err)
		}
		g.AddTx(tx)
	})
	if len(blocks) != 1 || len(receipts) != 1 || len(receipts[0]) != 1 || len(blocks[0].Transactions()) != 1 {
		t.Fatalf("unexpected fixture lengths: blocks=%d receiptBlocks=%d receipts=%d txs=%d", len(blocks), len(receipts), len(receipts[0]), len(blocks[0].Transactions()))
	}
	return genesis, blocks[0], receipts[0], blocks[0].Transactions()[0]
}

func makeMergedHistoryFixture(t *testing.T) (*core.Genesis, *types.Block, types.Receipts) {
	t.Helper()

	config := *params.MergedTestChainConfig
	genesis := &core.Genesis{Config: &config}
	_, blocks, receipts := core.GenerateChainWithGenesis(genesis, beacon.New(ethash.NewFaker()), 1, nil)
	if len(blocks) != 1 || len(receipts) != 1 {
		t.Fatalf("unexpected merged fixture lengths: blocks=%d receiptBlocks=%d", len(blocks), len(receipts))
	}
	return genesis, blocks[0], receipts[0]
}

func newHistoryImportChain(t *testing.T, genesis *core.Genesis, engine consensus.Engine) *core.BlockChain {
	t.Helper()

	db, err := rawdb.Open(rawdb.NewMemoryDatabase(), rawdb.OpenOptions{})
	if err != nil {
		t.Fatalf("failed to open import db: %v", err)
	}
	t.Cleanup(func() { db.Close() })

	genesis.MustCommit(db, triedb.NewDatabase(db, triedb.HashDefaults))
	imported, err := core.NewBlockChain(db, genesis, engine, nil)
	if err != nil {
		t.Fatalf("unable to initialize imported chain: %v", err)
	}
	t.Cleanup(imported.Stop)
	return imported
}