package main

import (
	"context"
	"flag"
	"io"
	"math/rand"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/kvproto/pkg/tikvpb"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	grpcStatus "google.golang.org/grpc/status"
)

// min and max delay time passed to the TiKV, in milliseconds.
const minDelay uint64 = 10
const maxDelay uint64 = 1000

var (
	concurrency int
	timeout     int
	runtime     int
)

func init() {
	flag.IntVar(&concurrency, "concurrency", 128, "concurrency to communicate with the TiKV")
	flag.IntVar(&timeout, "timeout", 10, "max duration in seconds for requests")
	flag.IntVar(&runtime, "time", 600, "test running time in seconds")
}

func main() {
	client.Info("[tikvclient] start...")
	ctx, cancel := context.WithCancel(context.Background())

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Infof("[tikvclient] got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()

	cli := client.New(ctx, client.ModeOnline)

	var rpcClient = tikv.NewTestRPCClient()
	defer rpcClient.Close()

	var workerChs = make([]chan struct{}, 0, concurrency)
	var errorCh = make(chan error, 1)
	for i := 0; i < concurrency; i++ {
		workerChs = append(workerChs, make(chan struct{}, 1))
		client.Infof("[tikvclient] test with tikv %s", cli.TiKVs()[0])
		go doTest(workerChs[i], errorCh, rpcClient, cli.TiKVs()[0])
	}

	timer := time.NewTimer(time.Duration(runtime) * time.Second)
	select {
	case err := <-errorCh:
		client.Infof("[tikvclient] exits because of error: %v", err)
		for _, workerCh := range workerChs {
			workerCh <- struct{}{}
		}
		time.Sleep(time.Second)
		return
	case <-timer.C:
		client.Infof("test is finished, now exits")
		return
	}
}

func doTest(ch <-chan struct{}, errorCh chan<- error, rpcClient tikv.Client, tikvAddr string) {
	for {
		select {
		case <-ch:
			return
		default:
		}
		req := &tikvrpc.Request{
			Type: tikvrpc.CmdEmpty,
			Empty: &tikvpb.BatchCommandsEmptyRequest{
				TestId:    getTestID(),
				DelayTime: randomDelay(),
			},
		}
		resp, err := rpcClient.SendRequest(context.Background(), tikvAddr, req, time.Duration(timeout)*time.Second)
		if err != nil {
			err = errors.Cause(err)
			if status, ok := grpcStatus.FromError(err); ok {
				client.Infof("gRPC error is allowed, status: %s", status.Message())
			} else if err == io.EOF {
				client.Infof("denyed by tcp layer is allowed")
				return
			} else {
				client.Errorf("unallowed error: %v", err)
				errorCh <- err
				return
			}
		} else if resp.Empty.TestId != req.Empty.TestId {
			client.Errorf("unmatched test id, want: %d, got: %d", req.Empty.TestId, resp.Empty.TestId)
			errorCh <- errors.New("unmatched test id")
			return
		}
	}
}

var idGen uint64 = 0

func getTestID() uint64 {
	return atomic.AddUint64(&idGen, 1)
}

func randomDelay() uint64 {
	r := rand.Int63n(int64(maxDelay - minDelay))
	return uint64(r) + minDelay
}