package scale_test

import (
	"context"
	"database/sql"
	"math/rand"
	"time"

	. "github.com/onsi/ginkgo"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/testcase/tiflash/pkg"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/ch"
)

/*
扩容或缩容 tikv/tilfash 不影响查询可用性
*/

var _ = Describe("TiFlash", func() {
	var tc resource.TiDBCluster
	var chWorkload *ch.ChWorkload
	var tidbClient *tidb.Client

	BeforeEach(func() {
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		chWorkload = ch.NewWorkload(suiteTestCtx).(*ch.ChWorkload)
		chWorkload.Prepare(suiteTestCtx)
		db := matcher.Try(tc.GetDB(chWorkload.DBName())).(*sql.DB)
		tidbClient = tidb.NewClient(db)
	})
	Context("when scale in or scale out", func() {
		It("should serve continuously", func() {
			// 高并发执行 10 分钟
			componentNum := matcher.Try(tidbClient.GetAllTypeInstanceNum(suiteTestCtx)).(map[string]int)
			log.Info("cluster info", zap.Any("component num", componentNum))
			ctxr, _ := context.WithTimeout(context.Background(), 5*time.Hour)
			g, ctxr2 := errgroup.WithContext(ctxr)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(chWorkload.RunAP(ctxr2, nil))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(chWorkload.RunTP(ctxr2, pkg.ErrorHandler(
					pkg.TiDBStaleSchemaHandler,
					pkg.DeadlockFoundHandler,
					pkg.LockWaitTimeoutHandler)))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				for {
					flashReplicas := 2 + uint32(rand.Intn(componentNum["tiflash"]))*3/2
					log.Info("scale tiflash", zap.Uint32("replicas", flashReplicas))
					matcher.Try(tc.Scale(resource.TiDBClusterScaleOptions{
						Target:   resource.TiFlash,
						Replicas: flashReplicas,
					}))
					kvReplicas := 3 + uint32(rand.Intn(componentNum["tikv"]))*3/2
					log.Info("scale tikv", zap.Uint32("replicas", kvReplicas))
					matcher.Try(tc.Scale(resource.TiDBClusterScaleOptions{
						Target:   resource.TiKV,
						Replicas: kvReplicas,
					}))
					select {
					case <-ctxr2.Done():
						return
					case <-time.After(30 * time.Minute):
						log.Info("start scale")
					}
				}
			}))
			matcher.Try(g.Wait())
		})
	})
})