package pdfun

import (
	"context"
	"database/sql"
	"io"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	. "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_TSO_SCALIBILITY_FUN_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	workloadType string,
	cmdpr *io.PipeReader,
) {
	tidbVersion := oltppkg.GetCompVersion(ctx, tc, "tidb")
	versionShort := strings.TrimSuffix(strings.TrimPrefix(tidbVersion, "v"), "-pre")
	if oltppkg.CompareVersion(versionShort, "5.3.0") < 0 {
		log.Info("Skip this case for current version is lower than 5.3.0", zap.String("tidbVersion", versionShort))
		return
	}
	runTime := 1 * time.Hour
	ctxAll, cancelAll := context.WithTimeout(ctx, 120*time.Minute)
	//Start oom and output check before all.
	oltppkg.LogAndBy("Step 1:Start oom and output check before all.")
	g1, ctxg1 := errgroup.WithContext(ctxAll)

	oltppkg.FunCommonCheck(ctx, tc, g1, runTime, cmdpr)

	oltppkg.LogAndBy("Enable tso follower")
	g2, ctxg2 := errgroup.WithContext(ctxg1)
	startTime := time.Now().Add(5 * time.Minute)
	oltppkg.WaitFor(ctxg2, "warm up workload", 10*time.Minute)
	db := Try(tc.GetDB("")).(*sql.DB)
	defer db.Close()
	log.Info("set global tidb_enable_tso_follower_proxy = 1")
	Try(db.Exec("set global tidb_enable_tso_follower_proxy = 1"))
	g2.Go(WithGinkgoRecover(func() {
		checkTime := 20 * time.Minute
		oltppkg.RepeatQpsDurCheckWithAbsTime(ctxg2, checkTime, startTime, tc, 0.8)
	}))
	Expect(g2.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Enable client batch")
	g3, ctxg3 := errgroup.WithContext(ctxg2)
	startTime = time.Now().Add(5 * time.Minute)
	oltppkg.WaitFor(ctxg3, "warm up workload", 10*time.Minute)
	log.Info("set global tidb_tso_client_batch_max_wait_time = 1")
	Try(db.Exec("set global tidb_tso_client_batch_max_wait_time = 1"))
	defer func() {
		Try(db.Exec("set global tidb_tso_client_batch_max_wait_time = 0"))
	}()
	g3.Go(WithGinkgoRecover(func() {
		checkTime := 20 * time.Minute
		oltppkg.RepeatQpsDurCheckWithAbsTime(ctxg3, checkTime, startTime, tc, 0.8)
	}))
	Expect(g3.Wait()).Should(BeNil())

	//Finish test
	cancelAll()
	Expect(g1.Wait()).Should(BeNil())

}