package main

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

	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/store/tikv"
)

type ConsistencyCheckCase struct {
	concurrency int
	db          *tikv.RawKVClient
	pdAddrs     []string
	tikvAddrs   []string
}

type Config struct {
	Concurrency int `toml:concurrency`
	PDAddrs     []string
	TiKVAddrs   []string
}

func NewConsistencyCheckCase(cfg *Config) *ConsistencyCheckCase {
	b := &ConsistencyCheckCase{
		concurrency: cfg.Concurrency,
		pdAddrs:     cfg.PDAddrs,
		db:          nil,
	}
	return b
}

func (c *ConsistencyCheckCase) String() string {
	return "consistency_check"
}

func (c *ConsistencyCheckCase) Initialize(ctx context.Context) error {
	db, err := tikv.NewRawKVClient(c.pdAddrs, config.Security{})
	if err != nil {
		return err
	}
	c.db = db
	InitClient(c.pdAddrs[0], c.tikvAddrs)

	if err := StartShuffleRegions(); err != nil {
		return err
	}
	return nil
}

func (c *ConsistencyCheckCase) Execute(ctx context.Context) error {
	client.Infof("[%s] start to test...", c)
	defer func() {
		client.Infof("[%s] test end...", c)
	}()

	r := rand.New(rand.NewSource(time.Now().Unix()))

	ctx, cancel := context.WithCancel(ctx)
	var wg sync.WaitGroup
	for i := 0; i < c.concurrency; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			ticker := time.NewTicker(10 * time.Second)
			defer ticker.Stop()

			var offset uint64
			prev := []byte(fmt.Sprintf("%d_", i))

			for {
				prefix := []byte(fmt.Sprintf("%d_%d_", i, offset))
				len := r.Intn(128)
				value := make([]byte, len+1)
				util.RandString(value, r)
				key := append(prefix, value...)

				if err := c.db.Put(key, value); err != nil {
					client.Fatalf("put value err %s", err)
				}

				if offset == math.MaxUint64 {
					return
				}
				offset++

				select {
				case <-ctx.Done():
					return
				case <-ticker.C:
					client.Infof("[%s] detroy range %v %v", c, prev, prefix)
					DestroyRange(ctx, prev, prefix)
					prev = prefix
				default:
				}
			}
		}(i)
	}

	cancel()
	wg.Wait()
	Close()

	return nil
}