package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/random"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/client/v3/concurrency"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

/*
GOOS=linux GOARCH=arm64 go build -v -o etcd-qps-arm64
docker buildx build --platform=linux/arm64 --push -t etcd-qps .

cat <<EOF | kubectl apply -f -
apiVersion: apps/v1
kind: Deployment
metadata:

	labels:
	  app: etcd-qps
	name: etcd-qps

spec:

	replicas: 1 # you can run multi replicas.
	selector:
	  matchLabels:
	    app: etcd-qps
	template:
	  metadata:
	    labels:
	      app: etcd-qps
	  spec:
	    containers:
	    - image: etcd-qps
	      imagePullPolicy: Always
	      name: etcd-qps
	      args:
	      - --endpoints=etcd-svc.tidb-admin:2379
	      - --conns=10
	      command:
	      - /etcd-qps

EOF
*/
var total int64
var ctx context.Context
var cancel context.CancelFunc
var wg sync.WaitGroup
var errch = make(chan error)

var etcdClient *clientv3.Client
var etcdTokenPrefix, etcdLockPrefix = *etcdPrefix + "token/", *etcdPrefix + "lock/"

var (
	//nolint:mnd // no need
	fakeToken = fmt.Sprintf("fake-token-%s", random.Random(4))
	//nolint:mnd // no need
	fakeEndpoints = fmt.Sprintf("fake-podname1,127.0.0.1;fake-random-%s,127.0.0.2", random.Random(4))
	fakeTokenKey  = etcdTokenPrefix + fakeToken + "/" + utils.KeyEndpoints
)

var (
	etcdEndpoints = flag.String("endpoints", "127.0.0.1:2379", "the etcd endpoints")
	etcdPrefix    = flag.String("etcd-prefix", "/etcd-qps/benchmark/", "the etcd prefix")
	totalConns    = flag.Uint("conns", 1, "the etcd total connections")
)

func tokenLocker(token string) (sync.Locker, error) {
	s, err := concurrency.NewSession(etcdClient)
	if err != nil {
		return nil, err
	}
	return concurrency.NewLocker(s, etcdLockPrefix+token), nil
}

func assignWithLock(token string) {
	var err error
	defer func() {
		if err != nil {
			select {
			case errch <- err:
			default:
			}
		}
	}()

	// Step 1: Simulate the locker
	locker, err := tokenLocker(token)
	if err != nil {
		return
	}
	locker.Lock()
	defer locker.Unlock()

	// Step 2: Simulate Get token(fake)
	_, err = etcdClient.Get(ctx, fakeTokenKey)
	if err != nil {
		return
	}

	// Step 3: Simulate Txn token
	key := etcdTokenPrefix + token + "/" + utils.KeyEndpoints
	endpoints := fmt.Sprintf("%s,127.0.0.1;%s,127.0.0.2", "podname1-"+token, "podname2-"+token)
	resp, err := etcdClient.Txn(ctx).
		If(clientv3.Compare(clientv3.Value(fakeTokenKey), "=", fakeEndpoints)).
		Then(clientv3.OpPut(key, endpoints)).
		Else(clientv3.OpGet(fakeTokenKey)).
		Commit()
	if err != nil {
		return
	}
	if !resp.Succeeded {
		err = errors.New("resp.Succeeded is false")
		return
	}

	atomic.AddInt64(&total, 1)
}

func preprocessing() {
	if _, err := etcdClient.KV.Put(ctx, fakeTokenKey, fakeEndpoints); err != nil {
		log.Fatal("Put fake token", err)
	}
}

func benchmark() {
	for i := 0; i < int(*totalConns); i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				//nolint:mnd // no need
				token := random.Random(15)
				assignWithLock(token)
				if ctx.Err() != nil {
					break
				}
			}
		}()
	}
}

func main() {
	var err error
	flag.Parse()

	if *totalConns < 1 {
		log.Fatal("--conns can't less than 1")
	}

	etcdClient, err = clientv3.New(clientv3.Config{
		Endpoints:   strings.Split(*etcdEndpoints, ";"),
		DialTimeout: 2 * time.Second, //nolint:mnd // no need
		DialOptions: []grpc.DialOption{grpc.WithBlock()},
		Logger:      zap.NewNop(),
	})
	if err != nil {
		log.Fatal("New etcd client", err)
	}

	defer func() {
		wg.Wait()

		_, err = etcdClient.KV.Delete(context.Background(), *etcdPrefix, clientv3.WithPrefix())
		if err != nil {
			log.Printf("clean key %s err: %v", *etcdPrefix, err)
		}

		etcdClient.Close() //nolint:errcheck // no need check error
	}()

	ctx, cancel = context.WithCancel(context.Background())
	defer cancel()

	preprocessing()

	start := time.Now()
	log.Printf("start benchmark(totalConns:%d) now %v ...", *totalConns, start)
	benchmark()

	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	//nolint: gomnd // not need
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()
	preTotal := int64(0)
LOOP:
	for {
		select {
		case <-ticker.C:
			nowTotal := atomic.LoadInt64(&total)
			//nolint: gomnd // not need
			log.Printf("etcd qps: %.2f", float64(nowTotal-preTotal)/float64(10))
			preTotal = nowTotal
		case err = <-errch:
			log.Printf("exit err: %v", err)
			break LOOP
		case sig := <-sigch:
			log.Printf("exit sig: %v", sig)
			break LOOP
		}
	}

	nowTotal := atomic.LoadInt64(&total)
	log.Printf("etcd avg qps: %.2f", float64(nowTotal)/time.Since(start).Seconds())
}