package stages

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/erigontech/erigon/cl/beacon/beaconevents"
	"github.com/erigontech/erigon/cl/clparams"
	"github.com/erigontech/erigon/cl/persistence/beacon_indicies"
	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/db/kv"
	"github.com/erigontech/erigon/db/kv/dbcfg"
	"github.com/erigontech/erigon/db/kv/memdb"
)

func TestUpdateCanonicalChainReorgEvent(t *testing.T) {
	db := memdb.NewTestDB(t, dbcfg.ChainDB)
	defer db.Close()

	ctx := context.Background()
	tx, err := db.BeginRw(ctx)
	require.NoError(t, err)
	defer tx.Rollback()

	root100 := common.Hash{0x10}
	root101a := common.Hash{0x11, 0xaa}
	root102a := common.Hash{0x12, 0xaa}
	root101b := common.Hash{0x11, 0xbb}
	root102b := common.Hash{0x12, 0xbb}

	state100 := common.Hash{0xa0}
	state101a := common.Hash{0xa1}
	state102a := common.Hash{0xa2}
	state101b := common.Hash{0xb1}
	state102b := common.Hash{0xb2}

	writeBlock := func(root, parentRoot, stateRoot common.Hash, slot uint64, canonical bool) {
		t.Helper()
		require.NoError(t, beacon_indicies.WriteHeaderSlot(tx, root, slot))
		require.NoError(t, beacon_indicies.WriteParentBlockRoot(ctx, tx, root, parentRoot))
		require.NoError(t, beacon_indicies.WriteStateRoot(tx, root, stateRoot))
		if canonical {
			require.NoError(t, beacon_indicies.MarkRootCanonical(ctx, tx, slot, root))
		}
	}

	writeBlock(root100, common.Hash{0x99}, state100, 100, true)
	writeBlock(root101a, root100, state101a, 101, true)
	writeBlock(root102a, root101a, state102a, 102, true)

	writeBlock(root101b, root100, state101b, 101, false)
	writeBlock(root102b, root101b, state102b, 102, false)

	reorg := drainReorgEvent(t, ctx, tx, 102, root102b)
	require.NotNil(t, reorg, "expected a chain_reorg event to be emitted")
	require.Equal(t, uint64(102), reorg.Slot, "reorg Slot")
	require.Equal(t, uint64(2), reorg.Depth, "reorg Depth should be oldHeadSlot - forkPointSlot")
	require.Equal(t, root102a, reorg.OldHeadBlock, "OldHeadBlock should be the previous canonical tip")
	require.Equal(t, root102b, reorg.NewHeadBlock, "NewHeadBlock")
	require.Equal(t, state102a, reorg.OldHeadState, "OldHeadState should match old canonical tip's state root")
	require.Equal(t, state102b, reorg.NewHeadState, "NewHeadState")
}

func TestUpdateCanonicalChainReorgShorterFork(t *testing.T) {
	db := memdb.NewTestDB(t, dbcfg.ChainDB)
	defer db.Close()

	ctx := context.Background()
	tx, err := db.BeginRw(ctx)
	require.NoError(t, err)
	defer tx.Rollback()

	root100 := common.Hash{0x10}
	root101a := common.Hash{0x11, 0xaa}
	root102a := common.Hash{0x12, 0xaa}
	root103a := common.Hash{0x13, 0xaa}
	root101b := common.Hash{0x11, 0xbb}
	root102b := common.Hash{0x12, 0xbb}

	state100 := common.Hash{0xa0}
	state103a := common.Hash{0xa3}
	state102b := common.Hash{0xb2}

	writeBlock := func(root, parentRoot, stateRoot common.Hash, slot uint64, canonical bool) {
		t.Helper()
		require.NoError(t, beacon_indicies.WriteHeaderSlot(tx, root, slot))
		require.NoError(t, beacon_indicies.WriteParentBlockRoot(ctx, tx, root, parentRoot))
		require.NoError(t, beacon_indicies.WriteStateRoot(tx, root, stateRoot))
		if canonical {
			require.NoError(t, beacon_indicies.MarkRootCanonical(ctx, tx, slot, root))
		}
	}

	writeBlock(root100, common.Hash{0x99}, state100, 100, true)
	writeBlock(root101a, root100, common.Hash{0xa1}, 101, true)
	writeBlock(root102a, root101a, common.Hash{0xa2}, 102, true)
	writeBlock(root103a, root102a, state103a, 103, true)

	writeBlock(root101b, root100, common.Hash{0xb1}, 101, false)
	writeBlock(root102b, root101b, state102b, 102, false)

	reorg := drainReorgEvent(t, ctx, tx, 102, root102b)
	require.NotNil(t, reorg, "expected a chain_reorg event to be emitted")
	require.Equal(t, uint64(102), reorg.Slot, "reorg Slot")
	require.Equal(t, uint64(3), reorg.Depth, "reorg Depth: old tip 103 - fork point 100 = 3")
	require.Equal(t, root103a, reorg.OldHeadBlock, "OldHeadBlock must be the actual old tip at slot 103, not 102a")
	require.Equal(t, root102b, reorg.NewHeadBlock, "NewHeadBlock")
	require.Equal(t, state103a, reorg.OldHeadState, "OldHeadState must match old tip's state root")
	require.Equal(t, state102b, reorg.NewHeadState, "NewHeadState")
}

func TestUpdateCanonicalChainReorgLongerFork(t *testing.T) {
	db := memdb.NewTestDB(t, dbcfg.ChainDB)
	defer db.Close()

	ctx := context.Background()
	tx, err := db.BeginRw(ctx)
	require.NoError(t, err)
	defer tx.Rollback()

	root100 := common.Hash{0x10}
	root101a := common.Hash{0x11, 0xaa}
	root102a := common.Hash{0x12, 0xaa}
	root101b := common.Hash{0x11, 0xbb}
	root102b := common.Hash{0x12, 0xbb}
	root103b := common.Hash{0x13, 0xbb}

	state102a := common.Hash{0xa2}
	state103b := common.Hash{0xb3}

	writeBlock := func(root, parentRoot, stateRoot common.Hash, slot uint64, canonical bool) {
		t.Helper()
		require.NoError(t, beacon_indicies.WriteHeaderSlot(tx, root, slot))
		require.NoError(t, beacon_indicies.WriteParentBlockRoot(ctx, tx, root, parentRoot))
		require.NoError(t, beacon_indicies.WriteStateRoot(tx, root, stateRoot))
		if canonical {
			require.NoError(t, beacon_indicies.MarkRootCanonical(ctx, tx, slot, root))
		}
	}

	writeBlock(root100, common.Hash{0x99}, common.Hash{0xa0}, 100, true)
	writeBlock(root101a, root100, common.Hash{0xa1}, 101, true)
	writeBlock(root102a, root101a, state102a, 102, true)

	writeBlock(root101b, root100, common.Hash{0xb1}, 101, false)
	writeBlock(root102b, root101b, common.Hash{0xb2}, 102, false)
	writeBlock(root103b, root102b, state103b, 103, false)

	reorg := drainReorgEvent(t, ctx, tx, 103, root103b)
	require.NotNil(t, reorg, "expected a chain_reorg event to be emitted")
	require.Equal(t, uint64(103), reorg.Slot, "reorg Slot")
	require.Equal(t, uint64(2), reorg.Depth, "reorg Depth: old tip 102 - fork point 100 = 2")
	require.Equal(t, root102a, reorg.OldHeadBlock, "OldHeadBlock must be the old tip at slot 102")
	require.Equal(t, root103b, reorg.NewHeadBlock, "NewHeadBlock")
	require.Equal(t, state102a, reorg.OldHeadState, "OldHeadState")
	require.Equal(t, state103b, reorg.NewHeadState, "NewHeadState")
}

func TestUpdateCanonicalChainNoReorg(t *testing.T) {
	db := memdb.NewTestDB(t, dbcfg.ChainDB)
	defer db.Close()

	ctx := context.Background()
	tx, err := db.BeginRw(ctx)
	require.NoError(t, err)
	defer tx.Rollback()

	root100 := common.Hash{0x10}
	root101 := common.Hash{0x11}
	root102 := common.Hash{0x12}

	writeBlock := func(root, parentRoot, stateRoot common.Hash, slot uint64, canonical bool) {
		t.Helper()
		require.NoError(t, beacon_indicies.WriteHeaderSlot(tx, root, slot))
		require.NoError(t, beacon_indicies.WriteParentBlockRoot(ctx, tx, root, parentRoot))
		require.NoError(t, beacon_indicies.WriteStateRoot(tx, root, stateRoot))
		if canonical {
			require.NoError(t, beacon_indicies.MarkRootCanonical(ctx, tx, slot, root))
		}
	}

	writeBlock(root100, common.Hash{0x99}, common.Hash{0xa0}, 100, true)
	writeBlock(root101, root100, common.Hash{0xa1}, 101, true)

	writeBlock(root102, root101, common.Hash{0xa2}, 102, false)

	reorg := drainReorgEvent(t, ctx, tx, 102, root102)
	require.Nil(t, reorg, "chain extension should NOT emit a chain_reorg event")
}

func TestUpdateCanonicalChainReorgOneSlot(t *testing.T) {
	db := memdb.NewTestDB(t, dbcfg.ChainDB)
	defer db.Close()

	ctx := context.Background()
	tx, err := db.BeginRw(ctx)
	require.NoError(t, err)
	defer tx.Rollback()

	root100 := common.Hash{0x10}
	root101a := common.Hash{0x11, 0xaa}
	root101b := common.Hash{0x11, 0xbb}

	state101a := common.Hash{0xa1}
	state101b := common.Hash{0xb1}

	writeBlock := func(root, parentRoot, stateRoot common.Hash, slot uint64, canonical bool) {
		t.Helper()
		require.NoError(t, beacon_indicies.WriteHeaderSlot(tx, root, slot))
		require.NoError(t, beacon_indicies.WriteParentBlockRoot(ctx, tx, root, parentRoot))
		require.NoError(t, beacon_indicies.WriteStateRoot(tx, root, stateRoot))
		if canonical {
			require.NoError(t, beacon_indicies.MarkRootCanonical(ctx, tx, slot, root))
		}
	}

	writeBlock(root100, common.Hash{0x99}, common.Hash{0xa0}, 100, true)
	writeBlock(root101a, root100, state101a, 101, true)
	writeBlock(root101b, root100, state101b, 101, false)

	reorg := drainReorgEvent(t, ctx, tx, 101, root101b)
	require.NotNil(t, reorg, "expected a chain_reorg event to be emitted")
	require.Equal(t, uint64(101), reorg.Slot, "reorg Slot")
	require.Equal(t, uint64(1), reorg.Depth, "reorg Depth: old tip 101 - fork point 100 = 1")
	require.Equal(t, root101a, reorg.OldHeadBlock, "OldHeadBlock must be root_101a")
	require.Equal(t, root101b, reorg.NewHeadBlock, "NewHeadBlock")
	require.Equal(t, state101a, reorg.OldHeadState, "OldHeadState")
	require.Equal(t, state101b, reorg.NewHeadState, "NewHeadState")
}

func drainReorgEvent(t *testing.T, ctx context.Context, tx kv.RwTx, headSlot uint64, headRoot common.Hash) *beaconevents.ChainReorgData {
	t.Helper()
	emitter := beaconevents.NewEventEmitter()
	ch := make(chan *beaconevents.EventStream, 16)
	sub := emitter.State().Subscribe(ch)
	defer sub.Unsubscribe()

	cfg := &Cfg{
		emitter:   emitter,
		beaconCfg: &clparams.MainnetBeaconConfig,
	}

	err := updateCanonicalChainInTheDatabase(ctx, tx, headSlot, headRoot, cfg)
	require.NoError(t, err)

	for {
		select {
		case evt := <-ch:
			if evt.Event == beaconevents.StateChainReorg {
				return evt.Data.(*beaconevents.ChainReorgData)
			}
		default:
			return nil
		}
	}
}
