package api_v2

import (
	"context"
	"math/rand"
	"strconv"
	"sync/atomic"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/testcase/rawkv/pkg/client"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/pkg/util/pdutil"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

func getAndClearRange(ctx context.Context, cli client.Client, keys [][]byte, clearRange bool) (startKey []byte, endKey []byte) {
	Expect(keys).NotTo(BeEmpty())

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

	if clearRange {
		Try(cli.DeleteRange(ctx, startKey, endKey))
	}
	return
}

func PointWorkload(ctx context.Context, cli client.Client, pdCli *pdutil.Client, key []byte, value0 []byte, value1 []byte) []byte {
	keyCopy := make([]byte, len(key))
	copy(keyCopy, key)

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

	// Verify order by Get
	res := Try(cli.Get(ctx, key)).([]byte)
	Expect(res).To(Equal(value1), "key: %v", string(key))

	// 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))
	res = Try(cli.Get(ctx, key)).([]byte)
	Expect(res).To(BeEmpty(), "key: %v", string(key))

	// Verify again
	Try(cli.Put(ctx, key, value1))
	Try(cli.Put(ctx, key, value0))
	res = Try(cli.Get(ctx, key)).([]byte)
	Expect(res).To(Equal(value0), "key: %v", string(key))

	// Verify key unchanged
	Expect(key).To(Equal(keyCopy))
	return res
}

// `keys` is expected to be in order.
func RangePointWorkload(ctx context.Context, cli client.Client, pdCli *pdutil.Client, keys [][]byte, values0 [][]byte, values1 [][]byte, clearRange bool) {
	startKey, endKey := getAndClearRange(ctx, cli, keys, clearRange)

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

	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(res_keys)).To(Equal(len(keys)), "startKey: %v, endKey: %v", string(startKey), string(endKey))
	Expect(res_keys).To(Equal(keys))
	Expect(res_values).To(Equal(values0))
}

// `keys` is expected to be in order.
func BatchWorkload(ctx context.Context, cli client.Client, pdCli *pdutil.Client, keys [][]byte, values0 [][]byte, values1 [][]byte, clearRange bool) {
	startKey, endKey := getAndClearRange(ctx, cli, keys, clearRange)

	// Verify order of BatchPut
	Try(cli.BatchPut(ctx, keys, values0))
	Try(cli.BatchPut(ctx, keys, values1))
	res := Try(cli.BatchGet(ctx, keys)).([][]byte)
	Expect(res).To(Equal(values1), "startKey: %v, endKey: %v, keys: %v", string(startKey), string(endKey), keys)

	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(res_keys)).To(Equal(len(keys)), "startKey: %v, endKey: %v", string(startKey), string(endKey))
	Expect(res_keys).To(Equal(keys))
	Expect(res_values).To(Equal(values1))

	// Verify BatchDelete
	Try(cli.BatchDelete(ctx, keys))
	res_keys, _, err = cli.Scan(ctx, startKey, endKey, len(keys)+1) // verify empty
	Expect(err).Should(Succeed())
	Expect(res_keys).To(BeEmpty(), "startKey: %v, endKey: %v", string(startKey), string(endKey))

	// Verify order of BatchPut again
	Try(cli.BatchPut(ctx, keys, values1))
	Try(cli.BatchPut(ctx, keys, values0))
	res = Try(cli.BatchGet(ctx, keys)).([][]byte)
	Expect(res).To(Equal(values0), "startKey: %v, endKey: %v, keys: %v", string(startKey), string(endKey), keys)

	// Verify ReverseScan
	res_keys, res_values, err = cli.ReverseScan(ctx, endKey, startKey, len(keys)+1) // +1 to verify that there is no more items than expected.
	Expect(err).Should(Succeed())
	Expect(len(res_keys)).To(Equal(len(keys)), "startKey: %v, endKey: %v", string(startKey), string(endKey))
	Expect(len(res_values)).To(Equal(len(keys)), "startKey: %v, endKey: %v", string(startKey), string(endKey))
	for i := 0; i < len(keys); i += 1 {
		j := len(keys) - 1 - i
		Expect(res_keys[i]).To(Equal(keys[j]))
		Expect(res_values[i]).To(Equal(values0[j]))
	}
}

const CAS_WORKLOAD_THREADS = 5

// CASAcumulateWorkload is not stable
// When cli.CompareAndSwap timeout, it is possible that the operation is actually successful.
// So the retry will add to value more than expected.
func CASAcumulateWorkload(ctx context.Context, cli client.Client, key []byte) {
	num := rand.Int()
	Try(cli.Delete(ctx, key))

	eg := new(errgroup.Group)
	for idx := 0; idx <= CAS_WORKLOAD_THREADS; idx++ {
		eg.Go(func() error {
			defer GinkgoRecover()

			var (
				err       error
				prevValue []byte
				newValue  []byte
			)
			succeed := false
			for !succeed {
				if prevValue == nil {
					newValue = []byte(strconv.Itoa(num))
				} else {
					prevNum := Try(strconv.Atoi(string(prevValue))).(int)
					newValue = []byte(strconv.Itoa(prevNum + 1))
				}
				prevValue, succeed, err = cli.CompareAndSwap(ctx, key, prevValue, newValue)
				Expect(err).To(Succeed())
			}

			return nil
		})
	}
	Try(eg.Wait())

	value := Try(cli.Get(ctx, key)).([]byte)
	expected := []byte(strconv.Itoa(num + CAS_WORKLOAD_THREADS))
	Expect(value).To(Equal(expected))
}

func CASLockingWorkload(ctx context.Context, cli client.Client, key []byte) {
	num := rand.Int()
	Try(cli.Delete(ctx, key))
	value := Try(cli.Get(ctx, key)).([]byte)
	Expect(value).Should(BeNil())

	var counter int32
	var idxLocked int32

	eg := new(errgroup.Group)
	for idx := 0; idx < CAS_WORKLOAD_THREADS; idx++ {
		idx := idx
		eg.Go(func() error {
			defer GinkgoRecover()

			var (
				err       error
				prevValue []byte
				newValue  []byte = []byte(strconv.Itoa(num + idx))
				succeed   bool   = false
			)
			for !succeed && prevValue == nil {
				prevValue, succeed, err = cli.CompareAndSwap(ctx, key, nil, newValue)
				Expect(err).To(Succeed())
			}

			if succeed {
				atomic.AddInt32(&counter, 1)
				atomic.StoreInt32(&idxLocked, int32(idx))
			}
			return nil
		})
	}
	Try(eg.Wait())

	if atomic.LoadInt32(&counter) == 0 {
		// No thread acquired the lock. It would be caused by previous request timeout. Check locking by value.
		value = Try(cli.Get(ctx, key)).([]byte)
		Expect(value).ShouldNot(BeNil())
		intValue := Try(strconv.Atoi(string(value))).(int)
		idxLocked = int32(intValue - num)

		Expect(idxLocked).To(BeNumerically(">=", 0))
		Expect(idxLocked).To(BeNumerically("<", CAS_WORKLOAD_THREADS))
		log.Warn("No thread succeed. Check locking by value", zap.Int32("idxLocked", idxLocked), zap.String("key", string(key)))
	} else {
		Expect(atomic.LoadInt32(&counter)).To(Equal(int32(1)))

		value := Try(cli.Get(ctx, key)).([]byte)
		expected := []byte(strconv.Itoa(num + int(idxLocked)))
		Expect(value).To(Equal(expected))
	}
}

func RangeCASWorkload(ctx context.Context, cli client.Client, _ *pdutil.Client, keys [][]byte, _ [][]byte, _ [][]byte, clearRange bool) {
	getAndClearRange(ctx, cli, keys, clearRange)

	for _, key := range keys {
		CASLockingWorkload(ctx, cli, key)
	}
}

func TTLWorkload(ctx context.Context, cli client.Client, key []byte, value []byte, ttl_in_seconds uint64) {
	Try(cli.PutWithTTL(ctx, key, value, ttl_in_seconds))
	res := Try(cli.Get(ctx, key)).([]byte)
	if res == nil {
		log.Warn("Not found after PutWithTTL. Would be caused by high load", zap.String("key", string(key)),
			zap.String("value", string(value)), zap.Uint64("ttl", ttl_in_seconds))
	} else {
		Expect(res).To(Equal(value))
	}

	ttl := Try(cli.GetKeyTTL(ctx, key)).(*uint64)
	if ttl == nil {
		log.Warn("Not found after PutWithTTL. Would be caused by high load", zap.String("key", string(key)),
			zap.String("value", string(value)), zap.Uint64("ttl", ttl_in_seconds))
	} else {
		Expect(*ttl).Should(BeNumerically("<=", ttl_in_seconds))
		time.Sleep(time.Duration(*ttl)*time.Second + time.Second/5)
	}

	res = Try(cli.Get(ctx, key)).([]byte)
	Expect(res).Should(BeNil())
}

const (
	TTLWorkloadConcurrency  = 5
	TTLWorkloadTTLInSeconds = 3
)

func RangeTTLWorkload(ctx context.Context, cli client.Client, _ *pdutil.Client, keys [][]byte, values0 [][]byte, _ [][]byte, clearRange bool) {
	startKey, endKey := getAndClearRange(ctx, cli, keys, clearRange)

	// There will be too many goroutine if run on every key
	// So limit to TTLWorkloadConcurrency
	batchSize := TTLWorkloadConcurrency
	if len(keys) < batchSize {
		batchSize = len(keys)
	}
	batchEndKey := append(keys[batchSize-1], 0)

	eg := new(errgroup.Group)
	for i := 0; i < batchSize; i++ {
		i := i
		eg.Go(func() error {
			defer GinkgoRecover()

			TTLWorkload(ctx, cli, keys[i], values0[i], TTLWorkloadTTLInSeconds)
			return nil
		})
	}
	Try(eg.Wait())

	// Should be empty
	res_keys, _, err := cli.Scan(ctx, startKey, batchEndKey, batchSize+1) // Limit `batchSize+1` to see the actual number of result.
	Expect(err).Should(Succeed())
	if len(res_keys) != 0 {
		// Would be not empty. Caused by local clock ?
		log.Warn("RangeTTLWorkload batch not empty. Try again.", zap.String("startKey", string(startKey)),
			zap.String("batchEndKey", string(batchEndKey)), zap.Int("len(res_keys)", len(res_keys)))

		time.Sleep(time.Duration(TTLWorkloadTTLInSeconds) * time.Second)
		res_keys, _, err = cli.Scan(ctx, startKey, batchEndKey, batchSize+1)
		Expect(err).Should(Succeed())
		Expect(len(res_keys)).To(Equal(0), "startKey: %v, endKey: %v", string(startKey), string(endKey))
	}

	// BatchPutWithTTL
	ttls := make([]uint64, 0, len(keys))
	for range keys {
		ttls = append(ttls, TTLWorkloadTTLInSeconds)
	}
	Try(cli.BatchPutWithTTL(ctx, keys, values0, ttls))

	// Should no more than `len(keys)`
	res_keys, _, err = cli.Scan(ctx, startKey, endKey, len(keys)+1)
	Expect(err).Should(Succeed())
	Expect(len(res_keys)).Should(BeNumerically("<=", len(keys)), "startKey: %v, endKey: %v", string(startKey), string(endKey))
	if len(res_keys) == 0 {
		log.Warn("Not found after BatchPutWithTTL. Would be caused by high load", zap.String("startKey", string(startKey)),
			zap.String("endKey", string(endKey)), zap.Uint64("ttl", TTLWorkloadTTLInSeconds), zap.Int("len(keys)", len(keys)))
	}

	time.Sleep(time.Duration(TTLWorkloadTTLInSeconds+1) * time.Second) // +1 second to avoid local clock deviation

	// Should be empty
	res_keys, _, err = cli.Scan(ctx, startKey, endKey, len(keys)+1)
	Expect(err).Should(Succeed())
	Expect(res_keys).Should(BeEmpty(), "keys: %v", keys)
}

type RangeWorkloadFunc func(ctx context.Context, cli client.Client, pCli *pdutil.Client, keys [][]byte, values0 [][]byte, values1 [][]byte, clearRange bool)

type WorkloadWeight struct {
	fn     RangeWorkloadFunc
	atomic bool
	weight int
}

type Workloads struct {
	weights []WorkloadWeight
	counter []int32
}

func NewWorkloads(weights []WorkloadWeight) Workloads {
	w := Workloads{
		weights: weights,
		counter: make([]int32, len(weights)),
	}
	w.init()
	return w
}

func (w *Workloads) init() {
	for i := range w.weights {
		if i > 0 {
			w.weights[i].weight += w.weights[i-1].weight
		}
	}
}

func (w *Workloads) Pick() (RangeWorkloadFunc, bool) {
	if len(w.weights) == 0 {
		return nil, false
	}
	n := rand.Intn(w.weights[len(w.weights)-1].weight)
	for i, weight := range w.weights {
		if n < weight.weight {
			atomic.AddInt32(&w.counter[i], 1)
			return weight.fn, weight.atomic
		}
	}
	return nil, false
}

func (w *Workloads) Counter() []int32 {
	return w.counter
}