package tikv_br

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/google/uuid"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/util"
	. "github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"

	pdclient "github.com/pingcap/endless/pkg/client/pd"
	_ "github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"github.com/tikv/client-go/v2/config"
	"github.com/tikv/client-go/v2/rawkv"
	"go.uber.org/zap"
)

const (
	s3Endpoint        = "http://minio.pingcap.net:9000"
	s3AccessKeyID     = "minioadmin"
	s3SecretAccessKey = "minioadmin"
	maxKeyCntPerBatch = 2048
)

func GenerateTestData(keyIndex int64) (keys []byte, values []byte) {
	key := []byte(fmt.Sprintf("indexInfo_:_pf01_:_APD0101_:_%019d", keyIndex))
	value := []byte(fmt.Sprintf("v_%0150d_%0150d", keyIndex, keyIndex))
	return key, value
}

func PreLoadData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client) error {
	errGroup := new(errgroup.Group)
	keyCntPerBatch := keyCnt / concurrency
	for i := int64(0); i < concurrency; i++ {
		startIdx := i * keyCntPerBatch
		endIdx := (i + 1) * keyCntPerBatch
		errGroup.Go(func() error {
			for j := startIdx; j < endIdx; j++ {
				key, value := GenerateTestData(j)
				err := cli.Put(ctx, key, value)
				if err != nil {
					return err
				}
			}
			return nil
		})
	}
	return errGroup.Wait()
}

func VerifyData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client) error {
	errGroup := new(errgroup.Group)
	keyCntPerBatch := keyCnt / concurrency
	for i := int64(0); i < concurrency; i++ {
		startIdx := i * keyCntPerBatch
		endIdx := (i + 1) * keyCntPerBatch
		errGroup.Go(func() error {
			for j := startIdx; j < endIdx; j++ {
				key, value := GenerateTestData(j)
				res_value, err := cli.Get(ctx, key)
				if err != nil {
					return err
				}
				if !bytes.Equal(res_value, value) {
					return errors.New("inconsistency value")
				}
			}
			return nil
		})
	}
	return errGroup.Wait()
}

// Example of running with parameters:
// `ginkgo -focus rawkv_br_basic_v1 -- TiKVBR -threads=50 -keys=10000`
var _ = ParameterizedDescribe("TiKVBR", func(flagSet *flag.FlagSet) {
	ctx := suiteTestCtx
	var (
		flagThreads  int64
		flagKeyCnts  int64
		flagDuration time.Duration
		s3Parms      S3Parms
		brPath       string

		srcTc     resource.TiDBCluster
		dstTc     resource.TiDBCluster
		br        resource.WorkloadNode
		srcPdAddr string
		dstPdAddr string
	)

	flagSet.Int64Var(&flagThreads, "threads", 2, "Number of threads")
	flagSet.Int64Var(&flagKeyCnts, "keys", 1000, "Number of preload keys")
	flagSet.DurationVar(&flagDuration, "duration", 10*time.Second, "Duration (e.g. 1m30s) to run")

	randomS3Dir := func(prefix string) string {
		uuid := uuid.New()
		key := uuid.String()
		return "s3://tmp/tikv-br-test/" + prefix + "/" + key
	}

	tikvBrTestImpl := func(srcClient, dstClient *rawkv.Client, dstApiVer, casePrefix string) {
		err := PreLoadData(ctx, flagKeyCnts, flagThreads, srcClient)
		Expect(err).Should(Succeed())
		fmt.Printf("load data succeeded, keycnt:%d, thread:%d.\n", flagKeyCnts, flagThreads)
		storage := randomS3Dir(casePrefix)
		TiKVBackupRawData(ctx, srcPdAddr, brPath, br, s3Parms, storage, dstApiVer)
		TiKVRestoreRawData(ctx, dstPdAddr, brPath, br, s3Parms, storage)
		// verify the correctness on target rawkv cluster
		err = VerifyData(ctx, flagKeyCnts, 600, dstClient)
		Expect(err).Should(Succeed())
	}

	BeforeEach(func() {
		srcTc = ctx.Resource("srctc").(resource.TiDBCluster)
		dstTc = ctx.Resource("dsttc").(resource.TiDBCluster)
		br = ctx.Resource("br").(resource.WorkloadNode)

		srcPdClient := Try(pdclient.NewPDClient(srcTc)).(*pdclient.Client)
		srcPdAddr = srcPdClient.GetPDAddr()

		dstPdClient := Try(pdclient.NewPDClient(dstTc)).(*pdclient.Client)
		dstPdAddr = dstPdClient.GetPDAddr()

		s3Parms = S3Parms{
			S3Endpoint:        s3Endpoint,
			S3AccessKeyID:     s3AccessKeyID,
			S3SecretAccessKey: s3SecretAccessKey,
		}
		// `tikv-br` binary path should be added to `PATH`
		brPath = "/disk1/home/haojinming/migration/br/bin/tikv-br"

		log.Info("BeforeEach", zap.Int64("keys", flagKeyCnts),
			zap.Duration("duration", flagDuration), zap.String("srcPdAddr", srcPdAddr),
			zap.String("dstPdAddr", dstPdAddr))
	})

	Context("ApiV2", func() {
		TestCase("TiKVBR", "rawkv_br_basic_v2", func() {
			srcCli, err := rawkv.NewClientWithOpts(ctx, []string{srcPdAddr}, rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
			Expect(err).Should(Succeed())
			defer srcCli.Close()
			log.Info("Client connected", zap.Uint64("ClusterID", srcCli.ClusterID()))

			dstCli, err := rawkv.NewClientWithOpts(ctx, []string{dstPdAddr}, rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
			Expect(err).Should(Succeed())
			defer dstCli.Close()

			tikvBrTestImpl(srcCli, dstCli, "V2", "rawkv_br_basic_v2")
		})

		TestCase("TiKVBR", "rawkv_br_basic_v1", func() {
			srcCli, err := rawkv.NewClient(ctx, []string{srcPdAddr}, config.DefaultConfig().Security)
			Expect(err).Should(Succeed())
			defer srcCli.Close()
			log.Info("Client connected", zap.Uint64("ClusterID", srcCli.ClusterID()))

			dstCli, err := rawkv.NewClient(ctx, []string{dstPdAddr}, config.DefaultConfig().Security)
			Expect(err).Should(Succeed())
			defer dstCli.Close()

			tikvBrTestImpl(srcCli, dstCli, "V1", "rawkv_br_basic_v1")
		})

		TestCase("TiKVBR", "rawkv_br_basic_v1ttl", func() {
			srcCli, err := rawkv.NewClient(ctx, []string{srcPdAddr}, config.DefaultConfig().Security)
			Expect(err).Should(Succeed())
			defer srcCli.Close()
			log.Info("Client connected", zap.Uint64("ClusterID", srcCli.ClusterID()))

			dstCli, err := rawkv.NewClient(ctx, []string{dstPdAddr}, config.DefaultConfig().Security)
			Expect(err).Should(Succeed())
			defer dstCli.Close()

			tikvBrTestImpl(srcCli, dstCli, "V1TTL", "rawkv_br_basic_v1")
		})

		TestCase("TiKVBR", "rawkv_br_v1_to_v2", func() {
			srcCli, err := rawkv.NewClient(ctx, []string{srcPdAddr}, config.DefaultConfig().Security)
			Expect(err).Should(Succeed())
			defer srcCli.Close()
			log.Info("Client connected", zap.Uint64("ClusterID", srcCli.ClusterID()))

			dstCli, err := rawkv.NewClientWithOpts(ctx, []string{dstPdAddr}, rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
			Expect(err).Should(Succeed())
			defer dstCli.Close()

			tikvBrTestImpl(srcCli, dstCli, "V2", "rawkv_br_v1_to_v2")
		})

		TestCase("TiKVBR", "rawkv_br_v1ttl_to_v2", func() {
			srcCli, err := rawkv.NewClient(ctx, []string{srcPdAddr}, config.DefaultConfig().Security)
			Expect(err).Should(Succeed())
			defer srcCli.Close()
			log.Info("Client connected", zap.Uint64("ClusterID", srcCli.ClusterID()))

			dstCli, err := rawkv.NewClientWithOpts(ctx, []string{dstPdAddr}, rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
			Expect(err).Should(Succeed())
			defer dstCli.Close()

			tikvBrTestImpl(srcCli, dstCli, "V2", "rawkv_br_v1ttl_to_v2")
		})
	})
})