// Copyright 2021 ChainSafe Systems (ON)
// SPDX-License-Identifier: LGPL-3.0-only

package babe

import (
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/ChainSafe/gossamer/dot/state"
	"github.com/ChainSafe/gossamer/dot/types"
	"github.com/ChainSafe/gossamer/lib/babe/mocks"
	"github.com/ChainSafe/gossamer/lib/common"
	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
	"github.com/ChainSafe/gossamer/pkg/scale"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
)

const testSlotDuration = time.Second

// this is generated by printing key ownership proof while running `test_generate_equivocation_report_blob`
// https://github.com/paritytech/substrate/blob/ded44948e2d5a398abcb4e342b0513cb690961bb/frame/grandpa/src/benchmarking.rs#L85
var testKeyOwnershipProof types.OpaqueKeyOwnershipProof = types.OpaqueKeyOwnershipProof([]byte{64, 138, 252, 29, 127, 102, 189, 129, 207, 47, 157, 60, 17, 138, 194, 121, 139, 92, 176, 175, 224, 16, 185, 93, 175, 251, 224, 81, 209, 61, 0, 71}) //nolint:lll

func newTestHeader(t *testing.T, digest ...any) *types.Header {
	t.Helper()
	header := types.NewEmptyHeader()
	header.Number = 1
	for _, d := range digest {
		err := header.Digest.Add(d)
		assert.NoError(t, err)
	}

	return header
}

func signAndAddSeal(t *testing.T, kp *sr25519.Keypair, header *types.Header, data []byte) {
	t.Helper()
	sig, err := kp.Sign(data)
	require.NoError(t, err)

	err = header.Digest.Add(types.SealDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              sig,
	})
	assert.NoError(t, err)
}

func newEncodedBabeDigest(t *testing.T, value any) []byte {
	t.Helper()
	babeDigest := types.NewBabeDigest()
	err := babeDigest.SetValue(value)
	require.NoError(t, err)

	enc, err := scale.Marshal(babeDigest)
	require.NoError(t, err)
	return enc
}

func encodeAndHashHeader(t *testing.T, header *types.Header) common.Hash {
	t.Helper()
	encHeader, err := scale.Marshal(*header)
	require.NoError(t, err)

	hash, err := common.Blake2bHash(encHeader)
	require.NoError(t, err)
	return hash
}

func newTestVerifier(kp *sr25519.Keypair, blockState state.BlockState, slotState SlotState,
	threshold *scale.Uint128, secSlots bool) *verifier {
	authority := types.NewAuthority(kp.Public(), uint64(1))
	info := &verifierInfo{
		authorities:    []types.AuthorityRaw{*authority.ToRaw(), *authority.ToRaw()},
		randomness:     Randomness{},
		threshold:      threshold,
		secondarySlots: secSlots,
	}
	return newVerifier(blockState, slotState, 1, info, testSlotDuration)
}

func Test_getAuthorityIndex(t *testing.T) {
	digest := types.NewDigest()
	err := digest.Add(types.SealDigest{
		ConsensusEngineID: types.ConsensusEngineID{},
	})
	assert.NoError(t, err)
	headerNoPre := types.NewEmptyHeader()
	headerNoPre.Digest = digest

	digest2 := types.NewDigest()
	err = digest2.Add(types.PreRuntimeDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              []byte{1},
	})
	assert.NoError(t, err)
	headerInvalidPre := types.NewEmptyHeader()
	headerInvalidPre.Digest = digest2

	// BabePrimaryPreDigest Case
	babeDigest := types.NewBabeDigest()
	err = babeDigest.SetValue(types.BabePrimaryPreDigest{AuthorityIndex: 21, SlotNumber: 1})
	assert.NoError(t, err)

	bdEnc, err := scale.Marshal(babeDigest)
	require.NoError(t, err)

	digestPrimary := types.NewDigest()
	err = digestPrimary.Add(types.PreRuntimeDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              bdEnc,
	})
	assert.NoError(t, err)
	headerPrimary := types.NewEmptyHeader()
	headerPrimary.Digest = digestPrimary

	// BabeSecondaryVRFPreDigest Case
	babeDigest2 := types.NewBabeDigest()
	err = babeDigest2.SetValue(types.BabeSecondaryVRFPreDigest{AuthorityIndex: 21, SlotNumber: 10})
	assert.NoError(t, err)

	bdEnc2, err := scale.Marshal(babeDigest2)
	require.NoError(t, err)

	digestSecondary := types.NewDigest()
	err = digestSecondary.Add(types.PreRuntimeDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              bdEnc2,
	})
	assert.NoError(t, err)
	headerSecondary := types.NewEmptyHeader()
	headerSecondary.Digest = digestSecondary

	// BabeSecondaryPlainPreDigest case
	babeDigest3 := types.NewBabeDigest()
	err = babeDigest3.SetValue(types.BabeSecondaryPlainPreDigest{AuthorityIndex: 21, SlotNumber: 100})
	assert.NoError(t, err)

	bdEnc3, err := scale.Marshal(babeDigest3)
	require.NoError(t, err)

	digestSecondaryPlain := types.NewDigest()
	err = digestSecondaryPlain.Add(types.PreRuntimeDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              bdEnc3,
	})
	assert.NoError(t, err)
	headerSecondaryPlain := types.NewEmptyHeader()
	headerSecondaryPlain.Digest = digestSecondaryPlain

	type args struct {
		header *types.Header
	}
	tests := []struct {
		name          string
		args          args
		expAuthIdx    uint32
		expSlotNumber uint64
		expErr        error
	}{
		{
			name:   "No Digest",
			args:   args{types.NewEmptyHeader()},
			expErr: fmt.Errorf("for block hash %s: %w", types.NewEmptyHeader().Hash(), errNoDigest),
		},
		{
			name:   "First Digest Invalid Type",
			args:   args{headerNoPre},
			expErr: errors.New("first digest item is not pre-digest"),
		},
		{
			name: "Invalid_Preruntime_Digest_Type",
			args: args{headerInvalidPre},
			expErr: errors.New("cannot decode babe header from pre-digest: decoding struct: unmarshalling field at" +
				" index 0: EOF"),
		},
		{
			name:          "BabePrimaryPreDigest_Type",
			args:          args{headerPrimary},
			expAuthIdx:    21,
			expSlotNumber: 1,
		},
		{
			name:          "BabeSecondaryVRFPreDigest_Type",
			args:          args{headerSecondary},
			expAuthIdx:    21,
			expSlotNumber: 10,
		},
		{
			name:          "BabeSecondaryPlainPreDigest_Type",
			args:          args{headerSecondaryPlain},
			expAuthIdx:    21,
			expSlotNumber: 100,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			authIdx, slotNumber, err := getAuthorityIndexAndSlot(tt.args.header)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}
			assert.Equal(t, tt.expAuthIdx, authIdx)
			assert.Equal(t, tt.expSlotNumber, slotNumber)
		})
	}
}

func Test_verifier_verifyPrimarySlotWinner(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockBlockState := NewMockBlockState(ctrl)

	// Generate keys
	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	auth := types.NewAuthority(kp.Public(), uint64(1))
	vi := &verifierInfo{
		authorities: []types.AuthorityRaw{*auth.ToRaw()},
		threshold:   &scale.Uint128{},
	}

	vi1 := &verifierInfo{
		authorities: []types.AuthorityRaw{*auth.ToRaw()},
		threshold:   scale.MaxUint128,
	}

	mockSlotState := NewMockSlotState(nil)

	output, proof, err := kp.VrfSign(makeTranscript(Randomness{}, uint64(1), 1))
	assert.NoError(t, err)

	type args struct {
		authorityIndex uint32
		slot           uint64
		vrfOutput      [sr25519.VRFOutputLength]byte
		vrfProof       [sr25519.VRFProofLength]byte
	}
	tests := []struct {
		name     string
		verifier verifier
		args     args
		exp      bool
		expErr   error
	}{
		{
			name:     "Over threshold",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, vi, time.Second),
			args: args{
				slot:      1,
				vrfOutput: [32]byte{},
				vrfProof:  [64]byte{},
			},
			expErr: ErrVRFOutputOverThreshold,
		},
		{
			name:     "VRF not verified",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, vi1, time.Second),
			args: args{
				slot:      1,
				vrfOutput: [32]byte{},
				vrfProof:  [64]byte{},
			},
		},
		{
			name:     "VRF verified",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, vi1, time.Second),
			args: args{
				slot:      1,
				vrfOutput: output,
				vrfProof:  proof,
			},
			exp: true,
		},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			res, err := tt.verifier.verifyPrimarySlotWinner(tt.args.authorityIndex, tt.args.slot,
				tt.args.vrfOutput, tt.args.vrfProof)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}
			assert.Equal(t, tt.exp, res)
		})
	}
}

func Test_verifier_verifyPreRuntimeDigest(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockBlockState := NewMockBlockState(ctrl)

	// Generate keys
	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	// BabePrimaryPreDigest case
	output, proof, err := kp.VrfSign(makeTranscript(Randomness{}, uint64(1), 1))
	assert.NoError(t, err)

	secDigest1 := types.BabePrimaryPreDigest{
		SlotNumber: 1,
		VRFOutput:  output,
		VRFProof:   proof,
	}
	prd1, err := secDigest1.ToPreRuntimeDigest()
	assert.NoError(t, err)

	auth := types.NewAuthority(kp.Public(), uint64(1))
	vi := &verifierInfo{
		authorities: []types.AuthorityRaw{*auth.ToRaw(), *auth.ToRaw()},
		threshold:   scale.MaxUint128,
	}

	mockSlotState := NewMockSlotState(nil)

	v := newVerifier(mockBlockState, mockSlotState, 1, vi, testSlotDuration)

	// Invalid
	v2 := newVerifier(mockBlockState, mockSlotState, 13, vi, testSlotDuration)

	// Above threshold case
	vi1 := &verifierInfo{
		authorities: []types.AuthorityRaw{*auth.ToRaw(), *auth.ToRaw()},
		threshold:   &scale.Uint128{},
	}

	v1 := newVerifier(mockBlockState, mockSlotState, 1, vi1, testSlotDuration)

	// BabeSecondaryVRFPreDigest case
	secVRFDigest := types.BabeSecondaryVRFPreDigest{
		SlotNumber: 1,
		VrfOutput:  output,
		VrfProof:   proof,
	}

	digestSecondaryVRF := types.NewBabeDigest()
	err = digestSecondaryVRF.SetValue(secVRFDigest)
	assert.NoError(t, err)

	bdEnc, err := scale.Marshal(digestSecondaryVRF)
	require.NoError(t, err)

	babePRD := types.NewBABEPreRuntimeDigest(bdEnc)

	authVRFSec := types.NewAuthority(kp.Public(), uint64(1))
	viVRFSec := &verifierInfo{
		authorities: []types.AuthorityRaw{*authVRFSec.ToRaw(), *authVRFSec.ToRaw()},
		threshold:   scale.MaxUint128,
	}

	viVRFSec2 := &verifierInfo{
		authorities:    []types.AuthorityRaw{*authVRFSec.ToRaw(), *authVRFSec.ToRaw()},
		threshold:      scale.MaxUint128,
		secondarySlots: true,
	}

	// BabeSecondaryPlainPreDigest case
	secDigest := types.BabeSecondaryPlainPreDigest{AuthorityIndex: 0, SlotNumber: uint64(1)}
	prd, err := secDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)

	authSec := types.NewAuthority(kp.Public(), uint64(1))
	viSec := &verifierInfo{
		authorities: []types.AuthorityRaw{*authSec.ToRaw(), *authSec.ToRaw()},
		threshold:   scale.MaxUint128,
	}

	viSec2 := &verifierInfo{
		authorities:    []types.AuthorityRaw{*authSec.ToRaw(), *authSec.ToRaw()},
		threshold:      scale.MaxUint128,
		secondarySlots: true,
	}

	type args struct {
		digest *types.PreRuntimeDigest
	}
	tests := []struct {
		name     string
		verifier verifier
		args     args
		exp      any
		expErr   error
	}{
		{
			name:     "Invalid PreRuntimeDigest",
			verifier: verifier{},
			args:     args{&types.PreRuntimeDigest{Data: []byte{0}}},
			expErr: errors.New(
				"unable to find VaryingDataTypeValue with index: for key 0 unable to find VaryingDataTypeValue with index"),
		},
		{
			name:     "Invalid BlockProducer Index",
			verifier: verifier{},
			args:     args{prd},
			expErr:   ErrInvalidBlockProducerIndex,
		},
		{
			name:     "BabePrimaryPreDigest Case OK",
			verifier: *v,
			args:     args{prd1},
			exp: types.BabePrimaryPreDigest{
				SlotNumber: 1,
				VRFOutput:  output,
				VRFProof:   proof,
			},
		},
		{
			name:     "BabePrimaryPreDigest Case output over threshold",
			verifier: *v1,
			args:     args{prd1},
			expErr:   errors.New("vrf output over threshold"),
		},
		{
			name:     "BabePrimaryPreDigest Case Invalid",
			verifier: *v2,
			args:     args{prd1},
			expErr:   ErrBadSlotClaim,
		},
		{
			name:     "BabeSecondaryPlainPreDigest SecondarySlot false",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, viSec, testSlotDuration),
			args:     args{prd},
			expErr:   ErrBadSlotClaim,
		},
		{
			name:     "BabeSecondaryPlainPreDigest invalid claim",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, viSec2, testSlotDuration),
			args:     args{prd},
			expErr:   errors.New("invalid secondary slot claim"),
		},
		{
			name:     "BabeSecondaryVRFPreDigest SecondarySlot false",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, viVRFSec, testSlotDuration),
			args:     args{babePRD},
			expErr:   ErrBadSlotClaim,
		},
		{
			name:     "BabeSecondaryVRFPreDigest invalid claim",
			verifier: *newVerifier(mockBlockState, mockSlotState, 1, viVRFSec2, testSlotDuration),
			args:     args{babePRD},
			expErr:   errors.New("invalid secondary slot claim"),
		},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			res, err := tt.verifier.verifyPreRuntimeDigest(tt.args.digest)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}
			assert.Equal(t, tt.exp, res)
		})
	}
}

func Test_verifier_verifyAuthorshipRight(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockBlockState := NewMockBlockState(ctrl)

	// Generate keys
	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	// Create a VRF output and proof
	output, proof, err := kp.VrfSign(makeTranscript(Randomness{}, uint64(1), 1))
	assert.NoError(t, err)

	testBabePrimaryPreDigest := types.BabePrimaryPreDigest{
		SlotNumber: 1,
		VRFOutput:  output,
		VRFProof:   proof,
	}
	testBabeSecondaryPlainPreDigest := types.BabeSecondaryPlainPreDigest{
		AuthorityIndex: 1,
		SlotNumber:     1,
	}
	testBabeSecondaryVRFPreDigest := types.BabeSecondaryVRFPreDigest{
		AuthorityIndex: 1,
		SlotNumber:     1,
		VrfOutput:      output,
		VrfProof:       proof,
	}
	testInvalidSeal := types.SealDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              []byte{1},
	}
	testInvalidPreRuntimeDigest := types.PreRuntimeDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              []byte{1},
	}

	// Primary Test Header
	encTestDigest := newEncodedBabeDigest(t, types.BabePrimaryPreDigest{AuthorityIndex: 0})
	testDigestPrimary := types.NewDigest()
	err = testDigestPrimary.Add(types.PreRuntimeDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              encTestDigest,
	})
	assert.NoError(t, err)

	testHeaderPrimary := types.NewEmptyHeader()
	testHeaderPrimary.Digest = testDigestPrimary

	// Secondary Plain Test Header
	testParentPrd, err := testBabeSecondaryPlainPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)

	testParentHeader := newTestHeader(t, *testParentPrd)
	testParentHash := encodeAndHashHeader(t, testParentHeader)
	testSecondaryPrd, err := testBabeSecondaryPlainPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)

	testSecPlainHeader := newTestHeader(t, *testSecondaryPrd)
	testSecPlainHeader.ParentHash = testParentHash

	// Secondary Vrf Test Header
	encParentVrfDigest := newEncodedBabeDigest(t, testBabeSecondaryVRFPreDigest)
	testParentVrfHeader := newTestHeader(t, *types.NewBABEPreRuntimeDigest(encParentVrfDigest))

	testVrfParentHash := encodeAndHashHeader(t, testParentVrfHeader)
	encVrfHeader := newEncodedBabeDigest(t, testBabeSecondaryVRFPreDigest)
	testSecVrfHeader := newTestHeader(t, *types.NewBABEPreRuntimeDigest(encVrfHeader))
	testSecVrfHeader.ParentHash = testVrfParentHash

	mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))

	mockSlotStateNoOp := NewMockSlotState(nil)

	// Case 0: First element not preruntime digest
	header0 := newTestHeader(t, testInvalidSeal, testInvalidSeal)

	// Case 1: Last element not seal
	header1 := newTestHeader(t, testInvalidPreRuntimeDigest, testInvalidPreRuntimeDigest)

	// Case 2: Fail to verify preruntime digest
	header2 := newTestHeader(t, testInvalidPreRuntimeDigest, testInvalidSeal)

	// Case 3: Invalid Seal Length
	babePrd, err := testBabePrimaryPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)
	header3 := newTestHeader(t, *babePrd, testInvalidSeal)
	babeVerifier := newTestVerifier(kp, mockBlockState, mockSlotStateNoOp, scale.MaxUint128, false)

	// Case 4: Invalid signature - BabePrimaryPreDigest
	babePrd2, err := testBabePrimaryPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)
	header4 := newTestHeader(t, *babePrd2)
	signAndAddSeal(t, kp, header4, []byte{1})
	babeVerifier2 := newTestVerifier(kp, mockBlockState, mockSlotStateNoOp, scale.MaxUint128, false)

	// Case 5: Invalid signature - BabeSecondaryPlainPreDigest
	babeSecPlainPrd, err := testBabeSecondaryPlainPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)
	header5 := newTestHeader(t, *babeSecPlainPrd)
	signAndAddSeal(t, kp, header5, []byte{1})
	babeVerifier3 := newTestVerifier(kp, mockBlockState, mockSlotStateNoOp, scale.MaxUint128, true)

	// Case 6: Invalid signature - BabeSecondaryVrfPreDigest
	encSecVrfDigest := newEncodedBabeDigest(t, testBabeSecondaryVRFPreDigest)
	assert.NoError(t, err)
	header6 := newTestHeader(t, *types.NewBABEPreRuntimeDigest(encSecVrfDigest))
	signAndAddSeal(t, kp, header6, []byte{1})
	babeVerifier4 := newTestVerifier(kp, mockBlockState, mockSlotStateNoOp, scale.MaxUint128, true)

	// Case 7: GetAuthorityIndex Err
	babeParentPrd, err := testBabePrimaryPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)
	babeParentHeader := newTestHeader(t, *babeParentPrd)

	parentHash := encodeAndHashHeader(t, babeParentHeader)
	babePrd3, err := testBabePrimaryPreDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)

	header7 := newTestHeader(t, *babePrd3)
	header7.ParentHash = parentHash

	hash := encodeAndHashHeader(t, header7)
	signAndAddSeal(t, kp, header7, hash[:])

	signerPublicKey := kp.Public().(*sr25519.PublicKey)
	signerAuthID := types.AuthorityID(signerPublicKey.AsBytes())

	slotStateMockErr := errors.New("slot state mock error")
	mockSlotState := NewMockSlotState(ctrl)
	mockSlotState.EXPECT().CheckEquivocation(gomock.Any(),
		testBabePrimaryPreDigest.SlotNumber, header7, signerAuthID).Return(nil, slotStateMockErr)

	babeVerifier5 := newTestVerifier(kp, mockBlockState, mockSlotState, scale.MaxUint128, false)

	tests := []struct {
		name     string
		verifier verifier
		header   *types.Header
		expErr   error
	}{
		{
			name:     "missing digest",
			verifier: verifier{},
			header:   types.NewEmptyHeader(),
			expErr:   errMissingDigestItems,
		},
		{
			name:     "first digest invalid",
			verifier: verifier{},
			header:   header0,
			expErr:   fmt.Errorf("%w: got types.SealDigest", types.ErrNoFirstPreDigest),
		},
		{
			name:     "last digest invalid",
			verifier: verifier{},
			header:   header1,
			expErr:   fmt.Errorf("%w: got types.PreRuntimeDigest", errLastDigestItemNotSeal),
		},
		{
			name:     "invalid preruntime digest data",
			verifier: verifier{},
			header:   header2,
			expErr: errors.New("failed to verify pre-runtime digest: decoding struct: unmarshalling field at index" +
				" 0: EOF"),
		},
		{
			name:     "invalid seal length",
			verifier: *babeVerifier,
			header:   header3,
			expErr:   errors.New("invalid signature length"),
		},
		{
			name:     "invalid seal signature - primary",
			verifier: *babeVerifier2,
			header:   header4,
			expErr:   ErrBadSignature,
		},
		{
			name:     "invalid seal signature - secondary plain",
			verifier: *babeVerifier3,
			header:   header5,
			expErr:   ErrBadSignature,
		},
		{
			name:     "invalid seal signature - secondary vrf",
			verifier: *babeVerifier4,
			header:   header6,
			expErr:   ErrBadSignature,
		},
		{
			name:     "check_equivocation_error",
			verifier: *babeVerifier5,
			header:   header7,
			expErr: fmt.Errorf("could not verify block equivocation: "+
				"checking equivocation: %w", slotStateMockErr),
		},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			err := tt.verifier.verifyAuthorshipRight(tt.header)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}

		})
	}
}

func Test_verifyBlockEquivocation(t *testing.T) {
	t.Parallel()
	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	slotStateMockErr := errors.New("slot state error")
	getRuntimeErr := errors.New("mock get runtime error")

	output, proof, err := kp.VrfSign(makeTranscript(Randomness{}, uint64(1), 1))
	assert.NoError(t, err)

	babePrimaryDigest := types.BabePrimaryPreDigest{
		AuthorityIndex: 0,
		SlotNumber:     1,
		VRFOutput:      output,
		VRFProof:       proof,
	}
	prd, err := babePrimaryDigest.ToPreRuntimeDigest()
	assert.NoError(t, err)
	defaultHeader := newTestHeader(t, *prd)

	cases := map[string]struct {
		header        *types.Header
		wantErr       error
		errString     string
		expected      bool
		buildVerifier func(t *testing.T) *verifier
	}{
		"empty_digest_failed_to_get_auth_index": {
			header: types.NewEmptyHeader(),
			errString: "failed to get authority index: for block hash " +
				"0xdcdd89927d8a348e00257e1ecc8617f45edb5118efff3ea2f9961b2ad9b7690a: no digest provided",
			wantErr: errNoDigest,
			buildVerifier: func(t *testing.T) *verifier {
				return &verifier{}
			},
		},
		"cannot_get_slot_from_header": {
			header: func() *types.Header {
				wrongDigest := types.NewGrandpaConsensusDigest()
				require.NoError(t, wrongDigest.SetValue(types.GrandpaForcedChange{}))

				data, err := scale.Marshal(wrongDigest)
				require.NoError(t, err)

				consensusDigest := types.ConsensusDigest{
					ConsensusEngineID: types.GrandpaEngineID,
					Data:              data,
				}

				return newTestHeader(t, consensusDigest)
			}(),
			errString: "failed to get authority index: first digest item is not pre-digest",
			wantErr:   types.ErrNoFirstPreDigest,
			buildVerifier: func(t *testing.T) *verifier {
				return &verifier{
					authorities: []types.AuthorityRaw{},
				}
			},
		},
		"auth_index_out_of_bound": {
			header:    defaultHeader,
			errString: "authority index doesn't exist",
			wantErr:   ErrAuthIndexOutOfBound,
			buildVerifier: func(t *testing.T) *verifier {
				return &verifier{
					authorities: []types.AuthorityRaw{},
				}
			},
		},
		"failed_to_check_equivocation": {
			header:    defaultHeader,
			errString: "checking equivocation: slot state error",
			wantErr:   slotStateMockErr,
			buildVerifier: func(t *testing.T) *verifier {
				ctrl := gomock.NewController(t)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))

				expectedAuthorityId := types.AuthorityID(kp.Public().Encode())
				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.
					EXPECT().
					CheckEquivocation(gomock.Any(), uint64(1),
						defaultHeader, expectedAuthorityId).
					Return(nil, slotStateMockErr)

				return &verifier{
					authorities: []types.AuthorityRaw{
						{
							Key:    [32]byte(kp.Public().Encode()),
							Weight: 1,
						},
					},
					blockState:   mockBlockState,
					slotState:    mockSlotState,
					slotDuration: 6 * time.Second,
				}
			},
		},
		"nil_equivocation_proof_from_check_equivocation": {
			header: defaultHeader,
			buildVerifier: func(t *testing.T) *verifier {
				ctrl := gomock.NewController(t)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))

				expectedAuthorityId := types.AuthorityID(kp.Public().Encode())
				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.
					EXPECT().
					CheckEquivocation(gomock.Any(), uint64(1),
						defaultHeader, expectedAuthorityId).
					Return(nil, nil)

				return &verifier{
					authorities: []types.AuthorityRaw{
						{
							Key:    [32]byte(kp.Public().Encode()),
							Weight: 1,
						},
					},
					blockState:   mockBlockState,
					slotState:    mockSlotState,
					slotDuration: 6 * time.Second,
				}
			},
		},
		"got_equivocation_proof": {
			expected: true,
			header:   defaultHeader,
			buildVerifier: func(t *testing.T) *verifier {
				ctrl := gomock.NewController(t)

				secondHeader := types.NewEmptyHeader()
				secondHeader.Number = 1
				secondHeader.Hash()

				expectedAuthorityId := types.AuthorityID(kp.Public().Encode())

				mockedEquivocationProof := &types.BabeEquivocationProof{
					Offender:     expectedAuthorityId,
					Slot:         1,
					FirstHeader:  *defaultHeader,
					SecondHeader: *secondHeader,
				}

				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.
					EXPECT().
					CheckEquivocation(gomock.Any(), uint64(1),
						defaultHeader, expectedAuthorityId).
					Return(mockedEquivocationProof, nil)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))
				mockBlockState.EXPECT().BestBlockHash().Return(defaultHeader.Hash())

				opaqueOwnershipProof := []byte{0, 1, 2, 3, 4, 5} // random data

				mockRuntimeInstance := mocks.NewMockInstance(ctrl)
				mockRuntimeInstance.EXPECT().
					BabeGenerateKeyOwnershipProof(uint64(1), expectedAuthorityId).
					Return(types.OpaqueKeyOwnershipProof(opaqueOwnershipProof), nil)

				mockRuntimeInstance.EXPECT().
					BabeSubmitReportEquivocationUnsignedExtrinsic(*mockedEquivocationProof, opaqueOwnershipProof).
					Return(nil)

				mockBlockState.EXPECT().GetRuntime(defaultHeader.Hash()).Return(mockRuntimeInstance, nil)

				return &verifier{
					authorities: []types.AuthorityRaw{
						{
							Key:    [32]byte(kp.Public().Encode()),
							Weight: 1,
						},
					},
					blockState:   mockBlockState,
					slotState:    mockSlotState,
					slotDuration: 6 * time.Second,
				}
			},
		},
		"failed_to_get_runtime_while_submiting_equivocation": {
			header:    defaultHeader,
			wantErr:   getRuntimeErr,
			errString: "submiting equivocation: getting runtime: mock get runtime error",
			buildVerifier: func(t *testing.T) *verifier {
				ctrl := gomock.NewController(t)

				secondHeader := types.NewEmptyHeader()
				secondHeader.Number = 1
				secondHeader.Hash()

				expectedAuthorityId := types.AuthorityID(kp.Public().Encode())

				mockedEquivocationProof := &types.BabeEquivocationProof{
					Offender:     expectedAuthorityId,
					Slot:         1,
					FirstHeader:  *defaultHeader,
					SecondHeader: *secondHeader,
				}

				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.
					EXPECT().
					CheckEquivocation(gomock.Any(), uint64(1),
						defaultHeader, expectedAuthorityId).
					Return(mockedEquivocationProof, nil)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))
				mockBlockState.EXPECT().BestBlockHash().Return(defaultHeader.Hash())
				mockBlockState.EXPECT().GetRuntime(defaultHeader.Hash()).Return(nil, getRuntimeErr)

				return &verifier{
					authorities: []types.AuthorityRaw{
						{
							Key:    [32]byte(kp.Public().Encode()),
							Weight: 1,
						},
					},
					blockState:   mockBlockState,
					slotState:    mockSlotState,
					slotDuration: 6 * time.Second,
				}
			},
		},
	}

	for tname, tt := range cases {
		tt := tt

		t.Run(tname, func(t *testing.T) {
			t.Parallel()
			verifier := tt.buildVerifier(t)
			out, err := verifier.verifyBlockEquivocation(tt.header)
			require.ErrorIs(t, err, tt.wantErr)
			if tt.errString != "" {
				require.EqualError(t, err, tt.errString)
			}

			require.Equal(t, out, tt.expected)
		})
	}
}

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

	keyPair, err := sr25519.GenerateKeypair()
	require.NoError(t, err)

	auth := types.NewAuthority(keyPair.Public(), uint64(1))
	vi := &verifierInfo{
		authorities: []types.AuthorityRaw{*auth.ToRaw(), *auth.ToRaw()},
		threshold:   scale.MaxUint128,
	}

	ctrl := gomock.NewController(t)
	mockBlockState := NewMockBlockState(ctrl)
	mockSlotState := NewMockSlotState(nil)

	verifier := newVerifier(mockBlockState, mockSlotState, 1, vi, testSlotDuration)

	const slot = uint64(1)
	const authorityIndex = uint32(1)
	const epochNumber = uint64(2)
	output, proof, err := keyPair.VrfSign(makeTranscript(Randomness{}, slot, epochNumber))
	require.NoError(t, err)

	testDigest := types.BabePrimaryPreDigest{
		AuthorityIndex: authorityIndex,
		SlotNumber:     slot,
		VRFOutput:      output,
		VRFProof:       proof,
	}
	preRuntimeDigest, err := testDigest.ToPreRuntimeDigest()
	require.NoError(t, err)

	firstHeader := newTestHeader(t, *preRuntimeDigest)
	firstHash := encodeAndHashHeader(t, firstHeader)
	signAndAddSeal(t, keyPair, firstHeader, firstHash[:])

	output2, proof2, err := keyPair.VrfSign(makeTranscript(Randomness{}, slot, epochNumber))
	require.NoError(t, err)

	digest2 := types.BabePrimaryPreDigest{
		AuthorityIndex: authorityIndex,
		SlotNumber:     slot,
		VRFOutput:      output2,
		VRFProof:       proof2,
	}
	preRuntimeDigest2, err := digest2.ToPreRuntimeDigest()
	require.NoError(t, err)

	secondHeader := newTestHeader(t, *preRuntimeDigest2)

	offenderPublicKey := verifier.authorities[authorityIndex].Key
	keyOwnershipProof := testKeyOwnershipProof
	mockRuntime := mocks.NewMockInstance(ctrl)

	equivocationProof := types.BabeEquivocationProof{
		Offender:     offenderPublicKey,
		Slot:         slot,
		FirstHeader:  *firstHeader,
		SecondHeader: *secondHeader,
	}
	mockRuntime.EXPECT().BabeGenerateKeyOwnershipProof(slot, offenderPublicKey).Return(keyOwnershipProof, nil)
	mockRuntime.EXPECT().BabeSubmitReportEquivocationUnsignedExtrinsic(equivocationProof, keyOwnershipProof).Return(nil)

	mockBlockState.EXPECT().BestBlockHash().Return(firstHash).Times(2)
	mockBlockState.EXPECT().GetRuntime(firstHash).Return(mockRuntime, nil)

	err = verifier.submitAndReportEquivocation(&equivocationProof)
	assert.NoError(t, err)

	// fails on not being able to get a runtime
	mockBlockState.EXPECT().GetRuntime(firstHash).Return(nil, errors.New("test error"))

	err = verifier.submitAndReportEquivocation(&equivocationProof)
	assert.EqualError(t, err, "getting runtime: test error")

}

func Test_verifier_verifyAuthorshipRightEquivocatory(t *testing.T) {
	ctrl := gomock.NewController(t)

	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	tests := []struct {
		name          string
		setupVerifier func(t *testing.T, header *types.Header) *verifier
		setupHeader   func(t *testing.T) *types.Header
		expErr        func(*types.Header) error
	}{
		{
			name: "equivocate_primary",
			setupHeader: func(t *testing.T) *types.Header {
				output, proof, err := kp.VrfSign(makeTranscript(Randomness{}, uint64(1), 1))
				assert.NoError(t, err)

				// BabePrimaryPreDigest case
				primaryDigest := types.BabePrimaryPreDigest{
					AuthorityIndex: 1,
					SlotNumber:     1,
					VRFOutput:      output,
					VRFProof:       proof,
				}
				primaryDigestRuntime, err := primaryDigest.ToPreRuntimeDigest()
				assert.NoError(t, err)

				header := newTestHeader(t, *primaryDigestRuntime)
				header.Hash()
				headerHash := encodeAndHashHeader(t, header)
				signAndAddSeal(t, kp, header, headerHash[:])

				return header
			},
			setupVerifier: func(t *testing.T, header *types.Header) *verifier {
				const slot = uint64(1)
				offenderPublicKey := types.AuthorityID(kp.Public().Encode())

				equivocationProof := &types.BabeEquivocationProof{
					Offender:     offenderPublicKey,
					Slot:         slot,
					FirstHeader:  *header,
					SecondHeader: *types.NewEmptyHeader(),
				}

				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.EXPECT().
					CheckEquivocation(gomock.Any(), slot, header, offenderPublicKey).
					Return(equivocationProof, nil)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))
				mockBlockState.EXPECT().BestBlockHash().Return(header.Hash())

				mockRuntime := mocks.NewMockInstance(gomock.NewController(t))
				mockRuntime.EXPECT().
					BabeGenerateKeyOwnershipProof(slot, offenderPublicKey).
					Return(testKeyOwnershipProof, nil)

				mockRuntime.EXPECT().
					BabeSubmitReportEquivocationUnsignedExtrinsic(
						*equivocationProof, testKeyOwnershipProof).
					Return(nil)

				mockBlockState.EXPECT().GetRuntime(header.Hash()).Return(mockRuntime, nil)
				auth := types.NewAuthority(kp.Public(), uint64(1))
				info := &verifierInfo{
					authorities: []types.AuthorityRaw{*auth.ToRaw(), *auth.ToRaw()},
					threshold:   scale.MaxUint128,
				}

				return newVerifier(mockBlockState, mockSlotState, 1, info, testSlotDuration)
			},
			expErr: func(h *types.Header) error {
				return fmt.Errorf("%w for block header %s", ErrProducerEquivocated, h.Hash())
			},
		},
		{
			name: "equivocate_secondary_plain",
			setupHeader: func(t *testing.T) *types.Header {
				babeSecondaryPlainPreDigest := types.BabeSecondaryPlainPreDigest{
					AuthorityIndex: 1,
					SlotNumber:     1,
				}

				babeSecPlainPrd2, err := babeSecondaryPlainPreDigest.ToPreRuntimeDigest()
				assert.NoError(t, err)
				header := newTestHeader(t, *babeSecPlainPrd2)
				header.Hash()

				hashEquivocatorySecondaryPlain := encodeAndHashHeader(t, header)
				signAndAddSeal(t, kp, header, hashEquivocatorySecondaryPlain[:])

				return header
			},
			setupVerifier: func(t *testing.T, header *types.Header) *verifier {
				const slot = uint64(1)
				offenderPublicKey := types.AuthorityID(kp.Public().Encode())

				equivocationProof := &types.BabeEquivocationProof{
					Offender:     offenderPublicKey,
					Slot:         slot,
					FirstHeader:  *header,
					SecondHeader: *types.NewEmptyHeader(),
				}

				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.EXPECT().
					CheckEquivocation(gomock.Any(), slot, header, offenderPublicKey).
					Return(equivocationProof, nil)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))
				mockBlockState.EXPECT().BestBlockHash().Return(header.Hash())

				mockRuntime := mocks.NewMockInstance(gomock.NewController(t))
				mockRuntime.EXPECT().
					BabeGenerateKeyOwnershipProof(slot, offenderPublicKey).
					Return(testKeyOwnershipProof, nil)

				mockRuntime.EXPECT().
					BabeSubmitReportEquivocationUnsignedExtrinsic(
						*equivocationProof, testKeyOwnershipProof).
					Return(nil)

				mockBlockState.EXPECT().GetRuntime(header.Hash()).Return(mockRuntime, nil)
				auth := types.NewAuthority(kp.Public(), uint64(1))
				info := &verifierInfo{
					authorities:    []types.AuthorityRaw{*auth.ToRaw(), *auth.ToRaw()},
					threshold:      scale.MaxUint128,
					secondarySlots: true,
					randomness:     Randomness{},
				}

				return newVerifier(mockBlockState, mockSlotState, 1, info, testSlotDuration)
			},
			expErr: func(h *types.Header) error {
				return fmt.Errorf("%w for block header %s", ErrProducerEquivocated, h.Hash())
			},
		},
		{
			name: "equivocate_secondary_vrf",
			expErr: func(h *types.Header) error {
				return fmt.Errorf("%w for block header %s", ErrProducerEquivocated, h.Hash())
			},
			setupHeader: func(t *testing.T) *types.Header {
				output, proof, err := kp.VrfSign(makeTranscript(Randomness{}, uint64(1), 1))
				assert.NoError(t, err)

				babeSecondaryVRFPreDigest := types.BabeSecondaryVRFPreDigest{
					AuthorityIndex: 1,
					SlotNumber:     1,
					VrfOutput:      output,
					VrfProof:       proof,
				}

				vrfDigest := newEncodedBabeDigest(t, babeSecondaryVRFPreDigest)
				header := newTestHeader(t, *types.NewBABEPreRuntimeDigest(vrfDigest))
				header.Hash()

				hashEquivocatorySecondaryVRF := encodeAndHashHeader(t, header)
				signAndAddSeal(t, kp, header, hashEquivocatorySecondaryVRF[:])

				return header
			},
			setupVerifier: func(t *testing.T, header *types.Header) *verifier {
				const slot = uint64(1)
				offenderPublicKey := types.AuthorityID(kp.Public().Encode())

				equivocationProof := &types.BabeEquivocationProof{
					Offender:     offenderPublicKey,
					Slot:         slot,
					FirstHeader:  *header,
					SecondHeader: *types.NewEmptyHeader(),
				}

				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.EXPECT().
					CheckEquivocation(gomock.Any(), slot, header, offenderPublicKey).
					Return(equivocationProof, nil)

				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.EXPECT().GenesisHash().Return(common.Hash([32]byte{}))
				mockBlockState.EXPECT().BestBlockHash().Return(header.Hash())

				mockRuntime := mocks.NewMockInstance(gomock.NewController(t))
				mockRuntime.EXPECT().
					BabeGenerateKeyOwnershipProof(slot, offenderPublicKey).
					Return(testKeyOwnershipProof, nil)

				mockRuntime.EXPECT().
					BabeSubmitReportEquivocationUnsignedExtrinsic(
						*equivocationProof, testKeyOwnershipProof).
					Return(nil)

				mockBlockState.EXPECT().GetRuntime(header.Hash()).Return(mockRuntime, nil)

				auth := types.NewAuthority(kp.Public(), uint64(1))
				info := &verifierInfo{
					authorities:    []types.AuthorityRaw{*auth.ToRaw(), *auth.ToRaw()},
					threshold:      scale.MaxUint128,
					secondarySlots: true,
					randomness:     Randomness{},
				}

				return newVerifier(mockBlockState, mockSlotState, 1, info, testSlotDuration)
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			argHeader := tt.setupHeader(t)
			verifier := tt.setupVerifier(t, argHeader)
			err := verifier.verifyAuthorshipRight(argHeader)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr(argHeader).Error())
				return
			}
			assert.NoError(t, err)
		})
	}
}

func TestVerificationManager_getVerifierInfo(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockEpochStateGetErr := NewMockEpochState(ctrl)
	mockEpochStateHasErr := NewMockEpochState(ctrl)
	mockEpochStateThresholdErr := NewMockEpochState(ctrl)
	mockEpochStateOk := NewMockEpochState(ctrl)

	testHeader := types.NewEmptyHeader()

	mockEpochStateGetErr.EXPECT().GetEpochDataRaw(uint64(0), testHeader).Return(nil, state.ErrEpochNotInMemory)

	mockEpochStateHasErr.EXPECT().GetEpochDataRaw(uint64(0), testHeader).Return(&types.EpochDataRaw{}, nil)
	mockEpochStateHasErr.EXPECT().GetConfigData(uint64(0), testHeader).Return(&types.ConfigData{}, state.ErrConfigNotFound)

	mockEpochStateThresholdErr.EXPECT().GetEpochDataRaw(uint64(0), testHeader).Return(&types.EpochDataRaw{}, nil)
	mockEpochStateThresholdErr.EXPECT().GetConfigData(uint64(0), testHeader).
		Return(&types.ConfigData{
			C1: 3,
			C2: 1,
		}, nil)

	mockEpochStateOk.EXPECT().GetEpochDataRaw(uint64(0), testHeader).Return(&types.EpochDataRaw{}, nil)
	mockEpochStateOk.EXPECT().GetConfigData(uint64(0), testHeader).
		Return(&types.ConfigData{
			C1: 1,
			C2: 3,
		}, nil)

	vm0 := &VerificationManager{epochState: mockEpochStateGetErr}
	vm1 := &VerificationManager{epochState: mockEpochStateHasErr}
	vm2 := &VerificationManager{epochState: mockEpochStateThresholdErr}
	vm3 := &VerificationManager{epochState: mockEpochStateOk}

	tests := []struct {
		name   string
		vm     *VerificationManager
		epoch  uint64
		exp    *verifierInfo
		expErr error
	}{
		{
			name:   "getEpochData error",
			vm:     vm0,
			expErr: fmt.Errorf("failed to get epoch data for epoch %d: %w", 0, state.ErrEpochNotInMemory),
		},
		{
			name:   "getConfigData error",
			vm:     vm1,
			expErr: fmt.Errorf("failed to get config data: %w", state.ErrConfigNotFound),
		},
		{
			name:   "calculate threshold error",
			vm:     vm2,
			expErr: errors.New("failed to calculate threshold: invalid C1/C2: greater than 1"),
		},
		{
			name: "happy_path",
			vm:   vm3,
			exp: &verifierInfo{
				threshold: scale.MaxUint128,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := tt.vm
			res, err := v.getVerifierInfo(tt.epoch, testHeader)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}
			assert.Equal(t, tt.exp, res)
		})
	}
}

func TestVerificationManager_VerifyBlock(t *testing.T) {
	// Generate keys
	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	defaultParentHeader := types.NewEmptyHeader()
	const defaultEpoch = 1
	const defaultSlotNumber = 0

	defaultEpoch1Authorities := []types.AuthorityRaw{
		{
			Key:    [32]byte(kp.Public().Encode()),
			Weight: 1,
		},
	}

	defaultConfigData := &types.ConfigData{
		C1:             1,
		C2:             1,
		SecondarySlots: 1,
	}

	threshold, err := CalculateThreshold(defaultConfigData.C1, defaultConfigData.C2, len(defaultEpoch1Authorities))
	require.NoError(t, err)

	defaultEpochData := &epochData{
		randomness:     [32]byte{},
		authorityIndex: 0,
		authorities:    defaultEpoch1Authorities,
		threshold:      threshold,
	}

	defaultPreRuntimeDigestForEpoch1, err := claimSlot(defaultEpoch, defaultSlotNumber, defaultEpochData, kp)
	require.NoError(t, err)

	digest := types.NewDigest()
	err = digest.Add(*defaultPreRuntimeDigestForEpoch1)
	require.NoError(t, err)

	headerWithPreRuntimeDigest := types.NewHeader(defaultParentHeader.Hash(),
		common.Hash{}, common.Hash{}, defaultParentHeader.Number+1, digest)

	seal := buildSealDigest(t, headerWithPreRuntimeDigest, kp)
	headerWithPreRuntimeDigest.Digest.Add(*seal)

	testBlockHeaderEmpty := types.NewEmptyHeader()
	testBlockHeaderEmpty.Number = 2

	errTestGetEpoch := errors.New("test get epoch error")
	errTestGetEpochData := errors.New("test get epoch data error")
	tests := []struct {
		name                     string
		header                   *types.Header
		setupVerificationManager func(t *testing.T, ctrl *gomock.Controller) *VerificationManager
		expErr                   error
	}{
		{
			name:   "get_epoch_error",
			header: testBlockHeaderEmpty,
			setupVerificationManager: func(t *testing.T, ctrl *gomock.Controller) *VerificationManager {
				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.
					EXPECT().
					GetHeader(testBlockHeaderEmpty.ParentHash).
					Return(defaultParentHeader, nil)

				mockEpochStateGetEpochErr := NewMockEpochState(ctrl)
				mockEpochStateGetEpochErr.EXPECT().GetEpochForBlock(testBlockHeaderEmpty).
					Return(uint64(0), errTestGetEpoch)

				return NewVerificationManager(mockBlockState,
					NewMockSlotState(nil), mockEpochStateGetEpochErr)
			},
			expErr: fmt.Errorf("getting epoch for block header: %w", errTestGetEpoch),
		},
		{
			name: "failed_to_get_epoch_data_raw",
			setupVerificationManager: func(t *testing.T, ctrl *gomock.Controller) *VerificationManager {
				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.
					EXPECT().
					GetHeader(testBlockHeaderEmpty.ParentHash).
					Return(defaultParentHeader, nil)

				mockBlockState.
					EXPECT().
					GenesisHash().
					Return(defaultParentHeader.Hash())

				mockEpochState := NewMockEpochState(ctrl)
				mockEpochState.EXPECT().GetSlotDuration().Return(6*time.Second, nil)
				mockEpochState.EXPECT().GetEpochForBlock(testBlockHeaderEmpty).Return(uint64(1), nil)
				mockEpochState.EXPECT().GetEpochDataRaw(uint64(1), testBlockHeaderEmpty).
					Return(nil, errTestGetEpochData)

				return NewVerificationManager(mockBlockState, NewMockSlotState(nil), mockEpochState)
			},
			header: testBlockHeaderEmpty,
			expErr: fmt.Errorf("getting verifier info: "+
				"failed to get epoch data for epoch 1: %w", errTestGetEpochData),
		},
		{
			name: "failed_to_get_epoch_config_data",
			setupVerificationManager: func(t *testing.T, ctrl *gomock.Controller) *VerificationManager {
				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.
					EXPECT().
					GetHeader(testBlockHeaderEmpty.ParentHash).
					Return(defaultParentHeader, nil)

				mockBlockState.
					EXPECT().
					GenesisHash().
					Return(defaultParentHeader.Hash())

				mockEpochState := NewMockEpochState(ctrl)
				mockEpochState.EXPECT().GetSlotDuration().Return(6*time.Second, nil)
				mockEpochState.EXPECT().GetEpochForBlock(testBlockHeaderEmpty).Return(uint64(1), nil)
				mockEpochState.EXPECT().GetEpochDataRaw(uint64(1), testBlockHeaderEmpty).
					Return(&types.EpochDataRaw{}, nil)
				mockEpochState.EXPECT().GetConfigData(uint64(1), testBlockHeaderEmpty).
					Return(nil, errTestGetEpochData)

				return NewVerificationManager(mockBlockState, NewMockSlotState(nil), mockEpochState)
			},
			header: testBlockHeaderEmpty,
			expErr: fmt.Errorf("getting verifier info: "+
				"failed to get config data: %w", errTestGetEpochData),
		},
		{
			name: "successfully_validate_block",
			setupVerificationManager: func(t *testing.T, ctrl *gomock.Controller) *VerificationManager {
				mockBlockState := NewMockBlockState(ctrl)
				mockBlockState.
					EXPECT().
					GetHeader(headerWithPreRuntimeDigest.ParentHash).
					Return(defaultParentHeader, nil)

				mockBlockState.
					EXPECT().
					GenesisHash().
					Return(defaultParentHeader.Hash()).
					Times(2)

				mockEpochState := NewMockEpochState(ctrl)
				mockEpochState.EXPECT().GetSlotDuration().Return(6*time.Second, nil)
				mockEpochState.EXPECT().GetEpochForBlock(headerWithPreRuntimeDigest).Return(uint64(1), nil)
				mockEpochState.EXPECT().GetEpochDataRaw(uint64(1), headerWithPreRuntimeDigest).
					Return(&types.EpochDataRaw{
						Authorities: defaultEpoch1Authorities,
						Randomness:  [32]byte{},
					}, nil)
				mockEpochState.EXPECT().GetConfigData(uint64(1), headerWithPreRuntimeDigest).
					Return(defaultConfigData, nil)

				mockSlotState := NewMockSlotState(ctrl)
				mockSlotState.EXPECT().
					CheckEquivocation(
						gomock.Any(), uint64(defaultSlotNumber),
						headerWithPreRuntimeDigest,
						[32]byte(kp.Public().Encode())).
					Return(nil, nil)

				return NewVerificationManager(mockBlockState, mockSlotState, mockEpochState)
			},
			header: headerWithPreRuntimeDigest,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := tt.setupVerificationManager(t, gomock.NewController(t))
			err := v.VerifyBlock(tt.header)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func buildSealDigest(t *testing.T, header *types.Header, kp *sr25519.Keypair) *types.SealDigest {
	t.Helper()

	encHeader, err := scale.Marshal(*header)
	require.NoError(t, err)

	hash, err := common.Blake2bHash(encHeader)
	require.NoError(t, err)

	sig, err := kp.Sign(hash[:])
	require.NoError(t, err)

	return &types.SealDigest{
		ConsensusEngineID: types.BabeEngineID,
		Data:              sig,
	}
}

func TestVerificationManager_SetOnDisabled(t *testing.T) {
	// Generate keys
	kp, err := sr25519.GenerateKeypair()
	assert.NoError(t, err)

	testHeader := types.NewEmptyHeader()
	testHeader.Number = 2

	ctrl := gomock.NewController(t)
	mockBlockStateEmpty := NewMockBlockState(ctrl)
	mockBlockStateIsDescendantErr := NewMockBlockState(ctrl)
	mockBlockStateAuthorityDisabled := NewMockBlockState(ctrl)
	mockBlockStateOk := NewMockBlockState(ctrl)

	mockEpochStateGetEpochErr := NewMockEpochState(ctrl)
	mockEpochStateGetEpochDataErr := NewMockEpochState(ctrl)
	mockEpochStateIndexLenErr := NewMockEpochState(ctrl)
	mockEpochStateSetDisabledProd := NewMockEpochState(ctrl)
	mockEpochStateOk := NewMockEpochState(ctrl)
	mockEpochStateOk2 := NewMockEpochState(ctrl)
	mockEpochStateOk3 := NewMockEpochState(ctrl)

	errTestGetEpoch := errors.New("test get epoch error")
	mockEpochStateGetEpochErr.EXPECT().GetEpochForBlock(types.NewEmptyHeader()).Return(uint64(0), errTestGetEpoch)

	mockEpochStateGetEpochDataErr.EXPECT().GetEpochForBlock(types.NewEmptyHeader()).Return(uint64(0), nil)
	errTestGetEpochData := errors.New("test get epoch data error")
	mockEpochStateGetEpochDataErr.EXPECT().
		GetEpochDataRaw(uint64(0), types.NewEmptyHeader()).
		Return(nil, errTestGetEpochData)

	mockEpochStateIndexLenErr.EXPECT().GetEpochForBlock(types.NewEmptyHeader()).Return(uint64(2), nil)

	mockEpochStateSetDisabledProd.EXPECT().GetEpochForBlock(types.NewEmptyHeader()).Return(uint64(2), nil)

	mockEpochStateOk.EXPECT().GetEpochForBlock(types.NewEmptyHeader()).Return(uint64(2), nil)
	errTestDescendant := errors.New("test descendant error")
	mockBlockStateIsDescendantErr.EXPECT().IsDescendantOf(gomock.Any(), gomock.Any()).Return(false, errTestDescendant)

	mockEpochStateOk2.EXPECT().GetEpochForBlock(testHeader).Return(uint64(2), nil)
	mockBlockStateAuthorityDisabled.EXPECT().IsDescendantOf(gomock.Any(), gomock.Any()).Return(true, nil)

	mockEpochStateOk3.EXPECT().GetEpochForBlock(testHeader).Return(uint64(2), nil)
	mockBlockStateOk.EXPECT().IsDescendantOf(gomock.Any(), gomock.Any()).Return(false, nil)

	mockSlotState := NewMockSlotState(nil)

	authority := types.NewAuthority(kp.Public(), uint64(1))
	info := &verifierInfo{
		authorities:    []types.AuthorityRaw{*authority.ToRaw(), *authority.ToRaw()},
		threshold:      scale.MaxUint128,
		secondarySlots: true,
	}

	disabledInfo := []*onDisabledInfo{
		{
			blockNumber: 2,
		},
	}

	vm0 := NewVerificationManager(mockBlockStateEmpty, mockSlotState, mockEpochStateGetEpochErr)
	vm1 := NewVerificationManager(mockBlockStateEmpty, mockSlotState, mockEpochStateGetEpochDataErr)
	vm1.epochInfo[1] = info

	vm2 := NewVerificationManager(mockBlockStateEmpty, mockSlotState, mockEpochStateIndexLenErr)
	vm2.epochInfo[2] = info

	vm3 := NewVerificationManager(mockBlockStateEmpty, mockSlotState, mockEpochStateSetDisabledProd)
	vm3.epochInfo[2] = info

	vm4 := NewVerificationManager(mockBlockStateIsDescendantErr, mockSlotState, mockEpochStateOk)
	vm4.epochInfo[2] = info
	vm4.onDisabled[2] = map[uint32][]*onDisabledInfo{}
	vm4.onDisabled[2][0] = disabledInfo

	vm5 := NewVerificationManager(mockBlockStateAuthorityDisabled, mockSlotState, mockEpochStateOk2)
	vm5.epochInfo[2] = info
	vm5.onDisabled[2] = map[uint32][]*onDisabledInfo{}
	vm5.onDisabled[2][0] = disabledInfo

	vm6 := NewVerificationManager(mockBlockStateOk, mockSlotState, mockEpochStateOk3)
	vm6.epochInfo[2] = info
	vm6.onDisabled[2] = map[uint32][]*onDisabledInfo{}
	vm6.onDisabled[2][0] = disabledInfo

	type args struct {
		index  uint32
		header *types.Header
	}
	tests := []struct {
		name   string
		vm     *VerificationManager
		args   args
		expErr error
	}{
		{
			name: "get_epoch_err",
			vm:   vm0,
			args: args{
				header: types.NewEmptyHeader(),
			},
			expErr: errTestGetEpoch,
		},
		{
			name: "get_epoch_data_err",
			vm:   vm1,
			args: args{
				header: types.NewEmptyHeader(),
			},
			expErr: fmt.Errorf("failed to get epoch data for epoch %d: %w", 0, errTestGetEpochData),
		},
		{
			name: "index_length_error",
			vm:   vm2,
			args: args{
				index:  10000,
				header: types.NewEmptyHeader(),
			},
			expErr: ErrInvalidBlockProducerIndex,
		},
		{
			name: "set_disabled_producers",
			vm:   vm3,
			args: args{
				header: types.NewEmptyHeader(),
			},
		},
		{
			name: "is_Descendant_of_err",
			vm:   vm4,
			args: args{
				header: types.NewEmptyHeader(),
			},
			expErr: errTestDescendant,
		},
		{
			name: "authority_already_disabled",
			vm:   vm5,
			args: args{
				header: testHeader,
			},
			expErr: ErrAuthorityAlreadyDisabled,
		},
		{
			name: "happy_path",
			vm:   vm6,
			args: args{
				header: testHeader,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := tt.vm
			err := v.SetOnDisabled(tt.args.index, tt.args.header)
			if tt.expErr != nil {
				assert.EqualError(t, err, tt.expErr.Error())
			} else {
				assert.NoError(t, err)
			}
		})
	}
}
