package tikv_cdc

import (
	"bytes"
	"context"
	"time"

	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	_ "github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/tikv/client-go/v2/rawkv"
)

const (
	MAX_KV_CNT_PER_BATCH int64 = 1024

	ClientRetryCnt      int           = 10
	ClientRetryInterval time.Duration = 300 * time.Millisecond
)

func min(left, right int64) int64 {
	if left < right {
		return left
	} else {
		return right
	}
}

type Workload struct {
	Func     WorkloadFunc
	IsDelete bool
}

var (
	PointPutWorkload    = Workload{PointLoadData, false}
	BatchPutWorkload    = Workload{BatchLoadData, false}
	CASWorkload         = Workload{CASLoadData, false}
	PointDeleteWorkload = Workload{PointDeleteData, true}
	BatchDeleteWorkload = Workload{BatchDeleteData, true}
)
var (
	Workloads = []Workload{
		BatchPutWorkload,
		BatchDeleteWorkload,
	}

	AtomicWorkloads = []Workload{
		BatchPutWorkload,
		BatchDeleteWorkload,
		CASWorkload,
		BatchDeleteWorkload,
	}
)

type WorkloadFunc func(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client, prefix string, timeSleepMs int64) error

func BatchLoadData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client, prefix string, timeSleepMs int64) error {
	log.Info("Run BatchLoadData")
	errGroup := new(errgroup.Group)
	keyCntPerBatch := keyCnt / concurrency
	timeSleepMs = timeSleepMs * MAX_KV_CNT_PER_BATCH

	for i := int64(0); i < concurrency; i++ {
		startIdx := i * keyCntPerBatch
		endIdx := (i + 1) * keyCntPerBatch
		errGroup.Go(func() error {
			for startIdx < endIdx {
				start := time.Now()
				batchCnt := min(MAX_KV_CNT_PER_BATCH, endIdx-startIdx)
				keys, values0, values1 := BatchGenerateData(prefix, startIdx, batchCnt)
				fn := func() error {
					err := cli.BatchPut(ctx, keys, values0)
					if err != nil {
						return errors.Trace(err)
					}
					err = cli.BatchPut(ctx, keys, values1)
					if err != nil {
						return errors.Trace(err)
					}
					return nil
				}
				err := ClientRetryWithError(fn, "BatchPut")
				if err != nil {
					return err
				}
				timeCost := time.Since(start).Milliseconds()
				if reminder := timeSleepMs - timeCost; reminder > 0 {
					time.Sleep(time.Duration(reminder) * time.Millisecond)
				}
				startIdx += MAX_KV_CNT_PER_BATCH
			}
			return nil
		})
	}
	return errGroup.Wait()
}

func PointLoadData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client, prefix string, timeSleepMs int64) error {
	log.Info("Run PointLoadData")

	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++ {
				start := time.Now()
				key, value0, value1 := GenerateTestData(prefix, j)

				fn := func() error {
					err := cli.Put(ctx, key, value0)
					if err != nil {
						return err
					}
					err = cli.Put(ctx, key, value1)
					if err != nil {
						return err
					}
					return nil
				}
				err := ClientRetryWithError(fn, "PointPut")
				if err != nil {
					return err
				}
				timeCost := time.Since(start).Milliseconds()
				if reminder := timeSleepMs - timeCost; reminder > 0 {
					time.Sleep(time.Duration(reminder) * time.Millisecond)
				}
			}
			return nil
		})
	}
	return errGroup.Wait()
}

func PointDeleteData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client, prefix string, perOpCost int64) error {
	log.Info("Run PointDeleteData")

	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++ {
				start := time.Now()
				key, _, _ := GenerateTestData(prefix, j)
				fn := func() error {
					err := cli.Delete(ctx, key)
					if err != nil {
						return err
					}
					return nil
				}
				err := ClientRetryWithError(fn, "PointDelete")
				if err != nil {
					return err
				}
				timeCost := time.Since(start).Milliseconds()
				if reminder := perOpCost - timeCost; reminder > 0 {
					time.Sleep(time.Duration(reminder) * time.Millisecond)
				}
			}
			return nil
		})
	}
	return errGroup.Wait()
}

func BatchDeleteData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client, prefix string, timeSleepMs int64) error {
	log.Info("Run BatchDeleteData")

	errGroup := new(errgroup.Group)
	keyCntPerBatch := keyCnt / concurrency
	timeSleepMs = timeSleepMs * MAX_KV_CNT_PER_BATCH

	for i := int64(0); i < concurrency; i++ {
		startIdx := i * keyCntPerBatch
		endIdx := (i + 1) * keyCntPerBatch
		errGroup.Go(func() error {
			for startIdx < endIdx {
				start := time.Now()
				batchCnt := min(MAX_KV_CNT_PER_BATCH, endIdx-startIdx)
				keys, _, _ := BatchGenerateData(prefix, startIdx, batchCnt)
				fn := func() error {
					err := cli.BatchDelete(ctx, keys)
					if err != nil {
						return errors.Trace(err)
					}
					return nil
				}
				err := ClientRetryWithError(fn, "BatchDelete")
				if err != nil {
					return err
				}
				timeCost := time.Since(start).Milliseconds()
				if reminder := timeSleepMs - timeCost; reminder > 0 {
					time.Sleep(time.Duration(reminder) * time.Millisecond)
				}
				startIdx += MAX_KV_CNT_PER_BATCH
			}
			return nil
		})
	}
	return errGroup.Wait()
}

func CASLoadData(ctx context.Context, keyCnt int64, concurrency int64, cli *rawkv.Client, prefix string, timeSleepMs int64) error {
	log.Info("Run CASLoadData")

	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++ {
				start := time.Now()
				fn := func() error {
					key, value0, value1 := GenerateTestData(prefix, j)
					err := cli.Put(ctx, key, value0)
					if err != nil {
						return errors.Trace(err)
					}
					preValue, ret, err := cli.CompareAndSwap(ctx, key, value0, value1)
					if err != nil {
						return errors.Trace(err)
					}
					if !ret || !bytes.Equal(preValue, value0) {
						return errors.Errorf("CAS put data error: preValue: %v, ret: %v, value0: %v, value1: %v", preValue, ret, value0, value1)
					}
					return nil
				}
				err := ClientRetryWithError(fn, "CASPut")
				if err != nil {
					return err
				}
				timeCost := time.Since(start).Milliseconds()
				if reminder := timeSleepMs - timeCost; reminder > 0 {
					time.Sleep(time.Duration(reminder) * time.Millisecond)
				}
			}
			return nil
		})
	}
	return errGroup.Wait()
}

func ClientRetryWithError(fn func() error, fnName string) (err error) {
	retry := 0
	for retry < ClientRetryCnt {
		err = fn()
		if err == nil {
			break
		}
		log.Warn("ClientRetryWithError get error", zap.Error(err), zap.Int("retry", retry), zap.String("func name", fnName))
		time.Sleep(ClientRetryInterval)
		retry += 1
	}
	return err
}