package scale

import (
	"context"
	"database/sql"
	"io"
	"net/url"
	"time"

	. "github.com/onsi/gomega"
	pdclient "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/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

func OLTP_SCALE_OUT_FUN_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	pr *io.PipeReader,
	caseCmdParms oltppkg.CaseCmdParms,
) {
	oltppkg.LogAndBy("Set store limit.")
	pd := Try(pdclient.NewPDClient(tc)).(*pdclient.Client)
	_, sumSize, err := oltppkg.GetStoreSizeByType(tc, "used_size")
	Expect(err).Should(BeNil())
	var balanceTime time.Duration
	if sumSize > 2*1024*1024 {
		balanceTime = 5 * time.Hour
		Try(pd.SetAllStoreLimit(200))
	} else {
		balanceTime = 3 * time.Hour
		Try(pd.SetAllStoreLimit(60))
	}
	defer pd.SetAllStoreLimit(15)

	dbUrl := Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
	db, err := sql.Open("mysql", resource.BuildDSN(dbUrl, "INFORMATION_SCHEMA"))
	Expect(err).Should(Succeed())
	defer db.Close()
	tidbClient := tidb.NewClient(db)

	initTidbNum := Try(tidbClient.GetInstanceNum("tidb")).(uint32)
	initPdNum := Try(tidbClient.GetInstanceNum("pd")).(uint32)
	initTikvNum := Try(tidbClient.GetInstanceNum("tikv")).(uint32)

	log.Info("initial number of instance", zap.Uint32("tidb-num", initTidbNum), zap.Uint32("pd-num", initPdNum), zap.Uint32("tikv-num", initTikvNum))

	targetNumber := map[string]uint32{
		"tikv": initTikvNum + 1,
		"tidb": initTidbNum + 1,
		"pd":   initPdNum + 1}

	oltppkg.LogAndBy("scale out TiDB, PD and TiKV")
	var runTime time.Duration = 60 * time.Minute
	ctxAll, cancel := context.WithTimeout(ctx, runTime)
	defer cancel()
	g, ctxg := errgroup.WithContext(ctxAll)

	oltppkg.FunCommonCheck(ctxg, tc, g, runTime, pr)

	g.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(ctxg, "scale out after 2min", 2*time.Minute)
		ScaleClusterByNumber(ctxg, tc, targetNumber, 4*time.Hour)
	}))

	Expect(g.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Waiting for balance.")
	oltppkg.WaitForBalance(ctx, tc, balanceTime, 5*time.Minute, 0.2)
}