package alphabet

import (
	"context"
	"fmt"
	"math/big"
	"testing"

	preimage "github.com/ethereum-optimism/optimism/op-preimage"

	"github.com/ethereum-optimism/optimism/op-challenger/game/fault/types"

	"github.com/ethereum/go-ethereum/common"
	"github.com/stretchr/testify/require"
)

func alphabetClaim(index *big.Int, claim *big.Int) common.Hash {
	return alphabetStateHash(BuildAlphabetPreimage(index, claim))
}

func TestAlphabetProvider_Prestate(t *testing.T) {
	depth := types.Depth(4)
	startingL2BlockNumber := big.NewInt(2)

	// Actual preimage values generated by the solidity AlphabetVM at each step.
	expectedPrestates := []string{
		"0000000000000000000000000000000000000000000000000000000000000060",
		"00000000000000000000000000000000000000000000000000000000000000210000000000000000000000000000000000000000000000000000000000000081",
		"00000000000000000000000000000000000000000000000000000000000000220000000000000000000000000000000000000000000000000000000000000082",
		"00000000000000000000000000000000000000000000000000000000000000230000000000000000000000000000000000000000000000000000000000000083",
		"00000000000000000000000000000000000000000000000000000000000000240000000000000000000000000000000000000000000000000000000000000084",
		"00000000000000000000000000000000000000000000000000000000000000250000000000000000000000000000000000000000000000000000000000000085",
	}

	ap := NewTraceProvider(startingL2BlockNumber, depth)

	for i, expected := range expectedPrestates {
		i, expected := i, expected
		t.Run(fmt.Sprintf("Step_%v", i), func(t *testing.T) {
			result, _, _, err := ap.GetStepData(context.Background(), types.NewPosition(4, big.NewInt(int64(i))))
			require.NoError(t, err)
			require.Equalf(t, expected, common.Bytes2Hex(result), "Incorrect prestate at trace index %v", i)
		})
	}
}

func TestAlphabetProvider_GetStepData_MaxLen(t *testing.T) {
	depth := types.Depth(4)
	startingL2BlockNumber := big.NewInt(2)
	ap := NewTraceProvider(startingL2BlockNumber, depth)

	// Step data for the max position is allowed
	maxLen := int64(1 << depth)
	maxPos := types.NewPosition(4, big.NewInt(maxLen))
	result, _, _, err := ap.GetStepData(context.Background(), maxPos)
	require.NoError(t, err)
	expected := "00000000000000000000000000000000000000000000000000000000000000300000000000000000000000000000000000000000000000000000000000000090"
	require.Equal(t, expected, common.Bytes2Hex(result))

	// Cannot step on a position greater than the max.
	oobPos := types.NewPosition(4, big.NewInt(int64(1<<depth)+1))
	_, _, _, err = ap.GetStepData(context.Background(), oobPos)
	require.ErrorIs(t, err, ErrIndexTooLarge)
}

// TestAlphabetProvider_Get_ClaimsByTraceIndex tests the Get function.
func TestAlphabetProvider_Get_ClaimsByTraceIndex(t *testing.T) {
	// Create a new alphabet provider.
	depth := types.Depth(3)
	startingL2BlockNumber := big.NewInt(1)
	sbn := new(big.Int).Lsh(startingL2BlockNumber, 4)
	startingTraceIndex := new(big.Int).Add(absolutePrestateInt, sbn)
	canonicalProvider := NewTraceProvider(startingL2BlockNumber, depth)

	// Build a list of traces.
	traces := []struct {
		traceIndex   types.Position
		expectedHash common.Hash
	}{
		{
			types.NewPosition(depth, big.NewInt(7)),
			alphabetClaim(new(big.Int).Add(sbn, big.NewInt(8)), new(big.Int).Add(startingTraceIndex, big.NewInt(8))),
		},
		{
			types.NewPosition(depth, big.NewInt(3)),
			alphabetClaim(new(big.Int).Add(sbn, big.NewInt(4)), new(big.Int).Add(startingTraceIndex, big.NewInt(4))),
		},
		{
			types.NewPosition(depth, big.NewInt(5)),
			alphabetClaim(new(big.Int).Add(sbn, big.NewInt(6)), new(big.Int).Add(startingTraceIndex, big.NewInt(6))),
		},
	}

	// Execute each trace and check the alphabet provider returns the expected hash.
	for i, trace := range traces {
		expectedHash, err := canonicalProvider.Get(context.Background(), trace.traceIndex)
		require.NoError(t, err)
		require.Equalf(t, trace.expectedHash, expectedHash, "Trace %v", i)
	}
}

// TestAlphabetProvider_GetStepData tests the GetStepData function.
func TestAlphabetProvider_GetStepData(t *testing.T) {
	depth := types.Depth(2)
	startingL2BlockNumber := big.NewInt(1)
	ap := NewTraceProvider(startingL2BlockNumber, depth)
	key := preimage.LocalIndexKey(L2ClaimBlockNumberLocalIndex).PreimageKey()
	expectedPreimageData := types.NewPreimageOracleData(key[:], startingL2BlockNumber.Bytes(), 0)

	tests := []struct {
		name                 string
		indexAtDepth         *big.Int
		expectedResult       []byte
		expectedPreimageData *types.PreimageOracleData
		expectedError        error
	}{
		{
			name:                 "AbsolutePrestate",
			indexAtDepth:         big.NewInt(0),
			expectedResult:       absolutePrestate,
			expectedPreimageData: expectedPreimageData,
			expectedError:        nil,
		},
		{
			name:                 "SecondStep",
			indexAtDepth:         big.NewInt(1),
			expectedResult:       BuildAlphabetPreimage(big.NewInt(17), big.NewInt(113)),
			expectedPreimageData: expectedPreimageData,
			expectedError:        nil,
		},
		{
			name:                 "LastStep",
			indexAtDepth:         big.NewInt(4),
			expectedResult:       BuildAlphabetPreimage(big.NewInt(20), big.NewInt(116)),
			expectedPreimageData: expectedPreimageData,
			expectedError:        nil,
		},
		{
			name:                 "IndexTooLarge",
			indexAtDepth:         big.NewInt(5),
			expectedResult:       nil,
			expectedPreimageData: nil,
			expectedError:        ErrIndexTooLarge,
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {

			result, proof, data, err := ap.GetStepData(context.Background(), types.NewPosition(depth, test.indexAtDepth))
			require.Equal(t, test.expectedResult, result)
			require.Empty(t, proof)
			require.Equal(t, test.expectedPreimageData, data)
			if test.expectedError != nil {
				require.ErrorIs(t, err, test.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

// TestAlphabetProvider_Get tests the Get function.
func TestAlphabetProvider_Get(t *testing.T) {
	tests := []struct {
		name          string
		depth         types.Depth
		indexAtDepth  *big.Int
		expectedClaim common.Hash
		expectedError error
	}{
		{
			name:          "AbsolutePrestate",
			depth:         types.Depth(2),
			indexAtDepth:  big.NewInt(0),
			expectedClaim: alphabetClaim(big.NewInt(17), big.NewInt(113)),
			expectedError: nil,
		},
		{
			name:          "SecondStep",
			depth:         types.Depth(2),
			indexAtDepth:  big.NewInt(1),
			expectedClaim: alphabetClaim(big.NewInt(18), big.NewInt(114)),
			expectedError: nil,
		},
		{
			name:          "LastStep",
			depth:         types.Depth(2),
			indexAtDepth:  big.NewInt(3),
			expectedClaim: alphabetClaim(big.NewInt(20), big.NewInt(116)),
			expectedError: nil,
		},
		{
			name:          "IndexTooLarge",
			depth:         types.Depth(2),
			indexAtDepth:  big.NewInt(5),
			expectedClaim: common.Hash{},
			expectedError: ErrIndexTooLarge,
		},
		{
			name:          "DepthTooLarge",
			depth:         types.Depth(3),
			indexAtDepth:  big.NewInt(0),
			expectedClaim: common.Hash{},
			expectedError: ErrIndexTooLarge,
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			startingL2BlockNumber := big.NewInt(1)
			ap := NewTraceProvider(startingL2BlockNumber, types.Depth(2))

			result, err := ap.Get(context.Background(), types.NewPosition(test.depth, test.indexAtDepth))
			require.Equal(t, test.expectedClaim, result)
			if test.expectedError != nil {
				require.ErrorIs(t, err, test.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}
