package debug

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/OffchainLabs/prysm/v6/api"
	"github.com/OffchainLabs/prysm/v6/api/server/structs"
	blockchainmock "github.com/OffchainLabs/prysm/v6/beacon-chain/blockchain/testing"
	dbtest "github.com/OffchainLabs/prysm/v6/beacon-chain/db/testing"
	doublylinkedtree "github.com/OffchainLabs/prysm/v6/beacon-chain/forkchoice/doubly-linked-tree"
	forkchoicetypes "github.com/OffchainLabs/prysm/v6/beacon-chain/forkchoice/types"
	"github.com/OffchainLabs/prysm/v6/beacon-chain/rpc/testutil"
	"github.com/OffchainLabs/prysm/v6/config/params"
	"github.com/OffchainLabs/prysm/v6/encoding/bytesutil"
	"github.com/OffchainLabs/prysm/v6/runtime/version"
	"github.com/OffchainLabs/prysm/v6/testing/assert"
	"github.com/OffchainLabs/prysm/v6/testing/require"
	"github.com/OffchainLabs/prysm/v6/testing/util"
	"github.com/ethereum/go-ethereum/common/hexutil"
)

func TestGetBeaconStateV2(t *testing.T) {
	ctx := t.Context()
	db := dbtest.SetupDB(t)

	t.Run("phase0", func(t *testing.T) {
		fakeState, err := util.NewBeaconState()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Phase0), resp.Version)
		st := &structs.BeaconState{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
	})
	t.Run("Altair", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateAltair()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Altair), resp.Version)
		st := &structs.BeaconStateAltair{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
	})
	t.Run("Bellatrix", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateBellatrix()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Bellatrix), resp.Version)
		st := &structs.BeaconStateBellatrix{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
	})
	t.Run("Capella", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateCapella()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Capella), resp.Version)
		st := &structs.BeaconStateCapella{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
	})
	t.Run("Deneb", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateDeneb()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Deneb), resp.Version)
		st := &structs.BeaconStateDeneb{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
	})
	t.Run("Electra", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateElectra()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Electra), resp.Version)
		st := &structs.BeaconStateElectra{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
	})
	t.Run("Fulu", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateFulu()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))
		chainService := &blockchainmock.ChainService{}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, version.String(version.Fulu), resp.Version)
		st := &structs.BeaconStateFulu{}
		require.NoError(t, json.Unmarshal(resp.Data, st))
		assert.Equal(t, "123", st.Slot)
		assert.Equal(t, int(params.BeaconConfig().MinSeedLookahead+1)*int(params.BeaconConfig().SlotsPerEpoch), len(st.ProposerLookahead))
	})
	t.Run("execution optimistic", func(t *testing.T) {
		parentRoot := [32]byte{'a'}
		blk := util.NewBeaconBlock()
		blk.Block.ParentRoot = parentRoot[:]
		root, err := blk.Block.HashTreeRoot()
		require.NoError(t, err)
		util.SaveBlock(t, ctx, db, blk)
		require.NoError(t, db.SaveGenesisBlockRoot(ctx, root))

		fakeState, err := util.NewBeaconStateBellatrix()
		require.NoError(t, err)
		chainService := &blockchainmock.ChainService{Optimistic: true}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
			BeaconDB:              db,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, true, resp.ExecutionOptimistic)
	})
	t.Run("finalized", func(t *testing.T) {
		parentRoot := [32]byte{'a'}
		blk := util.NewBeaconBlock()
		blk.Block.ParentRoot = parentRoot[:]
		root, err := blk.Block.HashTreeRoot()
		require.NoError(t, err)
		util.SaveBlock(t, ctx, db, blk)
		require.NoError(t, db.SaveGenesisBlockRoot(ctx, root))

		fakeState, err := util.NewBeaconStateBellatrix()
		require.NoError(t, err)
		headerRoot, err := fakeState.LatestBlockHeader().HashTreeRoot()
		require.NoError(t, err)
		chainService := &blockchainmock.ChainService{
			FinalizedRoots: map[[32]byte]bool{
				headerRoot: true,
			},
		}
		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
			FinalizationFetcher:   chainService,
			BeaconDB:              db,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetBeaconStateV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, true, resp.Finalized)
	})
}

func TestGetBeaconStateSSZV2(t *testing.T) {
	t.Run("Phase 0", func(t *testing.T) {
		fakeState, err := util.NewBeaconState()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))

		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		request.Header.Set("Accept", api.OctetStreamMediaType)
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		assert.Equal(t, version.String(version.Phase0), writer.Header().Get(api.VersionHeader))
		sszExpected, err := fakeState.MarshalSSZ()
		require.NoError(t, err)
		assert.DeepEqual(t, sszExpected, writer.Body.Bytes())
	})
	t.Run("Altair", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateAltair()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))

		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		request.Header.Set("Accept", api.OctetStreamMediaType)
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		assert.Equal(t, version.String(version.Altair), writer.Header().Get(api.VersionHeader))
		sszExpected, err := fakeState.MarshalSSZ()
		require.NoError(t, err)
		assert.DeepEqual(t, sszExpected, writer.Body.Bytes())
	})
	t.Run("Bellatrix", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateBellatrix()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))

		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		request.Header.Set("Accept", api.OctetStreamMediaType)
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		assert.Equal(t, version.String(version.Bellatrix), writer.Header().Get(api.VersionHeader))
		sszExpected, err := fakeState.MarshalSSZ()
		require.NoError(t, err)
		assert.DeepEqual(t, sszExpected, writer.Body.Bytes())
	})
	t.Run("Capella", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateCapella()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))

		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		request.Header.Set("Accept", api.OctetStreamMediaType)
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		assert.Equal(t, version.String(version.Capella), writer.Header().Get(api.VersionHeader))
		sszExpected, err := fakeState.MarshalSSZ()
		require.NoError(t, err)
		assert.DeepEqual(t, sszExpected, writer.Body.Bytes())
	})
	t.Run("Deneb", func(t *testing.T) {
		fakeState, err := util.NewBeaconStateDeneb()
		require.NoError(t, err)
		require.NoError(t, fakeState.SetSlot(123))

		s := &Server{
			Stater: &testutil.MockStater{
				BeaconState: fakeState,
			},
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/states/{state_id}", nil)
		request.SetPathValue("state_id", "head")
		request.Header.Set("Accept", api.OctetStreamMediaType)
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetBeaconStateV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		assert.Equal(t, version.String(version.Deneb), writer.Header().Get(api.VersionHeader))
		sszExpected, err := fakeState.MarshalSSZ()
		require.NoError(t, err)
		assert.DeepEqual(t, sszExpected, writer.Body.Bytes())
	})
}

func TestGetForkChoiceHeadsV2(t *testing.T) {
	expectedSlotsAndRoots := []struct {
		Slot string
		Root string
	}{{
		Slot: "0",
		Root: hexutil.Encode(bytesutil.PadTo([]byte("foo"), 32)),
	}, {
		Slot: "1",
		Root: hexutil.Encode(bytesutil.PadTo([]byte("bar"), 32)),
	}}

	chainService := &blockchainmock.ChainService{}
	s := &Server{
		HeadFetcher:           chainService,
		OptimisticModeFetcher: chainService,
	}

	request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/heads", nil)
	writer := httptest.NewRecorder()
	writer.Body = &bytes.Buffer{}

	s.GetForkChoiceHeadsV2(writer, request)
	require.Equal(t, http.StatusOK, writer.Code)
	resp := &structs.GetForkChoiceHeadsV2Response{}
	require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
	assert.Equal(t, 2, len(resp.Data))
	for _, sr := range expectedSlotsAndRoots {
		found := false
		for _, h := range resp.Data {
			if h.Slot == sr.Slot {
				found = true
				assert.Equal(t, sr.Root, h.Root)
			}
			assert.Equal(t, false, h.ExecutionOptimistic)
		}
		assert.Equal(t, true, found, "Expected head not found")
	}

	t.Run("optimistic head", func(t *testing.T) {
		chainService := &blockchainmock.ChainService{
			Optimistic:      true,
			OptimisticRoots: make(map[[32]byte]bool),
		}
		for _, sr := range expectedSlotsAndRoots {
			b, err := hexutil.Decode(sr.Root)
			require.NoError(t, err)
			chainService.OptimisticRoots[bytesutil.ToBytes32(b)] = true
		}
		s := &Server{
			HeadFetcher:           chainService,
			OptimisticModeFetcher: chainService,
		}

		request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v2/debug/beacon/heads", nil)
		writer := httptest.NewRecorder()
		writer.Body = &bytes.Buffer{}

		s.GetForkChoiceHeadsV2(writer, request)
		require.Equal(t, http.StatusOK, writer.Code)
		resp := &structs.GetForkChoiceHeadsV2Response{}
		require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
		assert.Equal(t, 2, len(resp.Data))
		for _, sr := range expectedSlotsAndRoots {
			found := false
			for _, h := range resp.Data {
				if h.Slot == sr.Slot {
					found = true
					assert.Equal(t, sr.Root, h.Root)
				}
				assert.Equal(t, true, h.ExecutionOptimistic)
			}
			assert.Equal(t, true, found, "Expected head not found")
		}
	})
}

func TestGetForkChoice(t *testing.T) {
	store := doublylinkedtree.New()
	fRoot := [32]byte{'a'}
	fc := &forkchoicetypes.Checkpoint{Epoch: 2, Root: fRoot}
	require.NoError(t, store.UpdateFinalizedCheckpoint(fc))
	s := &Server{ForkchoiceFetcher: &blockchainmock.ChainService{ForkChoiceStore: store}}

	request := httptest.NewRequest(http.MethodGet, "http://example.com/eth/v1/debug/fork_choice", nil)
	writer := httptest.NewRecorder()
	writer.Body = &bytes.Buffer{}

	s.GetForkChoice(writer, request)
	require.Equal(t, http.StatusOK, writer.Code)
	resp := &structs.GetForkChoiceDumpResponse{}
	require.NoError(t, json.Unmarshal(writer.Body.Bytes(), resp))
	require.Equal(t, "2", resp.FinalizedCheckpoint.Epoch)
}
