package scale

import (
	"context"
	"fmt"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/endless/pkg/chaos"
	. "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"
	"golang.org/x/sync/errgroup"
	k8swait "k8s.io/apimachinery/pkg/util/wait"
)

func OLTP_SCALE_STABILITY_FUN_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimeParms := oltppkg.NewRunTimePara(12*time.Hour, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 20*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	g0, ctxg0 := errgroup.WithContext(ctxAll)
	scaleOutPercent := 2
	defer cancelAll()

	RestoreCluster(ctx, tc, 2*time.Hour)
	defer func() {
		RestoreCluster(ctx, tc, 2*time.Hour)
		Try(oltppkg.ExecPdCtl(tc, "store limit all 15"))
	}()

	InitForScale(ctxg0, g0, tc, workloadRunParms, caseCmdParms, runTimeParms, GetScaleWorkloadCheckParms(caseCmdParms, true))

	oltppkg.LogAndBy("Step 2: Scale out and scale in the cluster")
	g1, ctxg1 := errgroup.WithContext(ctxg0)
	g1.Go(WithGinkgoRecover(func() {
		ScaleOutInAndCheck(ctxg1, tc, *runTimeParms, float64(scaleOutPercent))
		log.Info("scale in out finished")
	}))

	Expect(g1.Wait()).Should(BeNil())
	cancelAll()
	Expect(g0.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 3: Waiting for workload.")
	oltppkg.WaitForBalance(ctx, tc, 10*time.Minute, 2*time.Minute, 0.05)
	log.Info("test finished")
}

func OLTP_SCALE_STABILITY_REL_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimeParms := oltppkg.NewRunTimePara(12*time.Hour, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 20*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	g0, ctxg0 := errgroup.WithContext(ctxAll)
	defer cancelAll()

	scaleOutPercent := 2
	RestoreCluster(ctx, tc, 2*time.Hour)
	defer func() {
		RestoreCluster(ctx, tc, 2*time.Hour)
		Try(oltppkg.ExecPdCtl(tc, "store limit all 15"))
	}()

	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         false,
		PDRestartCheck:           true,
		TiDBRestartCheck:         true,
	}
	InitForScale(ctxg0, g0, tc, workloadRunParms, caseCmdParms, runTimeParms, GetScaleWorkloadCheckParms(caseCmdParms, false), restartCheckParms)

	oltppkg.LogAndBy("Step 2: Scale-out-Scale-in the cluster")
	startTime := time.Now().Add(runTimeParms.WarmUpTime / 2)
	g1, ctxg1 := errgroup.WithContext(ctxg0)
	canQuit := make(chan bool, 1)
	g1.Go(WithGinkgoRecover(func() {
		ScaleOutInAndCheck(ctxg1, tc, *runTimeParms, float64(scaleOutPercent))
		canQuit <- false
	}))

	g1.Go(WithGinkgoRecover(func() {
		runInterval := 30 * time.Minute
		k8swait.PollImmediateUntilWithContext(ctxg1, runInterval, func(ctx context.Context) (bool, error) {
			select {
			case <-canQuit:
				return true, nil
			default:
				initTikvNum := GetComponentNum(tc)["tikv"]
				c := Try(oltppkg.InitChaos(tc,
					resource.TiKV, "down one tikv",
					&resource.FailureExperimentSpec{},
					oltppkg.ToIntPtr(oltppkg.GetRandInt(0, int(initTikvNum))))).(chaos.Chaos)
				faultDuration := oltppkg.GetRandDuration(time.Minute, 10)
				oltppkg.InjectFaultByChaos(&c, ctxg1, faultDuration)
				return false, nil
			}
		})
	}))

	Expect(g1.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 3: Check qps and latency with that before do scale in.")
	oltppkg.WaitFor(ctxg0, "waiting for recovery", runTimeParms.RecoveryTime)
	oltppkg.RepeatQpsDurCheckWithAbsTime(ctxg0, 10*time.Minute, startTime, tc, 0.8)
	cancelAll()
	Expect(g0.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 4: Waiting for balance.")
	oltppkg.WaitForBalance(ctx, tc, 10*time.Minute, 2*time.Minute, 0.05)
}

func OLTP_SCALE_STABILITY_REL_002(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimeParms := oltppkg.NewRunTimePara(12*time.Hour, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 20*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	g0, ctxg0 := errgroup.WithContext(ctxAll)
	defer cancelAll()

	scaleOutPercent := 2
	RestoreCluster(ctx, tc, 2*time.Hour)
	defer func() {
		RestoreCluster(ctx, tc, 2*time.Hour)
		Try(oltppkg.ExecPdCtl(tc, "store limit all 15"))
	}()

	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         false,
		PDRestartCheck:           true,
		TiDBRestartCheck:         true,
	}
	InitForScale(ctxg0, g0, tc, workloadRunParms, caseCmdParms, runTimeParms, GetScaleWorkloadCheckParms(caseCmdParms, false), restartCheckParms)

	oltppkg.LogAndBy("Step 2: Scale out and scale in the cluster")
	startTime := time.Now().Add(runTimeParms.WarmUpTime / 2)
	g1, ctxg1 := errgroup.WithContext(ctxg0)
	canQuit := make(chan bool, 1)
	g1.Go(WithGinkgoRecover(func() {
		ScaleOutInAndCheck(ctxg1, tc, *runTimeParms, float64(scaleOutPercent))
		canQuit <- false
	}))

	g1.Go(WithGinkgoRecover(func() {
		runInterval := 30 * time.Minute
		k8swait.PollImmediateUntilWithContext(ctxg1, runInterval, func(ctx context.Context) (bool, error) {
			select {
			case <-canQuit:
				return true, nil
			default:
				c := chaos.Chaos{
					Name:           "one zone down",
					Tc:             tc,
					ExperimentSpec: &resource.FailureExperimentSpec{},
					SelectorOpts:   chaos.SelectAllTiKVInOneAZ(fmt.Sprintf("zone%d", oltppkg.GetRandInt(1, 4))),
				}
				faultDuration := oltppkg.GetRandDuration(time.Minute, 5)
				oltppkg.InjectFaultByChaos(&c, ctxg1, faultDuration)
				return false, nil
			}
		})
	}))

	Expect(g1.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 3: Check qps and latency with that before do scale in.")
	oltppkg.WaitFor(ctxg0, "waiting for recovery", runTimeParms.RecoveryTime)
	oltppkg.RepeatQpsDurCheckWithAbsTime(ctxg0, 10*time.Minute, startTime, tc, 0.8)
	cancelAll()
	Expect(g0.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 4: Waiting for balance.")
	oltppkg.WaitForBalance(ctx, tc, 10*time.Minute, 2*time.Minute, 0.05)
}

func OLTP_SCALE_STABILITY_REL_003(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimeParms := oltppkg.NewRunTimePara(12*time.Hour, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 20*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimeParms.TotalRunTime)
	g0, ctxg0 := errgroup.WithContext(ctxAll)
	defer cancelAll()

	scaleOutPercent := 2
	RestoreCluster(ctx, tc, 2*time.Hour)
	defer func() {
		RestoreCluster(ctx, tc, 2*time.Hour)
		Try(oltppkg.ExecPdCtl(tc, "store limit all 15"))
		Try(oltppkg.ExecPdCtl(tc, "store remove-tombstone"))
	}()

	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         true,
		PDRestartCheck:           false,
		TiDBRestartCheck:         true,
	}
	InitForScale(ctxg0, g0, tc, workloadRunParms, caseCmdParms, runTimeParms, GetScaleWorkloadCheckParms(caseCmdParms, false), restartCheckParms)

	oltppkg.LogAndBy("Step 2: Scale out and scale in the cluster")
	startTime := time.Now().Add(runTimeParms.WarmUpTime / 2)
	g1, ctxg1 := errgroup.WithContext(ctxg0)
	canQuit := make(chan bool, 1)
	g1.Go(WithGinkgoRecover(func() {
		ScaleOutInAndCheck(ctxg1, tc, *runTimeParms, float64(scaleOutPercent))
		canQuit <- false
	}))

	g1.Go(WithGinkgoRecover(func() {
		runInterval := 10 * time.Minute
		oltppkg.WaitFor(ctxg1, "waiting for workload running", runTimeParms.WarmUpTime+5*time.Minute)
		k8swait.PollImmediateUntilWithContext(ctxg1, runInterval, func(ctx context.Context) (bool, error) {
			select {
			case <-canQuit:
				log.Info("Got quit signal,scale-in-out finished")
				return true, nil
			default:
				c := Try(oltppkg.InitChaos(tc, resource.PD, "down pd leader", &resource.FailureExperimentSpec{},
					oltppkg.ToBoolPtr(true))).(chaos.Chaos)
				faultDuration := oltppkg.GetRandDuration(time.Minute, 2)
				oltppkg.InjectFaultByChaos(&c, ctxg1, faultDuration)
				return false, nil
			}
		})
	}))

	Expect(g1.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 3: Check qps and latency with that before do scale in.")
	oltppkg.WaitFor(ctxg0, "waiting for recovery", runTimeParms.RecoveryTime)
	oltppkg.RepeatQpsDurCheckWithAbsTime(ctxg0, 10*time.Minute, startTime, tc, 0.8)
	cancelAll()
	Expect(g0.Wait()).Should(BeNil())

	oltppkg.LogAndBy("Step 4: Waiting for balance.")
	oltppkg.WaitForBalance(ctx, tc, 10*time.Minute, 2*time.Minute, 0.05)
}

func RestoreCluster(ctx context.Context, tc resource.TiDBCluster, timeOut time.Duration) {
	Try(oltppkg.ExecPdCtl(tc, "store limit all 200"))
	log.Info(fmt.Sprintf("current component number %v", GetComponentNum(tc)))
	targetNumber := map[string]uint32{
		"tidb": 2,
		"pd":   3,
		"tikv": 3,
	}
	log.Info(fmt.Sprintf("target number %v", targetNumber))
	Try(ScaleClusterByNumber(ctx, tc, targetNumber, timeOut))
	Try(oltppkg.ExecPdCtl(tc, "store remove-tombstone"))
	oltppkg.WaitForBalance(ctx, tc, 3*time.Hour, 2*time.Minute, 0.1)
	Try(oltppkg.ExecPdCtl(tc, "store remove-tombstone"))
}