package client

import (
	"context"
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/pkg/errors"
	"github.com/prysmaticlabs/go-bitfield"
	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
	"github.com/prysmaticlabs/prysm/v5/crypto/bls"
	"github.com/prysmaticlabs/prysm/v5/encoding/bytesutil"
	ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
	"github.com/prysmaticlabs/prysm/v5/testing/assert"
	"github.com/prysmaticlabs/prysm/v5/testing/require"
	logTest "github.com/sirupsen/logrus/hooks/test"
	"go.uber.org/mock/gomock"
	"google.golang.org/protobuf/types/known/emptypb"
)

func TestSubmitSyncCommitteeMessage_ValidatorDutiesRequestFailure(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			hook := logTest.NewGlobal()
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{}}
			defer finish()

			m.validatorClient.EXPECT().GetSyncMessageBlockRoot(
				gomock.Any(), // ctx
				&emptypb.Empty{},
			).Return(&ethpb.SyncMessageBlockRootResponse{
				Root: bytesutil.PadTo([]byte{}, 32),
			}, nil)

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			validator.SubmitSyncCommitteeMessage(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not fetch validator assignment")
		})
	}
}

func TestSubmitSyncCommitteeMessage_BadDomainData(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			defer finish()
			hook := logTest.NewGlobal()
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}

			r := []byte{'a'}
			m.validatorClient.EXPECT().GetSyncMessageBlockRoot(
				gomock.Any(), // ctx
				&emptypb.Empty{},
			).Return(&ethpb.SyncMessageBlockRootResponse{
				Root: bytesutil.PadTo(r, 32),
			}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), gomock.Any()).
				Return(nil, errors.New("uh oh"))

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			validator.SubmitSyncCommitteeMessage(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not get sync committee domain data")
		})
	}
}

func TestSubmitSyncCommitteeMessage_CouldNotSubmit(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			defer finish()
			hook := logTest.NewGlobal()
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}

			r := []byte{'a'}
			m.validatorClient.EXPECT().GetSyncMessageBlockRoot(
				gomock.Any(), // ctx
				&emptypb.Empty{},
			).Return(&ethpb.SyncMessageBlockRootResponse{
				Root: bytesutil.PadTo(r, 32),
			}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			m.validatorClient.EXPECT().SubmitSyncMessage(
				gomock.Any(), // ctx
				gomock.AssignableToTypeOf(&ethpb.SyncCommitteeMessage{}),
			).Return(&emptypb.Empty{}, errors.New("uh oh") /* error */)

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			validator.SubmitSyncCommitteeMessage(context.Background(), 1, pubKey)

			require.LogsContain(t, hook, "Could not submit sync committee message")
		})
	}
}

func TestSubmitSyncCommitteeMessage_OK(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			defer finish()
			hook := logTest.NewGlobal()
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}

			r := []byte{'a'}
			m.validatorClient.EXPECT().GetSyncMessageBlockRoot(
				gomock.Any(), // ctx
				&emptypb.Empty{},
			).Return(&ethpb.SyncMessageBlockRootResponse{
				Root: bytesutil.PadTo(r, 32),
			}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			var generatedMsg *ethpb.SyncCommitteeMessage
			m.validatorClient.EXPECT().SubmitSyncMessage(
				gomock.Any(), // ctx
				gomock.AssignableToTypeOf(&ethpb.SyncCommitteeMessage{}),
			).Do(func(_ context.Context, msg *ethpb.SyncCommitteeMessage) {
				generatedMsg = msg
			}).Return(&emptypb.Empty{}, nil /* error */)

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			validator.SubmitSyncCommitteeMessage(context.Background(), 1, pubKey)

			require.LogsDoNotContain(t, hook, "Could not")
			require.Equal(t, primitives.Slot(1), generatedMsg.Slot)
			require.Equal(t, validatorIndex, generatedMsg.ValidatorIndex)
			require.DeepEqual(t, bytesutil.PadTo(r, 32), generatedMsg.BlockRoot)
		})
	}
}

func TestSubmitSignedContributionAndProof_ValidatorDutiesRequestFailure(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			hook := logTest.NewGlobal()
			validator, _, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not fetch validator assignment")
		})
	}
}

func TestSubmitSignedContributionAndProof_GetSyncSubcommitteeIndexFailure(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			hook := logTest.NewGlobal()
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			m.validatorClient.EXPECT().GetSyncSubcommitteeIndex(
				gomock.Any(), // ctx
				&ethpb.SyncSubcommitteeIndexRequest{
					Slot:      1,
					PublicKey: pubKey[:],
				},
			).Return(&ethpb.SyncSubcommitteeIndexResponse{}, errors.New("Bad index"))

			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not get sync subcommittee index")
		})
	}
}

func TestSubmitSignedContributionAndProof_NothingToDo(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			hook := logTest.NewGlobal()
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			m.validatorClient.EXPECT().GetSyncSubcommitteeIndex(
				gomock.Any(), // ctx
				&ethpb.SyncSubcommitteeIndexRequest{
					Slot:      1,
					PublicKey: pubKey[:],
				},
			).Return(&ethpb.SyncSubcommitteeIndexResponse{Indices: []primitives.CommitteeIndex{}}, nil)

			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Empty subcommittee index list, do nothing")
		})
	}
}

func TestSubmitSignedContributionAndProof_BadDomain(t *testing.T) {
	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			hook := logTest.NewGlobal()
			validator, m, validatorKey, finish := setup(t, isSlashingProtectionMinimal)
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			m.validatorClient.EXPECT().GetSyncSubcommitteeIndex(
				gomock.Any(), // ctx
				&ethpb.SyncSubcommitteeIndexRequest{
					Slot:      1,
					PublicKey: pubKey[:],
				},
			).Return(&ethpb.SyncSubcommitteeIndexResponse{Indices: []primitives.CommitteeIndex{1}}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, errors.New("bad domain response"))

			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not get selection proofs")
			require.LogsContain(t, hook, "bad domain response")
		})
	}
}

func TestSubmitSignedContributionAndProof_CouldNotGetContribution(t *testing.T) {
	hook := logTest.NewGlobal()
	// Hardcode secret key in order to have a valid aggregator signature.
	rawKey, err := hex.DecodeString("659e875e1b062c03f2f2a57332974d475b97df6cfc581d322e79642d39aca8fd")
	assert.NoError(t, err)
	validatorKey, err := bls.SecretKeyFromBytes(rawKey)
	assert.NoError(t, err)

	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			validator, m, validatorKey, finish := setupWithKey(t, validatorKey, isSlashingProtectionMinimal)
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			m.validatorClient.EXPECT().GetSyncSubcommitteeIndex(
				gomock.Any(), // ctx
				&ethpb.SyncSubcommitteeIndexRequest{
					Slot:      1,
					PublicKey: pubKey[:],
				},
			).Return(&ethpb.SyncSubcommitteeIndexResponse{Indices: []primitives.CommitteeIndex{1}}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			m.validatorClient.EXPECT().GetSyncCommitteeContribution(
				gomock.Any(), // ctx
				&ethpb.SyncCommitteeContributionRequest{
					Slot:      1,
					PublicKey: pubKey[:],
					SubnetId:  0,
				},
			).Return(nil, errors.New("Bad contribution"))

			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not get sync committee contribution")
		})
	}
}

func TestSubmitSignedContributionAndProof_CouldNotSubmitContribution(t *testing.T) {
	hook := logTest.NewGlobal()
	// Hardcode secret key in order to have a valid aggregator signature.
	rawKey, err := hex.DecodeString("659e875e1b062c03f2f2a57332974d475b97df6cfc581d322e79642d39aca8fd")
	assert.NoError(t, err)
	validatorKey, err := bls.SecretKeyFromBytes(rawKey)
	assert.NoError(t, err)

	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			validator, m, validatorKey, finish := setupWithKey(t, validatorKey, isSlashingProtectionMinimal)
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			m.validatorClient.EXPECT().GetSyncSubcommitteeIndex(
				gomock.Any(), // ctx
				&ethpb.SyncSubcommitteeIndexRequest{
					Slot:      1,
					PublicKey: pubKey[:],
				},
			).Return(&ethpb.SyncSubcommitteeIndexResponse{Indices: []primitives.CommitteeIndex{1}}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			aggBits := bitfield.NewBitvector128()
			aggBits.SetBitAt(0, true)
			m.validatorClient.EXPECT().GetSyncCommitteeContribution(
				gomock.Any(), // ctx
				&ethpb.SyncCommitteeContributionRequest{
					Slot:      1,
					PublicKey: pubKey[:],
					SubnetId:  0,
				},
			).Return(&ethpb.SyncCommitteeContribution{
				BlockRoot:       make([]byte, fieldparams.RootLength),
				Signature:       make([]byte, 96),
				AggregationBits: aggBits,
			}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			m.validatorClient.EXPECT().SubmitSignedContributionAndProof(
				gomock.Any(), // ctx
				gomock.AssignableToTypeOf(&ethpb.SignedContributionAndProof{
					Message: &ethpb.ContributionAndProof{
						AggregatorIndex: 7,
						Contribution: &ethpb.SyncCommitteeContribution{
							BlockRoot:         make([]byte, fieldparams.RootLength),
							Signature:         make([]byte, 96),
							AggregationBits:   bitfield.NewBitvector128(),
							Slot:              1,
							SubcommitteeIndex: 1,
						},
					},
				}),
			).Return(&emptypb.Empty{}, errors.New("Could not submit contribution"))

			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
			require.LogsContain(t, hook, "Could not submit signed contribution and proof")
		})
	}
}

func TestSubmitSignedContributionAndProof_Ok(t *testing.T) {
	// Hardcode secret key in order to have a valid aggregator signature.
	rawKey, err := hex.DecodeString("659e875e1b062c03f2f2a57332974d475b97df6cfc581d322e79642d39aca8fd")
	assert.NoError(t, err)
	validatorKey, err := bls.SecretKeyFromBytes(rawKey)
	assert.NoError(t, err)

	for _, isSlashingProtectionMinimal := range [...]bool{false, true} {
		t.Run(fmt.Sprintf("SlashingProtectionMinimal:%v", isSlashingProtectionMinimal), func(t *testing.T) {
			validator, m, validatorKey, finish := setupWithKey(t, validatorKey, isSlashingProtectionMinimal)
			validatorIndex := primitives.ValidatorIndex(7)
			committee := []primitives.ValidatorIndex{0, 3, 4, 2, validatorIndex, 6, 8, 9, 10}
			validator.duties = &ethpb.DutiesResponse{CurrentEpochDuties: []*ethpb.DutiesResponse_Duty{
				{
					PublicKey:      validatorKey.PublicKey().Marshal(),
					Committee:      committee,
					ValidatorIndex: validatorIndex,
				},
			}}
			defer finish()

			var pubKey [fieldparams.BLSPubkeyLength]byte
			copy(pubKey[:], validatorKey.PublicKey().Marshal())
			m.validatorClient.EXPECT().GetSyncSubcommitteeIndex(
				gomock.Any(), // ctx
				&ethpb.SyncSubcommitteeIndexRequest{
					Slot:      1,
					PublicKey: pubKey[:],
				},
			).Return(&ethpb.SyncSubcommitteeIndexResponse{Indices: []primitives.CommitteeIndex{1}}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			aggBits := bitfield.NewBitvector128()
			aggBits.SetBitAt(0, true)
			m.validatorClient.EXPECT().GetSyncCommitteeContribution(
				gomock.Any(), // ctx
				&ethpb.SyncCommitteeContributionRequest{
					Slot:      1,
					PublicKey: pubKey[:],
					SubnetId:  0,
				},
			).Return(&ethpb.SyncCommitteeContribution{
				BlockRoot:       make([]byte, fieldparams.RootLength),
				Signature:       make([]byte, 96),
				AggregationBits: aggBits,
			}, nil)

			m.validatorClient.EXPECT().
				DomainData(gomock.Any(), // ctx
					gomock.Any()). // epoch
				Return(&ethpb.DomainResponse{
					SignatureDomain: make([]byte, 32),
				}, nil)

			m.validatorClient.EXPECT().SubmitSignedContributionAndProof(
				gomock.Any(), // ctx
				gomock.AssignableToTypeOf(&ethpb.SignedContributionAndProof{
					Message: &ethpb.ContributionAndProof{
						AggregatorIndex: 7,
						Contribution: &ethpb.SyncCommitteeContribution{
							BlockRoot:         make([]byte, 32),
							Signature:         make([]byte, 96),
							AggregationBits:   bitfield.NewBitvector128(),
							Slot:              1,
							SubcommitteeIndex: 1,
						},
					},
				}),
			).Return(&emptypb.Empty{}, nil)

			validator.SubmitSignedContributionAndProof(context.Background(), 1, pubKey)
		})
	}
}