// Copyright 2023 TiKV Project Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client_test

import (
	"path"
	"strconv"
	"testing"
	"time"

	"github.com/pingcap/failpoint"
	"github.com/pingcap/kvproto/pkg/keyspacepb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/log"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	pd "github.com/tikv/pd/client"
	"github.com/tikv/pd/pkg/keyspace"
	"github.com/tikv/pd/pkg/utils/assertutil"
	"github.com/tikv/pd/pkg/utils/keypath"
	"github.com/tikv/pd/pkg/utils/testutil"
	"github.com/tikv/pd/server"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

// gcClientTestReceiver is the pdpb.PD_WatchGCSafePointV2Server mock for testing.
type gcClientTestReceiver struct {
	re *require.Assertions
	grpc.ServerStream
}

// Send is the mock implementation for pdpb.PD_WatchGCSafePointV2Server's Send.
// Instead of sending the response to the client, it will check the response.
// In testing, we will set all keyspace's safe point to be equal to its id,
// and this mock verifies that the response is correct.
func (s gcClientTestReceiver) Send(m *pdpb.WatchGCSafePointV2Response) error {
	log.Info("received", zap.Any("received", m.GetEvents()))
	for _, change := range m.GetEvents() {
		s.re.Equal(change.SafePoint, uint64(change.KeyspaceId))
	}
	return nil
}

type gcClientTestSuite struct {
	suite.Suite
	server              *server.GrpcServer
	client              pd.Client
	cleanup             testutil.CleanupFunc
	gcSafePointV2Prefix string
	keyspaces           []*keyspacepb.KeyspaceMeta
}

func TestGcClientTestSuite(t *testing.T) {
	suite.Run(t, new(gcClientTestSuite))
}

func (suite *gcClientTestSuite) SetupSuite() {
	re := suite.Require()
	var (
		err error
		gsi *server.Server
	)
	checker := assertutil.NewChecker()
	checker.FailNow = func() {}
	gsi, suite.cleanup, err = server.NewTestServer(re, checker)
	suite.server = &server.GrpcServer{Server: gsi}
	re.NoError(err)
	addr := suite.server.GetAddr()
	suite.client, err = pd.NewClientWithContext(suite.server.Context(), []string{addr}, pd.SecurityOption{})
	re.NoError(err)
	rootPath := path.Join("/pd", strconv.FormatUint(keypath.ClusterID(), 10))
	suite.gcSafePointV2Prefix = path.Join(rootPath, keypath.GCSafePointV2Prefix())
	// Enable the fail-point to skip checking keyspace validity.
	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/gc/checkKeyspace", "return(true)"))
	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)"))
	// Create keyspace and set `keyspace.SafePointVersion` is `v2`.
	suite.keyspaces = mustMakeTestKeyspaces(re, suite.server.Server, 100, true)
}

func (suite *gcClientTestSuite) TearDownSuite() {
	re := suite.Require()
	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/gc/checkKeyspace"))
	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion"))
	suite.cleanup()
	suite.client.Close()
}

func (suite *gcClientTestSuite) TearDownTest() {
	suite.CleanupEtcdGCPath()
}

func (suite *gcClientTestSuite) CleanupEtcdGCPath() {
	re := suite.Require()
	_, err := suite.server.GetClient().Delete(suite.server.Context(), suite.gcSafePointV2Prefix, clientv3.WithPrefix())
	re.NoError(err)
}

func (suite *gcClientTestSuite) TestWatch1() {
	re := suite.Require()
	receiver := gcClientTestReceiver{re: suite.Require()}
	go suite.server.WatchGCSafePointV2(&pdpb.WatchGCSafePointV2Request{
		Revision: 0,
	}, receiver)

	// Init gc safe points as index value of keyspace 0 ~ 5.

	for i := 1; i < 6; i++ {
		suite.mustUpdateSafePoint(re, uint32(i), uint64(i))
	}

	// delete gc safe points of keyspace 3 ~ 5.
	for i := 3; i < 6; i++ {
		suite.mustDeleteSafePoint(re, uint32(i))
	}

	// check gc safe point equal to keyspace id for keyspace 0 ~ 2 .

	for i := 1; i < 3; i++ {
		re.Equal(uint64(i), suite.mustLoadSafePoint(re, uint32(i)))
	}

	// check gc safe point is 0 for keyspace 3 ~ 5 after delete.
	for i := 3; i < 6; i++ {
		re.Equal(uint64(0), suite.mustLoadSafePoint(re, uint32(i)))
	}
}

func (suite *gcClientTestSuite) TestClientWatchWithRevision() {
	suite.testClientWatchWithRevision(false)
	suite.testClientWatchWithRevision(true)
}

// nolint:revive
func (suite *gcClientTestSuite) testClientWatchWithRevision(fromNewRevision bool) {
	re := suite.Require()
	testKeyspaceID := uint32(10)
	initGCSafePoint := uint64(50)
	updatedGCSafePoint := uint64(100)

	// Init gc safe point.
	suite.mustUpdateSafePoint(re, testKeyspaceID, initGCSafePoint)

	// Get the initial revision.
	initRevision := suite.mustGetRevision(re, testKeyspaceID)

	// Update the gc safe point.
	suite.mustUpdateSafePoint(re, testKeyspaceID, updatedGCSafePoint)

	// Get the revision of the updated gc safe point.
	updatedRevision := suite.mustGetRevision(re, testKeyspaceID)

	// Set the start revision of the watch request based on fromNewRevision.
	startRevision := initRevision
	if fromNewRevision {
		startRevision = updatedRevision
	}
	watchChan, err := suite.client.WatchGCSafePointV2(suite.server.Context(), startRevision)
	re.NoError(err)

	timer := time.NewTimer(time.Second)
	defer timer.Stop()
	isFirstUpdate := true
	runTest := false
	for {
		select {
		case <-timer.C:
			re.True(runTest)
			return
		case res := <-watchChan:
			for _, r := range res {
				re.Equal(r.GetKeyspaceId(), testKeyspaceID)
				if fromNewRevision {
					// If fromNewRevision, first response should be the updated gc safe point.
					re.Equal(r.GetSafePoint(), updatedGCSafePoint)
				} else if isFirstUpdate {
					isFirstUpdate = false
					re.Equal(r.GetSafePoint(), initGCSafePoint)
				} else {
					re.Equal(r.GetSafePoint(), updatedGCSafePoint)
					continue
				}
			}
			runTest = true
		}
	}
}

// mustUpdateSafePoint updates the gc safe point of the given keyspace id.
func (suite *gcClientTestSuite) mustUpdateSafePoint(re *require.Assertions, keyspaceID uint32, safePoint uint64) {
	_, err := suite.client.UpdateGCSafePointV2(suite.server.Context(), keyspaceID, safePoint)
	re.NoError(err)
}

// mustLoadSafePoint loads the gc safe point of the given keyspace id.
func (suite *gcClientTestSuite) mustLoadSafePoint(re *require.Assertions, keyspaceID uint32) uint64 {
	res, err := suite.server.GetSafePointV2Manager().LoadGCSafePoint(keyspaceID)
	re.NoError(err)
	return res.SafePoint
}

// mustDeleteSafePoint deletes the gc safe point of the given keyspace id.
func (suite *gcClientTestSuite) mustDeleteSafePoint(re *require.Assertions, keyspaceID uint32) {
	safePointPath := path.Join(suite.gcSafePointV2Prefix, keypath.EncodeKeyspaceID(keyspaceID))
	log.Info("test etcd path", zap.Any("path", safePointPath)) // TODO: Delete
	_, err := suite.server.GetClient().Delete(suite.server.Context(), safePointPath)
	re.NoError(err)
}

// mustGetRevision gets the revision of the given keyspace's gc safe point.
func (suite *gcClientTestSuite) mustGetRevision(re *require.Assertions, keyspaceID uint32) int64 {
	safePointPath := path.Join(suite.gcSafePointV2Prefix, keypath.EncodeKeyspaceID(keyspaceID))
	res, err := suite.server.GetClient().Get(suite.server.Context(), safePointPath)
	re.NoError(err)
	return res.Header.GetRevision()
}

// mustUpdateServiceSafePoint updates the service safe point of the given keyspace id.
func (suite *gcClientTestSuite) mustUpdateServiceSafePoint(keyspaceID uint32, safePoint uint64, ttl int64, serviceID string) uint64 {
	res, err := suite.client.UpdateServiceSafePointV2(suite.server.Context(), keyspaceID, serviceID, ttl, safePoint)
	suite.NoError(err)
	return res
}

// mustUpdateServiceSafePoint updates the service safe point of the given keyspace id.
func (suite *gcClientTestSuite) mustLoadServiceSafePoint(keyspaceID uint32, serviceID string) uint64 {
	res, err := suite.server.GetSafePointV2Manager().LoadServiceSafePointV2(keyspaceID, serviceID)
	suite.NoError(err)
	return res.SafePoint
}

func (suite *gcClientTestSuite) TestUpdateServiceSafePointV2() {
	serviceID := "test_service_id"
	ttl := int64(10000)

	keyspaceID01 := uint32(1)
	serviceSafePoint01 := uint64(1)
	_ = suite.mustUpdateServiceSafePoint(keyspaceID01, serviceSafePoint01, ttl, serviceID)
	minSafePoint01 := suite.mustLoadServiceSafePoint(keyspaceID01, serviceID)
	suite.Equal(serviceSafePoint01, minSafePoint01)

	keyspaceID02 := uint32(2)
	serviceSafePoint02 := uint64(2)
	_ = suite.mustUpdateServiceSafePoint(keyspaceID02, serviceSafePoint02, ttl, serviceID)
	minSafePoint02 := suite.mustLoadServiceSafePoint(keyspaceID02, serviceID)
	suite.Equal(serviceSafePoint02, minSafePoint02)
}

func (suite *gcClientTestSuite) TestGetGCSafePointV2() {
	re := suite.Require()

	// The keyspace.SafePointVersion of current keyspace is v2, when obtaining the GC V2 safe point,
	// if the GC safe point has not yet been generated, it should not get the GC safe point from GC V1.

	expectGCV1SafePoint := uint64(100)
	re.Equal(keyspace.KeyspaceGlobalSafePointVersionV2, suite.keyspaces[0].Config[keyspace.SafePointVersion])

	// Set GC V1 GC safe point.
	_, err := suite.server.Server.GetSafePointV1Manager().UpdateGCSafePoint(expectGCV1SafePoint)
	re.NoError(err)
	gcV1SafePoint, err := suite.server.Server.GetSafePointV1Manager().LoadGCSafePoint()
	re.NoError(err)
	re.Equal(expectGCV1SafePoint, gcV1SafePoint)

	// Check GC V2 GC safe point is not equal to GC V1 GC safe point.
	res, err := suite.server.Server.GetSafePointV2Manager().LoadGCSafePoint(suite.keyspaces[0].Id)
	re.NoError(err)
	re.NotEqual(expectGCV1SafePoint, res.SafePoint)
	re.Equal(uint64(0), res.SafePoint)
}