package scenario

import (
	"context"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pingcap/test-infra/sdk/resource"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/pkg/client/pd"
	"github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/endless/testcase/oltp/scenario/betting"
	"github.com/pingcap/endless/testcase/oltp/scenario/params"
	"github.com/pingcap/endless/testcase/oltp/scenario/qhad"
	"github.com/pingcap/endless/testcase/oltp/systemtest/scale"
)

func TScale001(rootCtx context.Context, config Config, tc resource.TiDBCluster, tidbClient *tidb.Client) {
	workloadName := config.WorkloadName

	By("initializing parameters")
	var parameters *params.Scale
	if workloadName == "betting" {
		parameters = Try(betting.InitializeScale001Params(rootCtx, tidbClient)).(*params.Scale)
	} else if workloadName == "qhad" {
		parameters = Try(qhad.InitializeScale001Params(rootCtx, tidbClient)).(*params.Scale)
	} else {
	}

	Scale(rootCtx, tc, parameters)
}

// TScale002 is the body of test case `oltp_scenario_scale_002`.
func TScale002(rootCtx context.Context, config Config, tc resource.TiDBCluster, tidbClient *tidb.Client) {
	workloadName := config.WorkloadName

	By("initializing parameters")
	var parameters *params.Scale
	if workloadName == "betting" {
		parameters = Try(betting.InitializeScale002Params(rootCtx, tidbClient)).(*params.Scale)
	} else if workloadName == "qhad" {
		parameters = Try(qhad.InitializeScale002Params(rootCtx, tidbClient)).(*params.Scale)
	} else {
	}

	Scale(rootCtx, tc, parameters)
}

// Scale is the main execution flow of the test case `oltp_scenario_scale_xxx`.
func Scale(ctx context.Context, tc resource.TiDBCluster, params *params.Scale) {
	// Params must not be nil.
	Expect(params).NotTo(BeNil())

	By("setting store limit")
	pd := Try(pd.NewPDClient(tc)).(*pd.Client)
	Try(pd.SetAllStoreLimit(params.StoreLimit))
	defer pd.SetAllStoreLimit(params.DefaultStoreLimit)

	By("scaling out the cluster")
	g, gctx := errgroup.WithContext(ctx)
	oltppkg.FunCommonCheck(gctx, tc, g, params.RunTime, nil)
	g.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(gctx, "Warm Up before scan", params.WarmUpTime)
		Try(scale.ScaleClusterByNumber(gctx, tc, params.ScaleOutTargetNumber, 4*time.Hour))
	}))
	Try(g.Wait())

	By("waiting for balance")
	oltppkg.WaitForBalance(ctx, tc, params.BalanceTimeout, 1*time.Minute, params.BalanceThreshold)
}