// Copyright 2025 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library 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.
//
// The go-ethereum library 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 the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package partial

import (
	"bytes"
	"testing"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/rawdb"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/core/types/bal"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/trie"
	"github.com/ethereum/go-ethereum/trie/trienode"
	"github.com/ethereum/go-ethereum/triedb"
	"github.com/holiman/uint256"
)

// testBALBuilder is a test helper for constructing BlockAccessLists.
// It wraps ConstructionBlockAccessList and provides convenience methods
// matching the test patterns (BalanceChange, NonceChange, StorageWrite, CodeChange).
type testBALBuilder struct {
	accesses bal.ConstructionBlockAccessList
}

func newTestBALBuilder() *testBALBuilder {
	return &testBALBuilder{
		accesses: make(bal.ConstructionBlockAccessList),
	}
}

func (b *testBALBuilder) ensureAccount(addr common.Address) *bal.ConstructionAccountAccesses {
	if _, ok := b.accesses[addr]; !ok {
		b.accesses[addr] = &bal.ConstructionAccountAccesses{}
	}
	return b.accesses[addr]
}

func (b *testBALBuilder) BalanceChange(txIdx uint16, addr common.Address, balance *uint256.Int) {
	acc := b.ensureAccount(addr)
	if acc.BalanceChanges == nil {
		acc.BalanceChanges = make(map[uint16]*uint256.Int)
	}
	acc.BalanceChanges[txIdx] = balance
}

func (b *testBALBuilder) NonceChange(addr common.Address, txIdx uint16, nonce uint64) {
	acc := b.ensureAccount(addr)
	if acc.NonceChanges == nil {
		acc.NonceChanges = make(map[uint16]uint64)
	}
	acc.NonceChanges[txIdx] = nonce
}

func (b *testBALBuilder) StorageWrite(txIdx uint16, addr common.Address, slot, value common.Hash) {
	acc := b.ensureAccount(addr)
	if acc.StorageWrites == nil {
		acc.StorageWrites = make(map[common.Hash]map[uint16]common.Hash)
	}
	if _, ok := acc.StorageWrites[slot]; !ok {
		acc.StorageWrites[slot] = make(map[uint16]common.Hash)
	}
	acc.StorageWrites[slot][txIdx] = value
}

func (b *testBALBuilder) CodeChange(addr common.Address, txIdx uint16, code []byte) {
	acc := b.ensureAccount(addr)
	if acc.CodeChanges == nil {
		acc.CodeChanges = make(map[uint16][]byte)
	}
	acc.CodeChanges[txIdx] = code
}

// Build converts the construction BAL to the encoding format via RLP round-trip.
func (b *testBALBuilder) Build(t *testing.T) *bal.BlockAccessList {
	t.Helper()

	var buf bytes.Buffer
	if err := b.accesses.EncodeRLP(&buf); err != nil {
		t.Fatalf("failed to encode BAL: %v", err)
	}

	var result bal.BlockAccessList
	if err := result.DecodeRLP(rlp.NewStream(bytes.NewReader(buf.Bytes()), 0)); err != nil {
		t.Fatalf("failed to decode BAL: %v", err)
	}
	return &result
}

// setupTestPartialState creates a test partial state with the given tracked contracts.
func setupTestPartialState(t *testing.T, trackedContracts []common.Address) (*PartialState, *triedb.Database, common.Hash) {
	t.Helper()

	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)

	filter := NewConfiguredFilter(trackedContracts)
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create empty state trie
	stateTrie, err := trie.NewStateTrie(trie.StateTrieID(types.EmptyRootHash), trieDB)
	if err != nil {
		t.Fatalf("failed to create state trie: %v", err)
	}
	emptyRoot, _ := stateTrie.Commit(false)

	return ps, trieDB, emptyRoot
}

// setupTestStateWithAccount creates a state trie with a single account.
func setupTestStateWithAccount(t *testing.T, trieDB *triedb.Database, addr common.Address, account *types.StateAccount) common.Hash {
	t.Helper()

	stateTrie, err := trie.NewStateTrie(trie.StateTrieID(types.EmptyRootHash), trieDB)
	if err != nil {
		t.Fatalf("failed to create state trie: %v", err)
	}

	if err := stateTrie.UpdateAccount(addr, account, 0); err != nil {
		t.Fatalf("failed to update account: %v", err)
	}

	root, nodeSet := stateTrie.Commit(false)
	if nodeSet != nil {
		merged := trienode.NewWithNodeSet(nodeSet)
		if err := trieDB.Update(root, types.EmptyRootHash, 0, merged, nil); err != nil {
			t.Fatalf("failed to update trieDB: %v", err)
		}
		if err := trieDB.Commit(root, false); err != nil {
			t.Fatalf("failed to commit trieDB: %v", err)
		}
	}

	return root
}

func TestApplyBALAndComputeRoot_EmptyBAL(t *testing.T) {
	ps, _, emptyRoot := setupTestPartialState(t, nil)

	// Apply empty BAL
	emptyBAL := bal.BlockAccessList{}
	accessList := &emptyBAL

	newRoot, _, err := ps.ApplyBALAndComputeRoot(emptyRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply empty BAL: %v", err)
	}

	// Empty BAL should result in same root
	if newRoot != emptyRoot {
		t.Errorf("expected empty root %x, got %x", emptyRoot, newRoot)
	}
}

func TestApplyBALAndComputeRoot_BalanceChange(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account
	initialBalance := uint256.NewInt(1000)
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  initialBalance,
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with balance change using ConstructionBlockAccessList
	newBalance := uint256.NewInt(2000)
	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, newBalance)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify new root is different
	if newRoot == parentRoot {
		t.Error("expected different root after balance change")
	}

	// Verify the account balance was updated
	newTrie, err := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	if err != nil {
		t.Fatalf("failed to open new trie: %v", err)
	}
	updatedAccount, err := newTrie.GetAccount(addr)
	if err != nil {
		t.Fatalf("failed to get updated account: %v", err)
	}
	if updatedAccount.Balance.Cmp(newBalance) != 0 {
		t.Errorf("expected balance %v, got %v", newBalance, updatedAccount.Balance)
	}
}

func TestApplyBALAndComputeRoot_NonceChange(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    5,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with nonce change
	cbal := newTestBALBuilder()
	cbal.NonceChange(addr, 0, 6)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify the account nonce was updated
	newTrie, err := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	if err != nil {
		t.Fatalf("failed to open new trie: %v", err)
	}
	updatedAccount, err := newTrie.GetAccount(addr)
	if err != nil {
		t.Fatalf("failed to get updated account: %v", err)
	}
	if updatedAccount.Nonce != 6 {
		t.Errorf("expected nonce 6, got %d", updatedAccount.Nonce)
	}
}

func TestApplyBALAndComputeRoot_StorageChange(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account (tracked contract)
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(0),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with storage change
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000042")

	cbal := newTestBALBuilder()
	cbal.StorageWrite(0, addr, slot, value)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify new root is different (storage changed)
	if newRoot == parentRoot {
		t.Error("expected different root after storage change")
	}

	// Verify the account storage root changed
	newTrie, err := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	if err != nil {
		t.Fatalf("failed to open new trie: %v", err)
	}
	updatedAccount, err := newTrie.GetAccount(addr)
	if err != nil {
		t.Fatalf("failed to get updated account: %v", err)
	}
	if updatedAccount.Root == types.EmptyRootHash {
		t.Error("expected non-empty storage root after storage change")
	}
}

func TestApplyBALAndComputeRoot_UntrackedContractStorageIgnored(t *testing.T) {
	trackedAddr := common.HexToAddress("0x1111111111111111111111111111111111111111")
	untrackedAddr := common.HexToAddress("0x2222222222222222222222222222222222222222")

	// Only track one contract
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{trackedAddr})

	// Create initial accounts
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}

	// Add both accounts to state
	stateTrie, _ := trie.NewStateTrie(trie.StateTrieID(types.EmptyRootHash), trieDB)
	stateTrie.UpdateAccount(trackedAddr, initialAccount, 0)
	stateTrie.UpdateAccount(untrackedAddr, initialAccount, 0)
	parentRoot, nodeSet := stateTrie.Commit(false)
	if nodeSet != nil {
		merged := trienode.NewWithNodeSet(nodeSet)
		trieDB.Update(parentRoot, types.EmptyRootHash, 0, merged, nil)
		trieDB.Commit(parentRoot, false)
	}

	// Create BAL with storage changes for both contracts
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000042")

	cbal := newTestBALBuilder()
	cbal.StorageWrite(0, trackedAddr, slot, value)
	cbal.StorageWrite(0, untrackedAddr, slot, value)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify tracked contract has storage
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	trackedAccount, _ := newTrie.GetAccount(trackedAddr)
	if trackedAccount.Root == types.EmptyRootHash {
		t.Error("tracked contract should have storage root")
	}

	// Verify untracked contract has NO storage (storage was ignored)
	untrackedAccount, _ := newTrie.GetAccount(untrackedAddr)
	if untrackedAccount.Root != types.EmptyRootHash {
		t.Error("untracked contract should have empty storage root")
	}
}

func TestApplyBALAndComputeRoot_NewAccount(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, emptyRoot := setupTestPartialState(t, []common.Address{addr})

	// Create BAL that creates a new account
	balance := uint256.NewInt(1000)

	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, balance)
	cbal.NonceChange(addr, 0, 1)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(emptyRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify new account was created
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, err := newTrie.GetAccount(addr)
	if err != nil {
		t.Fatalf("failed to get new account: %v", err)
	}
	if account == nil {
		t.Fatal("expected account to exist")
	}
	if account.Balance.Cmp(balance) != 0 {
		t.Errorf("expected balance %v, got %v", balance, account.Balance)
	}
	if account.Nonce != 1 {
		t.Errorf("expected nonce 1, got %d", account.Nonce)
	}
}

func TestApplyBALAndComputeRoot_CodeChange(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)
	filter := NewConfiguredFilter([]common.Address{addr})
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(0),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with code deployment
	code := []byte{0x60, 0x60, 0x60, 0x40, 0x52} // Some bytecode
	codeHash := crypto.Keccak256Hash(code)

	cbal := newTestBALBuilder()
	cbal.CodeChange(addr, 0, code)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify code hash was updated
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)
	if common.BytesToHash(account.CodeHash) != codeHash {
		t.Errorf("expected code hash %x, got %x", codeHash, account.CodeHash)
	}

	// Verify code was stored (tracked contract)
	storedCode := rawdb.ReadCode(db, codeHash)
	if storedCode == nil {
		t.Error("expected code to be stored for tracked contract")
	}
}

func TestApplyBALAndComputeRoot_MultipleTransactions(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with multiple balance/nonce changes (only last should apply)
	balance1 := uint256.NewInt(500)
	balance2 := uint256.NewInt(2000)
	balance3 := uint256.NewInt(1500) // Final balance

	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, balance1)
	cbal.BalanceChange(1, addr, balance2)
	cbal.BalanceChange(2, addr, balance3) // Final
	cbal.NonceChange(addr, 0, 1)
	cbal.NonceChange(addr, 1, 2)
	cbal.NonceChange(addr, 2, 3) // Final

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify only final values are applied
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)
	if account.Balance.Cmp(balance3) != 0 {
		t.Errorf("expected final balance %v, got %v", balance3, account.Balance)
	}
	if account.Nonce != 3 {
		t.Errorf("expected final nonce 3, got %d", account.Nonce)
	}
}

// ============================================================================
// Task 1: Edge Case Tests for ApplyBALAndComputeRoot
// ============================================================================

// TestApplyBALAndComputeRoot_StorageDeletion tests deleting a storage slot by writing zero value.
func TestApplyBALAndComputeRoot_StorageDeletion(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)
	filter := NewConfiguredFilter([]common.Address{addr})
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create initial account with storage
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	initialValue := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000042")

	// First, create account and add storage
	stateTrie, _ := trie.NewStateTrie(trie.StateTrieID(types.EmptyRootHash), trieDB)
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}

	// Create storage trie with initial value
	addrHash := crypto.Keccak256Hash(addr.Bytes())
	storageTrie, _ := trie.NewStateTrie(trie.StorageTrieID(types.EmptyRootHash, addrHash, types.EmptyRootHash), trieDB)
	storageTrie.UpdateStorage(addr, slot.Bytes(), initialValue.Bytes())
	storageRoot, storageNodes := storageTrie.Commit(false)

	initialAccount.Root = storageRoot
	stateTrie.UpdateAccount(addr, initialAccount, 0)
	parentRoot, accountNodes := stateTrie.Commit(false)

	// Commit storage and account nodes
	allNodes := trienode.NewMergedNodeSet()
	if storageNodes != nil {
		allNodes.Merge(storageNodes)
	}
	if accountNodes != nil {
		allNodes.Merge(accountNodes)
	}
	trieDB.Update(parentRoot, types.EmptyRootHash, 0, allNodes, nil)
	trieDB.Commit(parentRoot, false)

	// Create BAL that deletes the storage slot (write zero value)
	cbal := newTestBALBuilder()
	cbal.StorageWrite(0, addr, slot, common.Hash{}) // Zero value = delete

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify storage was deleted (root should be empty)
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)
	if account.Root != types.EmptyRootHash {
		t.Errorf("expected empty storage root after deletion, got %x", account.Root)
	}
}

// TestApplyBALAndComputeRoot_MultipleStorageWritesSameSlot tests last-write-wins for storage.
func TestApplyBALAndComputeRoot_MultipleStorageWritesSameSlot(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with multiple writes to same slot
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value1 := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value2 := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000002")
	value3 := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000003") // Final

	cbal := newTestBALBuilder()
	cbal.StorageWrite(0, addr, slot, value1)
	cbal.StorageWrite(1, addr, slot, value2)
	cbal.StorageWrite(2, addr, slot, value3) // Final value

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify only final value is stored
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)

	// Open storage trie and check value
	addrHash := crypto.Keccak256Hash(addr.Bytes())
	storageTrie, err := trie.NewStateTrie(trie.StorageTrieID(newRoot, addrHash, account.Root), trieDB)
	if err != nil {
		t.Fatalf("failed to open storage trie: %v", err)
	}

	storedValue, err := storageTrie.GetStorage(addr, slot.Bytes())
	if err != nil {
		t.Fatalf("failed to get storage: %v", err)
	}
	if common.BytesToHash(storedValue) != value3 {
		t.Errorf("expected final value %x, got %x", value3, storedValue)
	}
}

// TestApplyBALAndComputeRoot_AccountDeletion_EIP161 tests EIP-161 account deletion.
// An account should be deleted if: existed before, modified, and now empty.
func TestApplyBALAndComputeRoot_AccountDeletion_EIP161(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account with balance
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL that empties the account
	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, uint256.NewInt(0)) // Zero balance

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify account was deleted (EIP-161: empty account should be removed)
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, err := newTrie.GetAccount(addr)
	if err != nil {
		t.Fatalf("failed to get account: %v", err)
	}
	if account != nil {
		t.Errorf("expected account to be deleted (EIP-161), but it still exists with balance %v", account.Balance)
	}
}

// TestApplyBALAndComputeRoot_NeverExistedEmptyAccount tests that empty accounts that never existed
// are not added to the trie.
func TestApplyBALAndComputeRoot_NeverExistedEmptyAccount(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, emptyRoot := setupTestPartialState(t, []common.Address{addr})

	// Create BAL that "touches" an account but leaves it empty
	// This simulates an account that receives 0 balance and sends 0 balance
	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, uint256.NewInt(0)) // Zero balance on never-existed account

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(emptyRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Root should be the same as empty root (no account added)
	if newRoot != emptyRoot {
		t.Errorf("expected empty root (no account added), got different root")
	}

	// Verify account does not exist
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)
	if account != nil {
		t.Errorf("expected no account (never existed + empty), but found one")
	}
}

// TestApplyBALAndComputeRoot_CodeChangeUntracked tests that code hash is updated for untracked
// contracts but the code bytes are NOT stored.
func TestApplyBALAndComputeRoot_CodeChangeUntracked(t *testing.T) {
	trackedAddr := common.HexToAddress("0x1111111111111111111111111111111111111111")
	untrackedAddr := common.HexToAddress("0x2222222222222222222222222222222222222222")

	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)
	// Only track one contract
	filter := NewConfiguredFilter([]common.Address{trackedAddr})
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create initial accounts
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}

	stateTrie, _ := trie.NewStateTrie(trie.StateTrieID(types.EmptyRootHash), trieDB)
	stateTrie.UpdateAccount(trackedAddr, initialAccount, 0)
	stateTrie.UpdateAccount(untrackedAddr, initialAccount, 0)
	parentRoot, nodeSet := stateTrie.Commit(false)
	if nodeSet != nil {
		merged := trienode.NewWithNodeSet(nodeSet)
		trieDB.Update(parentRoot, types.EmptyRootHash, 0, merged, nil)
		trieDB.Commit(parentRoot, false)
	}

	// Create BAL with code changes for both
	code := []byte{0x60, 0x60, 0x60, 0x40, 0x52}
	codeHash := crypto.Keccak256Hash(code)

	cbal := newTestBALBuilder()
	cbal.CodeChange(trackedAddr, 0, code)
	cbal.CodeChange(untrackedAddr, 0, code)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify both accounts have updated code hash
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)

	trackedAccount, _ := newTrie.GetAccount(trackedAddr)
	if common.BytesToHash(trackedAccount.CodeHash) != codeHash {
		t.Errorf("tracked contract should have updated code hash")
	}

	untrackedAccount, _ := newTrie.GetAccount(untrackedAddr)
	if common.BytesToHash(untrackedAccount.CodeHash) != codeHash {
		t.Errorf("untracked contract should have updated code hash")
	}

	// Verify code is stored for tracked contract
	storedCode := rawdb.ReadCode(db, codeHash)
	if storedCode == nil {
		t.Error("code should be stored for tracked contract")
	}

	// Note: We can't directly test that code is NOT stored for untracked because
	// both contracts use the same code hash, and it's stored once for the tracked one.
	// The key invariant is that the code hash is correct for both.
}

// TestApplyBALAndComputeRoot_MixedChanges tests applying multiple types of changes to one account.
func TestApplyBALAndComputeRoot_MixedChanges(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)
	filter := NewConfiguredFilter([]common.Address{addr})
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    5,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Create BAL with balance, nonce, code, and storage changes
	newBalance := uint256.NewInt(2000)
	newNonce := uint64(10)
	code := []byte{0x60, 0x60, 0x60, 0x40, 0x52}
	codeHash := crypto.Keccak256Hash(code)
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000042")

	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, newBalance)
	cbal.NonceChange(addr, 0, newNonce)
	cbal.CodeChange(addr, 0, code)
	cbal.StorageWrite(0, addr, slot, value)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify all changes applied
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)

	if account.Balance.Cmp(newBalance) != 0 {
		t.Errorf("expected balance %v, got %v", newBalance, account.Balance)
	}
	if account.Nonce != newNonce {
		t.Errorf("expected nonce %d, got %d", newNonce, account.Nonce)
	}
	if common.BytesToHash(account.CodeHash) != codeHash {
		t.Errorf("expected code hash %x, got %x", codeHash, account.CodeHash)
	}
	if account.Root == types.EmptyRootHash {
		t.Error("expected non-empty storage root")
	}

	// Verify storage value
	addrHash := crypto.Keccak256Hash(addr.Bytes())
	storageTrie, _ := trie.NewStateTrie(trie.StorageTrieID(newRoot, addrHash, account.Root), trieDB)
	storedValue, _ := storageTrie.GetStorage(addr, slot.Bytes())
	if common.BytesToHash(storedValue) != value {
		t.Errorf("expected storage value %x, got %x", value, storedValue)
	}
}

// ============================================================================
// Task 3: Error Path Tests
// ============================================================================

// TestApplyBALAndComputeRoot_ErrorInvalidParentRoot tests error handling for invalid parent root.
func TestApplyBALAndComputeRoot_ErrorInvalidParentRoot(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, _, _ := setupTestPartialState(t, []common.Address{addr})

	// Use a non-existent root
	invalidRoot := common.HexToHash("0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef")

	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, uint256.NewInt(1000))
	accessList := cbal.Build(t)

	_, _, err := ps.ApplyBALAndComputeRoot(invalidRoot, common.Hash{}, accessList)
	if err == nil {
		t.Fatal("expected error for invalid parent root, got nil")
	}
	// Error should mention trie opening failure
	if !bytes.Contains([]byte(err.Error()), []byte("failed to open state trie")) {
		t.Errorf("expected 'failed to open state trie' error, got: %v", err)
	}
}

// ============================================================================
// Task 4: isEmptyAccount Tests
// ============================================================================

// TestIsEmptyAccount tests the EIP-161 empty account detection logic.
func TestIsEmptyAccount(t *testing.T) {
	ps, _, _ := setupTestPartialState(t, nil)

	tests := []struct {
		name     string
		account  *types.StateAccount
		expected bool
	}{
		{
			name: "completely empty account",
			account: &types.StateAccount{
				Nonce:    0,
				Balance:  uint256.NewInt(0),
				Root:     types.EmptyRootHash,
				CodeHash: types.EmptyCodeHash.Bytes(),
			},
			expected: true,
		},
		{
			name: "non-zero balance",
			account: &types.StateAccount{
				Nonce:    0,
				Balance:  uint256.NewInt(1),
				Root:     types.EmptyRootHash,
				CodeHash: types.EmptyCodeHash.Bytes(),
			},
			expected: false,
		},
		{
			name: "non-zero nonce",
			account: &types.StateAccount{
				Nonce:    1,
				Balance:  uint256.NewInt(0),
				Root:     types.EmptyRootHash,
				CodeHash: types.EmptyCodeHash.Bytes(),
			},
			expected: false,
		},
		{
			name: "non-empty storage root",
			account: &types.StateAccount{
				Nonce:    0,
				Balance:  uint256.NewInt(0),
				Root:     common.HexToHash("0x1234567890123456789012345678901234567890123456789012345678901234"),
				CodeHash: types.EmptyCodeHash.Bytes(),
			},
			expected: false,
		},
		{
			name: "non-empty code hash",
			account: &types.StateAccount{
				Nonce:    0,
				Balance:  uint256.NewInt(0),
				Root:     types.EmptyRootHash,
				CodeHash: common.HexToHash("0x1234567890123456789012345678901234567890123456789012345678901234").Bytes(),
			},
			expected: false,
		},
		{
			name: "large balance (uint256)",
			account: &types.StateAccount{
				Nonce:    0,
				Balance:  uint256.MustFromHex("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
				Root:     types.EmptyRootHash,
				CodeHash: types.EmptyCodeHash.Bytes(),
			},
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := ps.isEmptyAccount(tt.account)
			if result != tt.expected {
				t.Errorf("isEmptyAccount() = %v, expected %v", result, tt.expected)
			}
		})
	}
}

// ============================================================================
// Task 2: buildStateSet Tests (indirect verification)
// ============================================================================

// TestBuildStateSet_AccountModification verifies that modified accounts are correctly
// tracked in the StateSet by checking the resulting state.
func TestBuildStateSet_AccountModification(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    5,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Apply balance change
	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, uint256.NewInt(2000))
	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify the state was correctly updated (indirectly tests buildStateSet)
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)

	// The nonce should be preserved (not modified)
	if account.Nonce != 5 {
		t.Errorf("nonce should be preserved: expected 5, got %d", account.Nonce)
	}
	// Balance should be updated
	if account.Balance.Cmp(uint256.NewInt(2000)) != 0 {
		t.Errorf("balance should be updated: expected 2000, got %v", account.Balance)
	}
}

// TestBuildStateSet_StorageRLPEncoding verifies that storage values are correctly
// RLP-encoded in the StateSet.
func TestBuildStateSet_StorageRLPEncoding(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)
	filter := NewConfiguredFilter([]common.Address{addr})
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create initial account
	initialAccount := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Write storage value
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000042")

	cbal := newTestBALBuilder()
	cbal.StorageWrite(0, addr, slot, value)
	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify storage is readable
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)

	addrHash := crypto.Keccak256Hash(addr.Bytes())
	storageTrie, err := trie.NewStateTrie(trie.StorageTrieID(newRoot, addrHash, account.Root), trieDB)
	if err != nil {
		t.Fatalf("failed to open storage trie: %v", err)
	}

	storedValue, err := storageTrie.GetStorage(addr, slot.Bytes())
	if err != nil {
		t.Fatalf("failed to get storage: %v", err)
	}

	if common.BytesToHash(storedValue) != value {
		t.Errorf("storage value mismatch: expected %x, got %x", value, storedValue)
	}
}

// TestBuildStateSet_OriginTracking verifies that account origins are tracked correctly
// for PathDB compatibility.
func TestBuildStateSet_OriginTracking(t *testing.T) {
	addr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	ps, trieDB, _ := setupTestPartialState(t, []common.Address{addr})

	// Create initial account with specific values
	initialAccount := &types.StateAccount{
		Nonce:    10,
		Balance:  uint256.NewInt(5000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	parentRoot := setupTestStateWithAccount(t, trieDB, addr, initialAccount)

	// Modify the account
	cbal := newTestBALBuilder()
	cbal.BalanceChange(0, addr, uint256.NewInt(6000))
	cbal.NonceChange(addr, 0, 11)
	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify the new state is correct (origin tracking happens internally)
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)
	account, _ := newTrie.GetAccount(addr)

	if account.Nonce != 11 {
		t.Errorf("expected nonce 11, got %d", account.Nonce)
	}
	if account.Balance.Cmp(uint256.NewInt(6000)) != 0 {
		t.Errorf("expected balance 6000, got %v", account.Balance)
	}

	// The fact that this works with PathDB verifies origin tracking is correct
	// (PathDB requires origins for diff computation)
}

// TestApplyBALAndComputeRoot_MultipleAccountTypes tests processing multiple accounts with
// different modification patterns in one block.
func TestApplyBALAndComputeRoot_MultipleAccountTypes(t *testing.T) {
	addr1 := common.HexToAddress("0x1111111111111111111111111111111111111111") // Balance only
	addr2 := common.HexToAddress("0x2222222222222222222222222222222222222222") // Storage only
	addr3 := common.HexToAddress("0x3333333333333333333333333333333333333333") // New account

	db := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(db, triedb.HashDefaults)
	filter := NewConfiguredFilter([]common.Address{addr1, addr2, addr3})
	ps := NewPartialState(db, trieDB, filter, 256)

	// Create initial accounts for addr1 and addr2
	initialAccount1 := &types.StateAccount{
		Nonce:    0,
		Balance:  uint256.NewInt(1000),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}
	initialAccount2 := &types.StateAccount{
		Nonce:    5,
		Balance:  uint256.NewInt(500),
		Root:     types.EmptyRootHash,
		CodeHash: types.EmptyCodeHash.Bytes(),
	}

	stateTrie, _ := trie.NewStateTrie(trie.StateTrieID(types.EmptyRootHash), trieDB)
	stateTrie.UpdateAccount(addr1, initialAccount1, 0)
	stateTrie.UpdateAccount(addr2, initialAccount2, 0)
	parentRoot, nodeSet := stateTrie.Commit(false)
	if nodeSet != nil {
		merged := trienode.NewWithNodeSet(nodeSet)
		trieDB.Update(parentRoot, types.EmptyRootHash, 0, merged, nil)
		trieDB.Commit(parentRoot, false)
	}

	// Create BAL with different changes for each account
	cbal := newTestBALBuilder()

	// addr1: balance change
	cbal.BalanceChange(0, addr1, uint256.NewInt(2000))

	// addr2: storage write
	slot := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
	value := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000042")
	cbal.StorageWrite(0, addr2, slot, value)

	// addr3: new account
	cbal.BalanceChange(0, addr3, uint256.NewInt(3000))
	cbal.NonceChange(addr3, 0, 1)

	accessList := cbal.Build(t)

	newRoot, _, err := ps.ApplyBALAndComputeRoot(parentRoot, common.Hash{}, accessList)
	if err != nil {
		t.Fatalf("failed to apply BAL: %v", err)
	}

	// Verify all accounts
	newTrie, _ := trie.NewStateTrie(trie.StateTrieID(newRoot), trieDB)

	// addr1: balance changed
	acc1, _ := newTrie.GetAccount(addr1)
	if acc1.Balance.Cmp(uint256.NewInt(2000)) != 0 {
		t.Errorf("addr1: expected balance 2000, got %v", acc1.Balance)
	}

	// addr2: storage changed
	acc2, _ := newTrie.GetAccount(addr2)
	if acc2.Root == types.EmptyRootHash {
		t.Error("addr2: expected non-empty storage root")
	}

	// addr3: new account created
	acc3, _ := newTrie.GetAccount(addr3)
	if acc3 == nil {
		t.Fatal("addr3: expected account to exist")
	}
	if acc3.Balance.Cmp(uint256.NewInt(3000)) != 0 {
		t.Errorf("addr3: expected balance 3000, got %v", acc3.Balance)
	}
	if acc3.Nonce != 1 {
		t.Errorf("addr3: expected nonce 1, got %d", acc3.Nonce)
	}
}