package tikv_cdc

import (
	"bytes"
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/log"
	"github.com/tikv/client-go/v2/rawkv"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

type OpType int

const (
	DoWorkload OpType = iota
	Verify
	DoAndVerfiy
)

var (
	PrefixMap sync.Map
)

type WorkloadRunner struct {
	SrcClient       *rawkv.Client
	DstClient       *rawkv.Client
	KeyCnt          int64
	ThreadCnt       int64
	WaitSinkTimeout time.Duration
	TimeSleepMs     int64

	KeyPrefix string
	StartKey  []byte
	EndKey    []byte
}

func NewWorkloadRunner(ctx context.Context, srcPDAddr string, dstPDAddr string,
	keyCnt int64, threadCnt int64, waitSinkTimeout time.Duration) (*WorkloadRunner, error) {

	srcCli, err := rawkv.NewClientWithOpts(ctx, []string{srcPDAddr},
		rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
	if err != nil {
		return nil, errors.Trace(err)
	}
	dstCli, err := rawkv.NewClientWithOpts(ctx, []string{dstPDAddr},
		rawkv.WithAPIVersion(kvrpcpb.APIVersion_V2))
	if err != nil {
		return nil, errors.Trace(err)
	}

	var (
		keyPrefix string
		startKey  []byte
		endKey    []byte
	)
	rand.Seed(time.Now().UnixNano())
	for {
		keyPrefix = string(byte(rand.Intn(25)) + 'a')
		_, ok := PrefixMap.LoadOrStore(keyPrefix, 0)
		if !ok {
			startKey, _, _ = GenerateTestData(keyPrefix, 0)
			endKey, _, _ = GenerateTestData(keyPrefix, keyCnt)
			break
		}
	}

	log.Info("Start to clean src tikv data", zap.ByteString("start key", startKey),
		zap.ByteString("end key", endKey))
	srcCli.DeleteRange(ctx, startKey, endKey)
	log.Info("Start to clean dst tikv data", zap.ByteString("start key", startKey),
		zap.ByteString("end key", endKey))
	dstCli.DeleteRange(ctx, startKey, endKey)

	srcChecksum, err := GetChecksum(ctx, srcCli, startKey, endKey)
	if err != nil {
		return nil, errors.Trace(err)
	}
	dstChecksum, err := GetChecksum(ctx, dstCli, startKey, endKey)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if dstChecksum != srcChecksum {
		return nil, errors.Errorf("checksums are not equal: %v != %v", dstChecksum, srcChecksum)
	}

	workloadRunner := &WorkloadRunner{
		SrcClient:       srcCli,
		DstClient:       dstCli,
		KeyCnt:          keyCnt,
		ThreadCnt:       threadCnt,
		WaitSinkTimeout: waitSinkTimeout,
		TimeSleepMs:     0,
		KeyPrefix:       keyPrefix,
		StartKey:        startKey,
		EndKey:          endKey,
	}

	return workloadRunner, nil
}

func (r *WorkloadRunner) Close() {
	r.SrcClient.Close()
	r.DstClient.Close()
	PrefixMap.LoadAndDelete(r.KeyPrefix)
}

func (r *WorkloadRunner) SetAtomicForCAS(flag bool) {
	r.SrcClient.SetAtomicForCAS(flag)
}

func (r *WorkloadRunner) SetTimeSleep(timeSleepMs int64) {
	r.TimeSleepMs = timeSleepMs
}

func (r *WorkloadRunner) RunWorkload(ctx context.Context, workload Workload,
	op OpType) error {

	doWorkload := func() error {
		err := workload.Func(ctx, r.KeyCnt, r.ThreadCnt, r.SrcClient, r.KeyPrefix, r.TimeSleepMs)
		if err != nil {
			return errors.Trace(err)
		}
		log.Info("Load data succeeded", zap.Int64("keycnt", r.KeyCnt), zap.Int64("threadCnt", r.ThreadCnt))
		return nil
	}

	verify := func() error {
		err := r.WaitForTiKVCdcSinkFinish(ctx)
		if err != nil {
			return errors.Trace(err)
		}
		verifyFunc := VerifyPut
		if workload.IsDelete {
			verifyFunc = VerifyDelete
		}
		err = r.VerifyData(ctx, verifyFunc)
		if err != nil {
			return errors.Trace(err)
		}
		return nil
	}

	switch op {
	case DoAndVerfiy:
		err := doWorkload()
		if err != nil {
			return errors.Trace(err)
		}
		err = verify()
		if err != nil {
			return errors.Trace(err)
		}
	case DoWorkload:
		err := doWorkload()
		if err != nil {
			return errors.Trace(err)
		}
	case Verify:
		err := verify()
		if err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func (r *WorkloadRunner) VerifyData(ctx context.Context,
	verifyFunc func(context.Context, *rawkv.Client, string, int64) error) error {
	log.Info("Start to verify data")

	cli := r.DstClient
	keyCnt, concurrency := r.KeyCnt, r.ThreadCnt

	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++ {
				err := verifyFunc(ctx, cli, r.KeyPrefix, j)
				if err != nil {
					return errors.Trace(err)
				}
			}
			return nil
		})
	}
	return errGroup.Wait()
}

// WaitForTiKVCdcSinkFinish put a key to src cluster and wait it sink to dst cluster.
func (r *WorkloadRunner) WaitForTiKVCdcSinkFinish(ctx context.Context) error {
	log.Info("Wait for TiKV-CDC to finish")

	srcCli, dstCli := r.SrcClient, r.DstClient
	startKey, endKey := r.StartKey, r.EndKey
	timeout := r.WaitSinkTimeout

	srcChecksum, err := GetChecksum(ctx, srcCli, startKey, endKey)
	if err != nil {
		return errors.Trace(err)
	}
	curTime := time.Now()
	dstChecksum := rawkv.RawChecksum{}
	for time.Since(curTime) < timeout {
		dstChecksum, err = GetChecksum(ctx, dstCli, startKey, endKey)
		if err != nil {
			return errors.Trace(err)
		}
		if dstChecksum == srcChecksum {
			fmt.Printf("\t...Wait cdc finish and takes %f seconds, src checksum %v, dst checksum %v\n",
				time.Since(curTime).Seconds(), srcChecksum, dstChecksum)
			return nil
		}
		fmt.Printf("\t...Src checksum %v are not same with dst %v, wait for next time, cur takes %f s.\n",
			srcChecksum, dstChecksum, time.Since(curTime).Seconds())
		time.Sleep(time.Duration(10) * time.Second)
	}
	return errors.Errorf("Wait cdc finish timeout, %f s, src checksum:%v, dst checksum:%v",
		time.Since(curTime).Seconds(), srcChecksum, dstChecksum)
}

func VerifyPut(ctx context.Context, cli *rawkv.Client, prefix string, k int64) error {
	key, _, value := GenerateTestData(prefix, k)
	res_value, err := cli.Get(ctx, key)
	if err != nil {
		return errors.Trace(err)
	}
	if !bytes.Equal(res_value, value) {
		return errors.Errorf("inconsistency value, key:%s, expected:%s, got:%s",
			key, value, res_value)
	}
	return nil
}

func VerifyDelete(ctx context.Context, cli *rawkv.Client, prefix string, k int64) error {
	key, _, _ := GenerateTestData(prefix, k)
	res_value, err := cli.Get(ctx, key)
	if err != nil {
		return errors.Trace(err)
	}
	if len(res_value) != 0 {
		return errors.Errorf("inconsistency value, key:%s, expect empty, got:%s",
			key, res_value)
	}
	return nil
}

func GetChecksum(ctx context.Context, cli *rawkv.Client, startKey, endKey []byte) (rawkv.RawChecksum, error) {
	return cli.Checksum(ctx, startKey, endKey)
}