package sync

import (
	"fmt"
	"sort"
	"testing"

	"github.com/OffchainLabs/prysm/v7/beacon-chain/p2p"
	p2pTypes "github.com/OffchainLabs/prysm/v7/beacon-chain/p2p/types"
	"github.com/OffchainLabs/prysm/v7/beacon-chain/startup"
	"github.com/OffchainLabs/prysm/v7/config/params"
	"github.com/OffchainLabs/prysm/v7/consensus-types/blocks"
	types "github.com/OffchainLabs/prysm/v7/consensus-types/primitives"
	"github.com/OffchainLabs/prysm/v7/encoding/bytesutil"
	"github.com/OffchainLabs/prysm/v7/genesis"
	ethpb "github.com/OffchainLabs/prysm/v7/proto/prysm/v1alpha1"
	"github.com/OffchainLabs/prysm/v7/testing/require"
	"github.com/OffchainLabs/prysm/v7/testing/util"
	"github.com/libp2p/go-libp2p/core/network"
)

func (c *blobsTestCase) defaultOldestSlotByRoot(t *testing.T) types.Slot {
	oldest, err := BlobRPCMinValidSlot(c.clock.CurrentSlot())
	require.NoError(t, err)
	return oldest
}

func blobRootRequestFromSidecars(scs []blocks.ROBlob) any {
	req := make(p2pTypes.BlobSidecarsByRootReq, 0)
	for i := range scs {
		sc := scs[i]
		req = append(req, &ethpb.BlobIdentifier{BlockRoot: sc.BlockRootSlice(), Index: sc.Index})
	}
	return &req
}

func (c *blobsTestCase) filterExpectedByRoot(t *testing.T, scs []blocks.ROBlob, r any) []*expectedBlobChunk {
	rp, ok := r.(*p2pTypes.BlobSidecarsByRootReq)
	if !ok {
		panic("unexpected request type in filterExpectedByRoot")
	}
	req := *rp
	if uint64(len(req)) > params.BeaconConfig().MaxRequestBlobSidecars {
		return []*expectedBlobChunk{{
			code:    responseCodeInvalidRequest,
			message: p2pTypes.ErrBlobLTMinRequest.Error(),
		}}
	}
	sort.Sort(&req)
	var expect []*expectedBlobChunk
	blockOffset := 0
	if len(scs) == 0 {
		return expect
	}
	lastRoot := scs[0].BlockRoot()
	rootToOffset := make(map[[32]byte]int)
	rootToOffset[lastRoot] = 0
	scMap := make(map[[32]byte]map[uint64]blocks.ROBlob)
	for i := range scs {
		sc := scs[i]
		root := sc.BlockRoot()
		if root != lastRoot {
			blockOffset += 1
			rootToOffset[root] = blockOffset
		}
		lastRoot = root
		_, ok := scMap[root]
		if !ok {
			scMap[root] = make(map[uint64]blocks.ROBlob)
		}
		scMap[root][sc.Index] = sc
	}
	for i := range req {
		scid := req[i]
		rootMap, ok := scMap[bytesutil.ToBytes32(scid.BlockRoot)]
		if !ok {
			panic(fmt.Sprintf("test setup failure, no fixture with root %#x", scid.BlockRoot))
		}
		sc, idxOk := rootMap[scid.Index]
		if !idxOk {
			panic(fmt.Sprintf("test setup failure, no fixture at index %d with root %#x", scid.Index, scid.BlockRoot))
		}
		// Skip sidecars that are supposed to be missing.
		root := sc.BlockRoot()
		if c.missing[rootToOffset[root]] {
			continue
		}
		// If a sidecar is expired, we'll expect an error for the *first* index, and after that
		// we'll expect no further chunks in the stream, so filter out any further expected responses.
		// We don't need to check what index this is because we work through them in order and the first one
		// will set streamTerminated = true and skip everything else in the test case.
		if c.expired[rootToOffset[root]] {
			return append(expect, &expectedBlobChunk{
				sidecar: &sc,
				code:    responseCodeResourceUnavailable,
				message: p2pTypes.ErrBlobLTMinRequest.Error(),
			})
		}

		expect = append(expect, &expectedBlobChunk{
			sidecar: &sc,
			code:    responseCodeSuccess,
			message: "",
		})
	}
	return expect
}

func (c *blobsTestCase) runTestBlobSidecarsByRoot(t *testing.T) {
	if c.serverHandle == nil {
		c.serverHandle = func(s *Service) rpcHandler { return s.blobSidecarByRootRPCHandler }
	}
	if c.defineExpected == nil {
		c.defineExpected = c.filterExpectedByRoot
	}
	if c.requestFromSidecars == nil {
		c.requestFromSidecars = blobRootRequestFromSidecars
	}
	if c.topic == "" {
		c.topic = p2p.RPCBlobSidecarsByRootTopicV1
	}
	if c.oldestSlot == nil {
		c.oldestSlot = c.defaultOldestSlotByRoot
	}
	if c.streamReader == nil {
		c.streamReader = defaultExpectedRequirer
	}
	if c.clock == nil {
		de := params.BeaconConfig().DenebForkEpoch
		denebBuffer := params.BeaconConfig().MinEpochsForBlobsSidecarsRequest + 1000
		ce := de + denebBuffer
		cs := util.SlotAtEpoch(t, ce)
		c.clock = startup.NewClock(genesis.Time(), genesis.ValidatorsRoot(), startup.WithSlotAsNow(cs))
	}
	c.run(t)
}

func TestReadChunkEncodedBlobs(t *testing.T) {
	cases := []*blobsTestCase{
		{
			name:         "test successful read via requester",
			nblocks:      1,
			streamReader: readChunkEncodedBlobsAsStreamReader,
		},
		{
			name:         "test peer sending excess blobs",
			nblocks:      1,
			streamReader: readChunkEncodedBlobsLowMax,
		},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			c.runTestBlobSidecarsByRoot(t)
		})
	}
}

// Specifies a max expected chunk parameter of 1, so that a response with one or more blobs will give ErrInvalidFetchedData.
func readChunkEncodedBlobsLowMax(t *testing.T, s *Service, expect []*expectedBlobChunk) func(network.Stream) {
	encoding := s.cfg.p2p.Encoding()
	ctxMap, err := ContextByteVersionsForValRoot(s.cfg.clock.GenesisValidatorsRoot())
	require.NoError(t, err)
	vf := func(sidecar blocks.ROBlob) error {
		return nil
	}
	return func(stream network.Stream) {
		_, err := readChunkEncodedBlobs(stream, encoding, ctxMap, vf, 1)
		require.ErrorIs(t, err, errMaxRequestBlobSidecarsExceeded)
	}
}

func readChunkEncodedBlobsAsStreamReader(t *testing.T, s *Service, expect []*expectedBlobChunk) func(network.Stream) {
	encoding := s.cfg.p2p.Encoding()
	ctxMap, err := ContextByteVersionsForValRoot(s.cfg.clock.GenesisValidatorsRoot())
	require.NoError(t, err)
	vf := func(sidecar blocks.ROBlob) error {
		return nil
	}
	return func(stream network.Stream) {
		scs, err := readChunkEncodedBlobs(stream, encoding, ctxMap, vf, params.BeaconConfig().MaxRequestBlobSidecars)
		require.NoError(t, err)
		require.Equal(t, len(expect), len(scs))
		for i, sc := range scs {
			esc := expect[i].sidecar
			require.Equal(t, esc.Slot(), sc.Slot())
			require.Equal(t, esc.Index, sc.Index)
			require.Equal(t, esc.BlockRoot(), sc.BlockRoot())
		}
	}
}

func TestBlobsByRootValidation(t *testing.T) {
	params.SetupTestConfigCleanup(t)
	repositionFutureEpochs(params.BeaconConfig())

	de := params.BeaconConfig().DenebForkEpoch
	denebBuffer := params.BeaconConfig().MinEpochsForBlobsSidecarsRequest + 1000
	ce := de + denebBuffer
	cs := util.SlotAtEpoch(t, ce)
	clock := startup.NewClock(genesis.Time(), genesis.ValidatorsRoot(), startup.WithSlotAsNow(cs))

	dmc := defaultMockChain(t, ce)
	capellaSlot := util.SlotAtEpoch(t, params.BeaconConfig().CapellaForkEpoch)
	dmc.Slot = &capellaSlot
	dmc.FinalizedCheckPoint = &ethpb.Checkpoint{Epoch: params.BeaconConfig().CapellaForkEpoch}
	maxBlobs := params.BeaconConfig().MaxBlobsPerBlockAtEpoch(params.BeaconConfig().DenebForkEpoch)
	cases := []*blobsTestCase{
		{
			name:    "block before minimum_request_epoch",
			nblocks: 1,
			expired: map[int]bool{0: true},
			chain:   dmc,
			clock:   clock,
			err:     p2pTypes.ErrBlobLTMinRequest,
		},
		{
			name:    "blocks before and after minimum_request_epoch",
			nblocks: 2,
			expired: map[int]bool{0: true},
			chain:   dmc,
			clock:   clock,
			err:     p2pTypes.ErrBlobLTMinRequest,
		},
		{
			name:    "one after minimum_request_epoch then one before",
			nblocks: 2,
			expired: map[int]bool{1: true},
			chain:   dmc,
			clock:   clock,
			err:     p2pTypes.ErrBlobLTMinRequest,
		},
		{
			name:    "block with all indices missing between 2 full blocks",
			nblocks: 3,
			missing: map[int]bool{1: true},
			total:   func(i int) *int { return &i }(2 * int(maxBlobs)),
		},
		{
			name:    "exceeds req max",
			nblocks: int(params.BeaconConfig().MaxRequestBlobSidecars) + 1,
			err:     p2pTypes.ErrRateLimited,
		},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			c.clock = clock
			c.runTestBlobSidecarsByRoot(t)
		})
	}
}

func TestBlobsByRootOK(t *testing.T) {
	cases := []*blobsTestCase{
		{
			name:    "0 blob",
			nblocks: 0,
		},
		{
			name:    "1 blob",
			nblocks: 1,
		},
		{
			name:    "2 blob",
			nblocks: 2,
		},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			c.runTestBlobSidecarsByRoot(t)
		})
	}
}

func TestValidateBlobByRootRequest(t *testing.T) {
	params.SetupTestConfigCleanup(t)
	cfg := params.BeaconConfig()

	// Helper function to create blob identifiers
	createBlobIdents := func(count int) p2pTypes.BlobSidecarsByRootReq {
		idents := make([]*ethpb.BlobIdentifier, count)
		for i := range count {
			idents[i] = &ethpb.BlobIdentifier{
				BlockRoot: make([]byte, 32),
				Index:     uint64(i),
			}
		}
		return idents
	}

	tests := []struct {
		name        string
		blobIdents  p2pTypes.BlobSidecarsByRootReq
		slot        types.Slot
		expectedErr error
	}{
		{
			name:        "pre-Electra: at max limit",
			blobIdents:  createBlobIdents(int(cfg.MaxRequestBlobSidecars)),
			slot:        util.SlotAtEpoch(t, cfg.ElectraForkEpoch-1),
			expectedErr: nil,
		},
		{
			name:        "pre-Electra: exceeds max limit by 1",
			blobIdents:  createBlobIdents(int(cfg.MaxRequestBlobSidecars) + 1),
			slot:        util.SlotAtEpoch(t, cfg.ElectraForkEpoch-1),
			expectedErr: p2pTypes.ErrMaxBlobReqExceeded,
		},
		{
			name:        "Electra: at max limit",
			blobIdents:  createBlobIdents(int(cfg.MaxRequestBlobSidecarsElectra)),
			slot:        util.SlotAtEpoch(t, cfg.ElectraForkEpoch),
			expectedErr: nil,
		},
		{
			name:        "Electra: exceeds Electra max limit by 1",
			blobIdents:  createBlobIdents(int(cfg.MaxRequestBlobSidecarsElectra) + 1),
			slot:        util.SlotAtEpoch(t, cfg.ElectraForkEpoch),
			expectedErr: p2pTypes.ErrMaxBlobReqExceeded,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := validateBlobByRootRequest(tt.blobIdents, tt.slot)
			if tt.expectedErr != nil {
				require.ErrorIs(t, err, tt.expectedErr)
				return
			}

			require.NoError(t, err)
		})
	}
}
