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

import (
	"context"
	"sync/atomic"
	"testing"
	"time"

	"github.com/holiman/uint256"
	"github.com/stretchr/testify/require"
	"golang.org/x/sync/errgroup"

	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/execution/types"
	"github.com/erigontech/erigon/node/gointerfaces/sentryproto"
)

func TestTrackingFetcherFetchHeadersUpdatesPeerTracker(t *testing.T) {
	t.Parallel()

	peerId1 := PeerIdFromUint64(1)
	requestId1 := uint64(1234)
	mockInboundMessages1 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_HEADERS_66,
			PeerId: peerId1.H512(),
			Data:   newMockBlockHeadersPacket66Bytes(t, requestId1, 2),
		},
	}
	mockRequestResponse1 := requestResponseMock{
		requestId:                   requestId1,
		mockResponseInboundMessages: mockInboundMessages1,
		wantRequestPeerId:           peerId1,
		wantRequestOriginNumber:     1,
		wantRequestAmount:           2,
	}
	requestId2 := uint64(1235)
	mockInboundMessages2 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_HEADERS_66,
			PeerId: peerId1.H512(),
			// peer returns 0 headers for requestId2 - peer does not have this header range
			Data: newMockBlockHeadersPacket66Bytes(t, requestId2, 0),
		},
	}
	mockRequestResponse2 := requestResponseMock{
		requestId:                   requestId2,
		mockResponseInboundMessages: mockInboundMessages2,
		wantRequestPeerId:           peerId1,
		wantRequestOriginNumber:     3,
		wantRequestAmount:           2,
	}

	test := newTrackingFetcherTest(t, newMockRequestGenerator(requestId1, requestId2))
	test.mockSentryStreams(mockRequestResponse1, mockRequestResponse2)
	test.run(func(ctx context.Context, t *testing.T) {
		test.simulateDefaultPeerEvents()
		var peerIds []*PeerId // peers which may have blocks 1 and 2
		require.Eventuallyf(t, func() bool {
			peerIds = test.peerTracker.ListPeersMayHaveBlockNum(2)
			return len(peerIds) == 2
		}, time.Second, 100*time.Millisecond, "expected number of initial peers never satisfied: want=2, have=%d", len(peerIds))

		headers, err := test.trackingFetcher.FetchHeaders(ctx, 1, 3, peerId1) // fetch headers 1 and 2
		headersData := headers.Data
		require.NoError(t, err)
		require.Len(t, headersData, 2)
		require.Equal(t, uint64(1), headersData[0].Number.Uint64())
		require.Equal(t, uint64(2), headersData[1].Number.Uint64())

		peerIds = test.peerTracker.ListPeersMayHaveBlockNum(4) // peers which may have blocks 1,2,3,4
		require.Len(t, peerIds, 2)

		var errIncompleteHeaders *ErrIncompleteHeaders
		headers, err = test.trackingFetcher.FetchHeaders(ctx, 3, 5, peerId1) // fetch headers 3 and 4
		require.ErrorAs(t, err, &errIncompleteHeaders)                       // peer 1 does not have headers 3 and 4
		require.Equal(t, uint64(3), errIncompleteHeaders.start)
		require.Equal(t, uint64(2), errIncompleteHeaders.requested)
		require.Equal(t, uint64(0), errIncompleteHeaders.received)
		require.Equal(t, uint64(3), errIncompleteHeaders.LowestMissingBlockNum())
		require.Nil(t, headers.Data)

		// should be one peer less now given that we know that peer 1 does not have block num 4
		peerIds = test.peerTracker.ListPeersMayHaveBlockNum(4)
		require.Len(t, peerIds, 1)
	})
}

func TestTrackingFetcherFetchHeadersBackwardsUpdatesPeerTracker(t *testing.T) {
	t.Parallel()

	peerId1 := PeerIdFromUint64(1)
	requestId1 := uint64(1235)
	mockInboundMessages1 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_HEADERS_66,
			PeerId: peerId1.H512(),
			// peer returns 0 headers for requestId2 - peer does not have this header range [10,15)
			Data: newMockBlockHeadersPacket66Bytes(t, requestId1, 0),
		},
	}
	mockRequestResponse1 := requestResponseMock{
		requestId:                   requestId1,
		mockResponseInboundMessages: mockInboundMessages1,
		wantRequestPeerId:           peerId1,
		wantRequestOriginNumber:     10,
		wantRequestAmount:           5,
	}
	mockHeaders := newMockBlockHeaders(12)
	mockHash := mockHeaders[11].Hash()
	requestId2 := uint64(1236)
	mockInboundMessages2 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_HEADERS_66,
			PeerId: peerId1.H512(),
			// peer returns 3 headers from block height 12 backwards using a reverse lookup by hash
			Data: blockHeadersPacket66Bytes(t, requestId2, []*types.Header{mockHeaders[11], mockHeaders[10], mockHeaders[9]}),
		},
	}
	mockRequestResponse2 := requestResponseMock{
		requestId:                   requestId2,
		mockResponseInboundMessages: mockInboundMessages2,
		wantRequestPeerId:           peerId1,
		wantRequestOriginHash:       mockHash,
		wantRequestAmount:           3,
		wantReverse:                 true,
	}

	test := newTrackingFetcherTest(t, newMockRequestGenerator(requestId1, requestId2))
	test.mockSentryStreams(mockRequestResponse1, mockRequestResponse2)
	test.run(func(ctx context.Context, t *testing.T) {
		test.simulateDefaultPeerEvents()
		var peerIds []*PeerId // peers which may have blocks 1 and 2
		require.Eventuallyf(t, func() bool {
			peerIds = test.peerTracker.ListPeersMayHaveBlockNum(2)
			return len(peerIds) == 2
		}, time.Second, 100*time.Millisecond, "expected number of initial peers never satisfied: want=2, have=%d", len(peerIds))

		// fetch headers [10,15) - peer doesn't have them
		var errIncompleteHeaders *ErrIncompleteHeaders
		headers, err := test.trackingFetcher.FetchHeaders(ctx, 10, 15, peerId1)
		require.ErrorAs(t, err, &errIncompleteHeaders)
		require.Equal(t, uint64(10), errIncompleteHeaders.start)
		require.Equal(t, uint64(5), errIncompleteHeaders.requested)
		require.Equal(t, uint64(0), errIncompleteHeaders.received)
		require.Equal(t, uint64(10), errIncompleteHeaders.LowestMissingBlockNum())
		require.Nil(t, headers.Data)

		peerIds = test.peerTracker.ListPeersMayHaveBlockNum(10) // peer1 doesn't - so one peer less
		require.Len(t, peerIds, 1)

		// now fetch headers backwards by hash H which starts at block 12 - this time the peer has this
		headers, err = test.trackingFetcher.FetchHeadersBackwards(ctx, mockHash, 3, peerId1)
		require.NoError(t, err)
		require.Len(t, headers.Data, 3)
		require.Equal(t, uint64(10), headers.Data[0].Number.Uint64())
		require.Equal(t, uint64(11), headers.Data[1].Number.Uint64())
		require.Equal(t, uint64(12), headers.Data[2].Number.Uint64())

		peerIds = test.peerTracker.ListPeersMayHaveBlockNum(10) // peer1 has block num 10 now - so one peer more
		require.Len(t, peerIds, 2)
	})
}

func TestTrackingFetcherFetchBodiesUpdatesPeerTracker(t *testing.T) {
	if testing.Short() {
		t.Skip()
	}

	t.Parallel()

	peerId1 := PeerIdFromUint64(1)
	peerId2 := PeerIdFromUint64(2)
	requestId1 := uint64(1234)
	requestId2 := uint64(1235)
	requestId3 := uint64(1236)
	mockHeaders := []*types.Header{{Number: *uint256.NewInt(1)}}
	mockHashes := []common.Hash{mockHeaders[0].Hash()}
	mockInboundMessages1 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_BODIES_66,
			PeerId: peerId1.H512(),
			Data:   newMockBlockBodiesPacketBytes(t, requestId1),
		},
	}
	mockRequestResponse1 := requestResponseMock{
		requestId:                   requestId1,
		mockResponseInboundMessages: mockInboundMessages1,
		wantRequestPeerId:           peerId1,
		wantRequestHashes:           mockHashes,
	}
	mockInboundMessages2 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_BODIES_66,
			PeerId: peerId2.H512(),
			Data:   nil, // response timeout
		},
	}
	mockRequestResponse2 := requestResponseMock{
		requestId:                   requestId2,
		mockResponseInboundMessages: mockInboundMessages2,
		wantRequestPeerId:           peerId2,
		wantRequestHashes:           mockHashes,
		responseDelay:               600 * time.Millisecond,
	}
	mockInboundMessages3 := []*sentryproto.InboundMessage{
		{
			Id:     sentryproto.MessageId_BLOCK_BODIES_66,
			PeerId: peerId2.H512(),
			Data:   nil, // response timeout
		},
	}
	mockRequestResponse3 := requestResponseMock{
		requestId:                   requestId3,
		mockResponseInboundMessages: mockInboundMessages3,
		wantRequestPeerId:           peerId2,
		wantRequestHashes:           mockHashes,
		responseDelay:               600 * time.Millisecond,
	}

	test := newTrackingFetcherTest(t, newMockRequestGenerator(requestId1, requestId2, requestId3))
	test.mockSentryStreams(mockRequestResponse1, mockRequestResponse2, mockRequestResponse3)
	test.run(func(ctx context.Context, t *testing.T) {
		test.simulateDefaultPeerEvents()
		var peerIds []*PeerId // peers which may have block 1
		require.Eventuallyf(t, func() bool {
			peerIds = test.peerTracker.ListPeersMayHaveBlockNum(1)
			return len(peerIds) == 2
		}, time.Second, 100*time.Millisecond, "expected number of initial peers never satisfied: want=2, have=%d", len(peerIds))

		bodies, err := test.trackingFetcher.FetchBodies(ctx, mockHeaders, peerId1)
		require.ErrorIs(t, err, &ErrMissingBodies{})
		require.Nil(t, bodies.Data)

		peerIds = test.peerTracker.ListPeersMayHaveBlockNum(1) // only peerId2 may have block 1, peerId does not
		require.Len(t, peerIds, 1)

		bodies, err = test.trackingFetcher.FetchBodies(ctx, mockHeaders, peerId2)
		require.ErrorIs(t, err, context.DeadlineExceeded)
		require.Nil(t, bodies.Data)

		peerIds = test.peerTracker.ListPeersMayHaveBlockNum(1) // neither peerId1 nor peerId2 have block num 1
		require.Empty(t, peerIds)
	})
}

func newTrackingFetcherTest(t *testing.T, requestIdGenerator RequestIdGenerator) *trackingFetcherTest {
	fetcherTest := newFetcherTest(t, requestIdGenerator)
	logger := fetcherTest.logger
	messageListener := fetcherTest.messageListener
	peerTracker := NewPeerTracker(logger, messageListener, WithPreservingPeerShuffle)
	trackingFetcher := NewTrackingFetcher(fetcherTest.fetcher, peerTracker)
	return &trackingFetcherTest{
		fetcherTest:     fetcherTest,
		trackingFetcher: trackingFetcher,
		peerTracker:     peerTracker,
	}
}

type trackingFetcherTest struct {
	*fetcherTest
	trackingFetcher *TrackingFetcher
	peerTracker     *PeerTracker
}

func (tft *trackingFetcherTest) run(f func(ctx context.Context, t *testing.T)) {
	var done atomic.Bool
	tft.t.Run("start", func(t *testing.T) {
		go func() {
			defer done.Store(true)
			eg, ctx := errgroup.WithContext(tft.ctx)
			eg.Go(func() error {
				return tft.peerTracker.Run(ctx)
			})
			eg.Go(func() error {
				return tft.messageListener.Run(ctx)
			})
			err := eg.Wait()
			require.ErrorIs(t, err, context.Canceled)
		}()
		waitCond := func() bool {
			// wait for the tracker to be initialised before simulating peer events,
			// otherwise the tests may be flake-y because simulated test events in each test
			// may be consumed by the listener background loop before the tracker has
			// registered its peer event observer which makes the test unreliable
			// == 2 since there is one for the background loop and one for the new observer with replay=true
			return tft.peerEventsSubsCount() == 2
		}
		require.Eventually(t, waitCond, time.Second, 100*time.Millisecond, "expected peer tracker to be initialised")
	})

	tft.t.Run("test", func(t *testing.T) {
		f(tft.ctx, t)
	})

	tft.t.Run("stop", func(t *testing.T) {
		tft.ctxCancel()
		require.Eventually(t, done.Load, time.Second, 5*time.Millisecond)
	})
}
