package pdfun

import (
	"context"
	"net/url"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/endless/pkg/chaos"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/matcher/metrics"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/endless/testcase/oltp/systemtest/pdrel"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/core"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	"golang.org/x/sync/errgroup"
	k8swait "k8s.io/apimachinery/pkg/util/wait"
)

//kill and restart br while running importing, check if pd scheduler can be resumed after one ttl
func OLTP_TOOLS_BR_REL_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
	br resource.BR,
	s3Parms oltppkg.S3Parms,
) {
	oltppkg.StopWorkload(br, "br")
	runTimeParms := oltppkg.NewRunTimePara(60*time.Minute, 5*time.Minute, 10*time.Minute, 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	defer cancelAll()
	g, ctxg := errgroup.WithContext(ctxAll)

	ResumeAllSchdulers(tc)
	defer func() {
		ResumeAllSchdulers(tc)
	}()

	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck:   true,
		IgnoreErr: GetIngoreErrors("fun"),
	}
	oltppkg.FunCommonCheck(ctxg, tc, g, runTimeParms.TotalRunTime, workloadRunParms.Cmdpr, workloadCheckParms)

	//Inject chaos
	oltppkg.LogAndBy("kill and restart br while running importing, check if pd scheduler can be resumed after one ttl")
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)

	brcmd := oltppkg.NewBr("", "restore full", s3Parms.S3AccessKeyID, s3Parms.S3SecretAccessKey)
	storageUrl := "s3://benchmark/location-6.0"
	dbName := "location"
	oltppkg.RecreateDatabase(ctxg, tc, dbName)
	brSysbenchCmd := brcmd.Pd(pdURL.String()).Storage(storageUrl).S3endpoint(s3Parms.S3Endpoint).CheckReq(false).Build()

	g.Go(WithGinkgoRecover(func() {
		interval := 5 * time.Minute
	Loop:
		for {
			select {
			case <-ctxg.Done():
				break Loop
			case <-time.After(interval):
				oltppkg.ExecWorkloadCmd(ctxg, br, brSysbenchCmd, false, nil)
				oltppkg.WaitFor(ctxg, "waiting for br import some data", runTimeParms.WarmUpTime)
				oltppkg.StopWorkload(br, "br")
				oltppkg.WaitFor(context.Background(), "waiting for check scheduler status", 6*time.Minute)
				NoPausedSchedulerCheck(tc)
			}
		}
	}))

	Expect(g.Wait()).Should(BeNil())
	log.Info("import again")
	Try(oltppkg.ExecWorkloadCmd(ctx, br, brSysbenchCmd, true, nil))
	NoPausedSchedulerCheck(tc)
	CheckIfPdConfigResetToDefault(tc)
}

//down pd leader during run br
func OLTP_TOOLS_BR_REL_002(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
	br resource.BR,
	s3Parms oltppkg.S3Parms,
) {
	oltppkg.StopWorkload(br, "br")
	runTimeParms := oltppkg.NewRunTimePara(60*time.Minute, 5*time.Minute, 10*time.Minute, 5*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	defer cancelAll()
	g, ctxg := errgroup.WithContext(ctxAll)

	ResumeAllSchdulers(tc)
	defer func() {
		ResumeAllSchdulers(tc)
	}()

	log.Info("Will inject chaos for all pod restart")
	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck:   true,
		IgnoreErr: GetIngoreErrors("rel"),
	}
	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         true,
		PDRestartCheck:           false,
		TiDBRestartCheck:         true,
	}
	oltppkg.FunCommonCheck(ctxg, tc, g, runTimeParms.TotalRunTime, workloadRunParms.Cmdpr, workloadCheckParms, restartCheckParms)

	//Inject chaos
	oltppkg.LogAndBy("Inject all pod restart")
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)

	brcmd := oltppkg.NewBr("", "restore full", s3Parms.S3AccessKeyID, s3Parms.S3SecretAccessKey)
	storageUrl := "s3://benchmark/location-6.0"
	dbName := "location"
	oltppkg.RecreateDatabase(ctxg, tc, dbName)
	brSysbenchCmd := brcmd.Pd(pdURL.String()).Storage(storageUrl).S3endpoint(s3Parms.S3Endpoint).CheckReq(false).Build()

	oltppkg.WaitFor(ctxg, "waiting for running", runTimeParms.WarmUpTime)
	g.Go(WithGinkgoRecover(func() {
		conditionWrapper := func(context.Context) (done bool, err error) {
			cmd, err := oltppkg.ExecWorkloadCmd(ctxg, br, brSysbenchCmd, false, nil)
			if cmd.Wait() == nil && err == nil {
				return true, err
			}
			return false, nil
		}
		k8swait.PollImmediateUntilWithContext(ctxg, 1*time.Minute, conditionWrapper)
	}))
	g.Go(WithGinkgoRecover(func() {
		conditionWrapper := pdrel.GetPdFaultConditionFun(ctxg, tc, "Down pd leader", true, &resource.FailureExperimentSpec{}, 2*time.Minute, 2*time.Minute)
		k8swait.PollImmediateUntilWithContext(ctxg, 1*time.Minute, conditionWrapper)
	}))

	Expect(g.Wait()).Should(BeNil())
	log.Info("import again")
	Try(oltppkg.ExecWorkloadCmd(ctx, br, brSysbenchCmd, true, nil))
	NoPausedSchedulerCheck(tc)
	CheckIfPdConfigResetToDefault(tc)
}

//kill and restart lightning while running importing, check if pd scheduler can be resumed after one ttl
func OLTP_TOOLS_LIGHTNING_REL_001(
	ctx core.TestContext,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
	s3Parms oltppkg.S3Parms,
) {
	lightning := ctx.Resource("lightning").(resource.Lightning)
	runTimeParms := oltppkg.NewRunTimePara(30*time.Minute, 10*time.Minute, 10*time.Minute, 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	defer cancelAll()
	g, ctxg := errgroup.WithContext(ctxAll)

	LightningCommonCheck(ctxg, tc, caseCmdParms, workloadRunParms, g, runTimeParms)

	//Inject chaos
	oltppkg.LogAndBy("Inject lightning down")
	g.Go(WithGinkgoRecover(func() {
		workloadRunParms.Pause(ctxg, runTimeParms.TotalRunTime)
	}))
	g.Go(WithGinkgoRecover(func() {
		interval := runTimeParms.TotalRunTime
		k8swait.PollImmediateUntilWithContext(ctxg, interval, func(context.Context) (done bool, err error) {
			LightningImport(ctxg, tc, lightning, s3Parms, false)
			oltppkg.WaitFor(ctxg, "waiting for lightning import some data", runTimeParms.WarmUpTime)
			oltppkg.StopWorkload(lightning, "lightning")
			oltppkg.WaitFor(context.Background(), "waiting for check scheduler status", 6*time.Minute)
			NoPausedSchedulerCheck(tc)
			CheckIfPdConfigResetToDefault(tc)
			return false, nil
		})
	}))
	Expect(g.Wait()).Should(BeNil())
}

//down one tikv during run lightning
func OLTP_TOOLS_LIGHTNING_REL_002(
	ctx core.TestContext,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
	s3Parms oltppkg.S3Parms,
) {
	lightning := ctx.Resource("lightning").(resource.Lightning)
	runTimeParms := oltppkg.NewRunTimePara(30*time.Minute, 10*time.Minute, 10*time.Minute, 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	defer cancelAll()
	g, ctxg := errgroup.WithContext(ctxAll)

	LightningCommonCheck(ctxg, tc, caseCmdParms, workloadRunParms, g, runTimeParms)

	//Inject chaos
	oltppkg.LogAndBy("Inject pd leader down during run lightning")

	g.Go(WithGinkgoRecover(func() {
		workloadRunParms.Pause(ctxg, runTimeParms.TotalRunTime)
	}))
	c := chaos.Chaos{
		Name:           "pd leader down",
		Tc:             tc,
		ExperimentSpec: &resource.FailureExperimentSpec{},
		SelectorOpts: resource.TiDBClusterSelectorOpts{
			TiDBClusterPDSelectorSpec: &resource.TiDBClusterPDSelectorSpec{
				IsLeader: oltppkg.ToBoolPtr(true),
			},
		},
	}
	g.Go(WithGinkgoRecover(func() {
		interval := 40 * time.Minute
		k8swait.PollImmediateUntilWithContext(ctxg, interval, func(context.Context) (done bool, err error) {
			LightningImport(ctxg, tc, lightning, s3Parms, true)
			oltppkg.WaitFor(ctxg, "waiting for lightning import some data", runTimeParms.WarmUpTime)
			oltppkg.InjectFaultByChaos(&c, ctx, runTimeParms.FaultDuration)
			oltppkg.StopWorkload(lightning, "lightning")
			oltppkg.WaitFor(context.Background(), "waiting for check scheduler status", 6*time.Minute)
			NoPausedSchedulerCheck(tc)
			CheckIfPdConfigResetToDefault(tc)
			return false, nil
		})
	}))
	Expect(g.Wait()).Should(BeNil())
}

//Region should be balance after importing via lightning
func OLTP_TOOLS_LIGHTNING_FUN_001(
	ctx core.TestContext,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
	s3Parms oltppkg.S3Parms,
) {
	lightning := ctx.Resource("lightning").(resource.Lightning)
	runTimeParms := oltppkg.NewRunTimePara(30*time.Minute, 10*time.Minute, 10*time.Minute, 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	defer cancelAll()
	g, ctxg := errgroup.WithContext(ctxAll)
	LightningCommonCheck(ctxg, tc, caseCmdParms, workloadRunParms, g, runTimeParms)

	//Inject chaos
	oltppkg.LogAndBy("start to import data via lightning")
	g.Go(WithGinkgoRecover(func() {
		workloadRunParms.Pause(ctxg, runTimeParms.TotalRunTime)
	}))
	g.Go(WithGinkgoRecover(func() {
		Try(LightningImport(ctxg, tc, lightning, s3Parms, true))
	}))
	g.Go(WithGinkgoRecover(func() {
		threshHold := 0.3
		log.Info("start score check")
		ConsistentlyWithMetrics(tc, runTimeParms.TotalRunTime).
			ShouldNot(
				SatisfyAny(
					StoreScoreDifferentExceeds(threshHold).WithCtx(ctxg).For(30*time.Second),
					StoreRegionSizeDifferentExceeds(threshHold).WithCtx(ctxg).For(30*time.Second),
				))
		log.Info("finish score check")
	}))

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

func LightningImport(ctx context.Context,
	tc resource.TiDBCluster,
	lightning resource.WorkloadNode,
	s3Parms oltppkg.S3Parms,
	syncFlag bool,
) (*workloadnode.Cmd, error) {
	defer func() {
		oltppkg.ExecWorkloadCmd(ctx, lightning, "rm -rf /tmp/tidb_lightning_checkpoint.pb", true, nil)
	}()
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	dbUrl := Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
	dbName := "location"
	storageUrl := "s3://benchmark/location-dumpling"
	lgcmd := oltppkg.NewLightning().LightningPath("").
		StorageUrl(storageUrl).TidbHost(dbUrl.Hostname()).TidbPort(dbUrl.Port()).
		S3Parms(s3Parms).SortedKvDir("/tmp/sorted-kv-dir").PdAddr(pdURL.String())
	lgRunCmd := lgcmd.Build()
	oltppkg.RecreateDatabase(ctx, tc, dbName)
	oltppkg.ExecWorkloadCmd(ctx, lightning, "rm -rf /tmp/tidb_lightning_checkpoint.pb", true, nil)
	return oltppkg.ExecWorkloadCmd(ctx, lightning, lgRunCmd, syncFlag, nil)
}
func LightningCommonCheck(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
	g *errgroup.Group,
	runTimeParms *oltppkg.RunTimePara,
) {
	ResumeAllSchdulers(tc)
	oltppkg.WaitForEmptyRegionMerge(ctx, tc, 2*time.Hour, 1*time.Minute, 100, 0.1)

	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck:   true,
		IgnoreErr: GetIngoreErrors("fun"),
	}
	qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
		IsCheck: false,
	}
	durJitterCheckParms := oltppkg.DurJitterCheckParameters{
		IsCheck: false,
	}
	oltppkg.FunCommonCheck(ctx, tc, g, runTimeParms.TotalRunTime, workloadRunParms.Cmdpr, workloadCheckParms,
		qpsJitterCheckParms, durJitterCheckParms)
}