package api_v2

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	pdclient "github.com/pingcap/endless/pkg/client/pd"
	_ "github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"github.com/tikv/client-go/v2/oracle"
	"github.com/tikv/client-go/v2/rawkv"
	pd "github.com/tikv/pd/client"
	"go.uber.org/zap"
)

const (
	DEFAULT_GC_PADDING_LEN = 1000
)

// Example of running with parameters:
// `ginkgo -focus correctness_01 -- RawKV -threads=50 -keys=10000 -duration=1m`
var _ = ParameterizedDescribe("RawKVGC", func(flagSet *flag.FlagSet) {
	ctx := suiteTestCtx
	var (
		flagThreads             int64
		flagKeysPerThread       int64
		flagPadding             []byte
		flagDuration            time.Duration
		flagGCSafePointDuration time.Duration

		tc     resource.TiDBCluster
		pdAddr string
	)

	flagSet.Int64Var(&flagThreads, "threads", 20, "Number of threads")
	flagSet.Int64Var(&flagKeysPerThread, "keys", 5, "Number of keys per thread")
	flagSet.DurationVar(&flagDuration, "duration", 120*time.Second, "Duration (e.g. 1m30s) to run")
	flagSet.DurationVar(&flagGCSafePointDuration, "gc_safepoint_duration", 5*time.Second, "Duration (e.g. 1m30s) to run")

	var paddingLen int
	flagSet.IntVar(&paddingLen, "pad", DEFAULT_GC_PADDING_LEN, "Value padding length")
	flagPadding = []byte(strings.Repeat("0", paddingLen))

	// Do some write and read operation and do some random sleep in the process.
	gcWorkload := func(ctx context.Context, cli *rawkv.Client, key []byte, value0 []byte, value1 []byte) []byte {

		Try(cli.Put(ctx, key, value0))
		Try(cli.Put(ctx, key, value1))

		time.Sleep(time.Duration(rand.Intn(10)) * time.Second)
		// Verify order by Get
		res := Try(cli.Get(ctx, key)).([]byte)
		Expect(res).To(Equal(value1))

		time.Sleep(time.Duration(rand.Intn(10)) * time.Second)
		// Verify order by Scan
		keys, values, err := cli.Scan(ctx, key, append(key, 0), 100)
		Expect(err).Should(Succeed())
		Expect(len(keys)).To(Equal(1))
		Expect(keys[0]).To(Equal(key))
		Expect(values[0]).To(Equal(value1))

		// Verify Delete
		Try(cli.Delete(ctx, key))
		time.Sleep(time.Duration(rand.Intn(10)) * time.Second)
		res = Try(cli.Get(ctx, key)).([]byte)
		Expect(res).To(BeEmpty())

		// Verify again
		Try(cli.Put(ctx, key, value1))
		Try(cli.Put(ctx, key, value0))

		time.Sleep(time.Duration(rand.Intn(10)) * time.Second)
		res = Try(cli.Get(ctx, key)).([]byte)
		Expect(res).To(Equal(value0))

		return res
	}

	gcRangeWorkload := func(ctx context.Context, cli *rawkv.Client, keys [][]byte, values0 [][]byte, values1 [][]byte) {
		Expect(keys).NotTo(BeEmpty())

		startKey := keys[0]
		endKey := append(keys[len(keys)-1], 0) // append 0 to make endKey inclusive

		Try(cli.DeleteRange(ctx, startKey, endKey))

		for i, key := range keys {
			_ = gcWorkload(ctx, cli, key, values0[i], values1[i])
		}

		time.Sleep(time.Duration(rand.Intn(10)) * time.Second)
		res_keys, res_values, err := cli.Scan(ctx, startKey, endKey, len(keys)+1) // +1 to verify that there is no more items than expected.
		Expect(err).Should(Succeed())
		Expect(len(keys)).To(Equal(len(res_keys)))
		Expect(res_keys).To(Equal(keys))
		Expect(res_values).To(Equal(values0))
	}

	generateTestDataImpl := func(keyRange int64, keysPerThread int64, threads int64, padding []byte, idx int64) (keys [][]byte, values0 [][]byte, values1 [][]byte) {
		step := int64(keyRange / keysPerThread / threads)
		if step <= 0 {
			step = 1
		}

		keys = make([][]byte, 0, keysPerThread)
		values0 = make([][]byte, 0, keysPerThread)
		values1 = make([][]byte, 0, keysPerThread)

		start := idx * keysPerThread * step
		end := (idx + 1) * keysPerThread * step
		for k := start; k < end; k += step {
			keys = append(keys, []byte(fmt.Sprintf("k%016d", k)))
			values0 = append(values0, append([]byte(fmt.Sprintf("%v_", k)), padding...))
			values1 = append(values1, append([]byte(fmt.Sprintf("%v_", k+1)), padding...))
		}
		return
	}

	generateTestData := func(idx int64) (keys [][]byte, values0 [][]byte, values1 [][]byte) {
		return generateTestDataImpl(KEY_RANGE, flagKeysPerThread, flagThreads, flagPadding, idx)
	}

	schedulerRoutine := func() {
		// TODO: add operator of leader transfer, region split, and region merge.
		time.Sleep(flagDuration)
		log.Info("scheduler routine finished")
	}

	BeforeEach(func() {
		tc = ctx.Resource("tc").(resource.TiDBCluster)

		pdClient := Try(pdclient.NewPDClient(tc)).(*pdclient.Client)
		pdAddr = pdClient.GetPDAddr()

		log.Info("BeforeEach", zap.Int64("threads", flagThreads), zap.Int64("keysPerThread", flagKeysPerThread),
			zap.Duration("duration", flagDuration), zap.String("pdAddr", pdAddr))
	})

	Context("ApiV2", func() {

		TestCase("Correctness", "correctness_gc", func() {
			var wg sync.WaitGroup
			closeCh := make(chan struct{})

			for idx := int64(0); idx < int64(flagThreads); idx++ {
				wg.Add(1)
				idx := idx
				go func() {
					defer GinkgoRecover()
					defer wg.Done()

					keys, values0, values1 := generateTestData(idx)
					startKey := keys[0]
					endKey := keys[len(keys)-1]

				Loop:
					for {
						ctx, cancel := context.WithCancel(context.Background())
						cli, err := rawkv.NewClientWithOpts(ctx, []string{pdAddr}, rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
						Expect(err).Should(Succeed())
						defer cli.Close()

						log.Info("worker loop", zap.Uint64("clusterID", cli.ClusterID()), zap.Int64("thread", idx),
							zap.String("start", string(startKey)), zap.String("end", string(endKey)))

						gcRangeWorkload(ctx, cli, keys, values0, values1)
						log.Info("worker loop end")

						cancel()

						select {
						case <-closeCh:
							break Loop
						default:
						}
					}
				}()
			}

			wg.Add(1)
			go func() {
				defer GinkgoRecover()
				defer wg.Done()
			Loop:
				for {
					timeout := time.Second * 10
					gcOffset := time.Second * 1
					ctx, cancel := context.WithTimeout(context.Background(), timeout)
					defer cancel()
					var (
						pdclt_gc_safepoint, err = pd.NewClientWithContext(ctx, []string{pdAddr}, pd.SecurityOption{
							CAPath:   "",
							CertPath: "",
							KeyPath:  "",
						})
					)
					if err != nil {
						log.Panic("create pd client failed", zap.Error(err))
					}

					p, l, err := pdclt_gc_safepoint.GetTS(ctx)
					if err != nil {
						log.Panic("get ts failed", zap.Error(err))
					}
					now := oracle.ComposeTS(p, l)
					nowMinusOffset := oracle.GetTimeFromTS(now).Add(-gcOffset)
					newSP := oracle.ComposeTS(oracle.GetPhysical(nowMinusOffset), 0)
					_, err = pdclt_gc_safepoint.UpdateGCSafePoint(ctx, newSP)
					if err != nil {
						log.Panic("update safe point failed", zap.Error(err))
					}
					log.Info("update GC safe point", zap.Uint64("SP", newSP), zap.Uint64("now", now))

					time.Sleep(flagGCSafePointDuration)

					cancel()

					select {
					case <-closeCh:
						break Loop
					default:
					}
				}

			}()

			// scheduler routine
			wg.Add(1)
			go func() {
				defer GinkgoRecover()
				defer wg.Done()

				schedulerRoutine()
				close(closeCh)
			}()

			wg.Wait()
		})

	})
})