package upgrade

import (
	"context"
	"fmt"
	"net/url"
	"strings"

	"time"

	. "github.com/onsi/gomega"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/matcher/metrics"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/endless/testcase/oltp/systemtest/scale"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

type UpgradeOptions struct {
	Tc                           resource.TiDBCluster
	CaseCmdParms                 oltppkg.CaseCmdParms
	WorkloadRunParms             oltppkg.WorkloadRunParms
	RunTimeParms                 oltppkg.RunTimePara
	TestType                     string
	RestartCheckTimeAfterUpgrade time.Duration
	SlaCheckTime                 time.Duration
	SlaCheckThreshhold           float64
	BenchToolSet                 resource.WorkloadNode
}

func (opt *UpgradeOptions) UpgradeTest(ctx context.Context) {

	if opt.WorkloadRunParms.Cmdpr == nil {
		oltppkg.LogAndBy("Start workload")
		opt.WorkloadRunParms = oltppkg.WorkloadRunParms{}
		opt.WorkloadRunParms.InitAndStart(ctx, opt.Tc, opt.BenchToolSet, opt.CaseCmdParms, false)
	}

	upgradeConfig := opt.CaseCmdParms.UpgradeParms
	log.Info(fmt.Sprintf("upgrade config is %v", upgradeConfig))
	ctxAll, cancelAll := context.WithTimeout(ctx, opt.RunTimeParms.TotalRunTime)

	g1, ctxg1 := errgroup.WithContext(ctxAll)
	workloadWarmUpTime := 10 * time.Minute
	timeStart := time.Now().Add(workloadWarmUpTime / 2)
	opt.CommonCheckBeforeUpgrade(ctxg1, g1)

	oltppkg.WaitFor(ctxAll, "waiting for workload warmup", workloadWarmUpTime)

	ctx2, cancel2 := context.WithTimeout(ctx, opt.RunTimeParms.WarmUpTime)
	defer cancel2()
	g2, ctxg2 := errgroup.WithContext(ctx2)
	Try(oltppkg.ExecPdCtl(opt.Tc, "store limit all 200"))
	defer func() {
		Try(oltppkg.ExecPdCtl(opt.Tc, "store limit all 15"))
	}()

	curComponentNum := scale.GetComponentNum(opt.Tc)
	switch opt.TestType {
	case "tombstone":
		targetNumber := map[string]uint32{
			"tikv": curComponentNum["tikv"] - 1}
		Try(scale.ScaleClusterByNumber(ctxg2, opt.Tc, targetNumber, opt.RunTimeParms.TotalRunTime))
		opt.SlaCheckThreshhold = 0.5
	case "during-scale-in":
		targetNumber := map[string]uint32{
			"tikv": curComponentNum["tikv"] - 1}
		g2.Go(WithGinkgoRecover(func() {
			Try(scale.ScaleClusterByNumber(ctxg2, opt.Tc, targetNumber, opt.RunTimeParms.TotalRunTime))
		}))
		opt.SlaCheckThreshhold = 0.5
	case "during-scale-out":
		targetNumber := map[string]uint32{
			"tikv": curComponentNum["tikv"] + 1}
		g2.Go(WithGinkgoRecover(func() {
			Try(scale.ScaleClusterByNumber(ctxg2, opt.Tc, targetNumber, opt.RunTimeParms.TotalRunTime))
		}))
		opt.SlaCheckThreshhold = 1.1
	default:
	}

	for _, comtype := range []string{"tikv", "tidb", "pd"} {
		upgradeFinishFlag := make(chan bool, 1)
		ctx3, cancel3 := context.WithCancel(ctxg2)
		defer cancel3()
		threshHold, ifCheck := GetUpgradeSlaThreshhold(comtype, opt.CaseCmdParms.WorkloadType.Type[0])
		g2.Go(WithGinkgoRecover(func() {
			Try(UpgradeComponent(ctx3, opt.Tc, comtype, opt.RunTimeParms.TotalRunTime, upgradeConfig.TargetVersion))
			upgradeFinishFlag <- true
			cancel3()
		}))
		if opt.TestType == "nothing" && ifCheck {
			g2.Go(WithGinkgoRecover(func() {
				log.Info("qps check during upgrade started")
				oltppkg.RepeatQpsDurCheckWithAbsTime(ctx3, opt.RunTimeParms.TotalRunTime, timeStart, opt.Tc, threshHold)
				log.Info("qps check during upgrade finished")
			}))
		}
		var timeOut time.Duration
		if comtype == "tikv" {
			timeOut = 60 * time.Minute
		} else {
			timeOut = 30 * time.Minute
		}
		Eventually(upgradeFinishFlag, timeOut, 1*time.Minute).Should(Receive())
		opt.RestartCheckAfterUpgrade(ctxg2, upgradeFinishFlag, g2, comtype)
		if comtype == "tidb" && strings.HasPrefix(opt.CaseCmdParms.WorkloadType.Type[0], "sysbench") {
			//restart workload for sysbench after tidb upgrade
			opt.WorkloadRunParms.Stop()
			opt.WorkloadRunParms.Start(ctx, opt.Tc, opt.CaseCmdParms, false)
		}
		oltppkg.WaitFor(ctxg2, "waiting for next component upgrade", 5*time.Minute)
		cancel2()
	}

	Expect(g2.Wait()).Should(BeNil())

	opt.WorkloadRunParms.Stop()
	opt.WorkloadRunParms.Start(ctx, opt.Tc, opt.CaseCmdParms, false)

	oltppkg.WaitForBalance(ctxg1, opt.Tc, opt.RunTimeParms.RecoveryTime, 1*time.Minute, 0.1)
	opt.SlaCheckWithAfterUpgrade(ctxAll, timeStart)

	cancelAll()
	Expect(g1.Wait()).Should(BeNil())
	log.Info("Finish the test, stop the workload")
	opt.WorkloadRunParms.Stop()
}
func GetUpgradeSlaThreshhold(compType string, workloadType string) (float64, bool) {
	if compType == "tikv" {
		return 0.5, true
	} else if compType == "tidb" && strings.HasPrefix(workloadType, "sysbench") {
		return 0, false
	} else if compType == "tidb" {
		return 0.5, true
	} else {
		return 0.3, true
	}
}
func (opt *UpgradeOptions) CommonCheckBeforeUpgrade(ctx context.Context,
	g *errgroup.Group,
) {
	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck: true,
		//9001: PD server timeout
		//8027: Information schema is out of date: schema failed to update in 1 lease
		//9005: ingore 9005 during upgrade for https://internal.pingcap.net/jira/projects/INCALL/issues/INCALL-127
		IgnoreErr: append(strings.Split(opt.CaseCmdParms.TpccParms.IgnoreErrors, ","), "9005", "8027", "9001", "9002", "invalid connection", "err dial tcp"),
	}
	qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
		IsCheck: false,
	}
	durJitterCheckParms := oltppkg.DurJitterCheckParameters{
		IsCheck: false,
	}
	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         false,
		PDRestartCheck:           false,
		TiDBRestartCheck:         false,
	}
	//Skip applyDbMutexLockDurCheckParams for there is issue: https://internal.pingcap.net/jira/browse/INCALL-151
	applyDbMutexLockDurCheckParams := oltppkg.ApplyDbMutexLockDurCheckParameters{
		IsCheck: false,
	}
	oltppkg.FunCommonCheck(ctx, opt.Tc, g, opt.RunTimeParms.TotalRunTime, opt.WorkloadRunParms.Cmdpr,
		workloadCheckParms, qpsJitterCheckParms, durJitterCheckParms, restartCheckParms, applyDbMutexLockDurCheckParams)
	Try(oltppkg.ExecPdCtl(opt.Tc, "store remove-tombstone"))
}
func (opt *UpgradeOptions) RestartCheckAfterUpgrade(ctx context.Context, stopChan chan bool, g *errgroup.Group, compType string) {
	g.Go(WithGinkgoRecover(func() {
		select {
		case <-ctx.Done():
			break
		case <-stopChan:
			log.Info("Check there is no restart after do upgrade.")
			oltppkg.WaitFor(context.Background(), "waiting to check restart after upgrade finished.", 2*time.Minute)
			//time.Sleep(3 * time.Minute)
			log.Info("begin to check after upgrade")
			switch compType {
			case "tikv":
				ConsistentlyWithMetrics(opt.Tc, opt.RestartCheckTimeAfterUpgrade).
					ShouldNot(metrics.TiKVRestart().WithCtx(ctx))
			case "tidb":
				ConsistentlyWithMetrics(opt.Tc, opt.RestartCheckTimeAfterUpgrade).
					ShouldNot(metrics.TiDBRestart().WithCtx(ctx))
			case "pd":
				ConsistentlyWithMetrics(opt.Tc, opt.RestartCheckTimeAfterUpgrade).
					ShouldNot(metrics.PDRestart().WithCtx(ctx))
			default:
			}

		}
	}))
}
func (opt *UpgradeOptions) CommonCheckAfterUpgrade(ctx context.Context) {
	upgradeConfig := opt.CaseCmdParms.UpgradeParms
	oltppkg.CheckClusterVersion(ctx, opt.Tc, upgradeConfig.TargetVersion,
		map[string]string{
			"pd":   upgradeConfig.MasterPdVersion,
			"tikv": upgradeConfig.MasterTikvVersion,
			"tidb": upgradeConfig.TargetTidbImage},
	)
	promUrl := Try(opt.Tc.ServiceURL(resource.Prometheus)).(*url.URL)
	if opt.TestType == "tombstone" || opt.TestType == "during-scale-in" {
		Consistently(scale.GetTombstoneNum(promUrl), 5*time.Minute, 1*time.Minute).Should(BeNumerically("==", float64(1)))
		Try(oltppkg.ExecPdCtl(opt.Tc, "store remove-tombstone"))
	}
	if opt.TestType == "during-scale-out" {
		Consistently(scale.GetPreparingNum(promUrl), 5*time.Minute, 1*time.Minute).Should(BeNumerically(">=", float64(0)))
	}
}
func (opt *UpgradeOptions) SlaCheckWithAfterUpgrade(ctx context.Context, timeStart time.Time) {

	workloadType := opt.CaseCmdParms.WorkloadType.Type[0]
	runTime3 := opt.SlaCheckTime
	ctx3, cancel3 := context.WithTimeout(ctx, runTime3)
	defer cancel3()
	g3, ctxg3 := errgroup.WithContext(ctx3)
	//Check qps drop with the qps before upgrade
	g3.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(ctxg3, "Wait for workload balance", 20*time.Minute)
		oltppkg.RepeatQpsDurCheckWithAbsTime(ctxg3, 30*time.Minute, timeStart, opt.Tc, opt.SlaCheckThreshhold)
	}))
	g3.Go(WithGinkgoRecover(func() {
		if strings.HasPrefix(workloadType, "sysbench") {
			oltppkg.GetTableAnalyzeInfo(ctxg3, opt.Tc, opt.CaseCmdParms.SysbenchParms.Db, "sbtest11")
		}
	}))
	oltppkg.LogAndBy("Wait for workload balance on all tikvs.")
	oltppkg.WaitForWorkloadBalance(ctxg3, opt.Tc, workloadType, g3, 2*time.Minute, 20*time.Minute, 0.8)
	Expect(g3.Wait()).Should(BeNil())
	log.Info("SlaCheckWithAfterUpgrade finished.")
}

//UpgradeClusterByVersion Only support when baseImage is set in testbed.
func UpgradeClusterByVersion(ctx context.Context, tc resource.TiDBCluster, targetVersion string, masterMap map[string]string, timeout time.Duration, workloadRunParms oltppkg.WorkloadRunParms) error {
	for _, comtype := range []string{"tikv", "tidb", "pd"} {
		target := oltppkg.GetResourceTypeByString(comtype)
		targetVersionShort := strings.TrimPrefix(targetVersion, "v")
		if strings.Contains(targetVersionShort, "-") {
			targetVersionShort = strings.Split(targetVersionShort, "-")[0]
		}
		if targetVersionShort == "master" {
			targetVersionShort = masterMap[comtype]
		}
		curVersion := oltppkg.GetCompVersion(ctx, tc, comtype)
		if strings.Contains(curVersion, targetVersionShort) {
			log.Info("Seems no need to upgrade.", zap.String("comtype", comtype), zap.String("targetVersion", targetVersionShort), zap.String("curVersion", curVersion))
			continue
		}
		log.Info("Now begin to upgrade ", zap.String("comptype", comtype), zap.String("targetVersion", targetVersion))
		err := tc.Upgrade(
			resource.TiDBClusterUpgradeOptions{
				ComponentUpgradeOption: &resource.TiDBClusterComponentUpgradeOption{
					Target: target,
					//Image:  image,
					Version: targetVersion,
				},
				Timeout: timeout})
		if err != nil {
			return err
		}
		log.Info("Upgrade succeeded.", zap.String("comptype", comtype))
		oltppkg.WaitFor(ctx, "waiting for next component upgrade", 10*time.Minute)
	}
	return nil
}

func UpgradeCluster(ctx context.Context, tc resource.TiDBCluster, timeout time.Duration, upgradeConfig oltppkg.UpgradeParms, workloadRunParms oltppkg.WorkloadRunParms) error {
	masterMap := map[string]string{"pd": upgradeConfig.MasterPdVersion,
		"tikv": upgradeConfig.MasterTikvVersion,
		"tidb": upgradeConfig.MasterTidbVersion,
	}
	if upgradeConfig.UpgradeByVersion {
		return UpgradeClusterByVersion(ctx, tc, upgradeConfig.TargetVersion, masterMap, timeout, workloadRunParms)
	} else {
		targetImage := map[string]string{"tidb": upgradeConfig.TargetTidbImage, "tikv": upgradeConfig.TargetTikvImage, "pd": upgradeConfig.TargetPdImage}
		return UpgradeClusterByImage(ctx, tc, targetImage, masterMap, timeout, workloadRunParms)
	}
}

func UpgradeComponent(ctx context.Context, tc resource.TiDBCluster, comtype string, timeout time.Duration, targetVersion string) error {
	target := oltppkg.GetResourceTypeByString(comtype)
	log.Info("Now begin to upgrade ", zap.String("comptype", comtype), zap.String("targetVersion", targetVersion))
	err := tc.Upgrade(
		resource.TiDBClusterUpgradeOptions{
			ComponentUpgradeOption: &resource.TiDBClusterComponentUpgradeOption{
				Target:  target,
				Version: targetVersion,
			},
			Timeout: timeout})
	if err != nil {
		return err
	}
	log.Info("Upgrade succeeded.", zap.String("comptype", comtype))
	return nil
}

//UpgradeClusterByImage only support when baseImage is not set in testbed, cannot use it if use v0.0.0-20220221083224.
func UpgradeClusterByImage(ctx context.Context, tc resource.TiDBCluster, targetImage map[string]string, masterMap map[string]string, timeout time.Duration, workloadRunParms oltppkg.WorkloadRunParms) error {
	for _, comtype := range []string{"tikv", "tidb", "pd"} {
		if targetImage[comtype] == "" {
			log.Info("No image provided, no need to upgrade.")
			continue
		}
		log.Info(fmt.Sprintf(" %v", targetImage))
		target := oltppkg.GetResourceTypeByString(comtype)
		image := targetImage[comtype]
		log.Info(image)
		targetVersion := strings.Split(image, ":")[1]
		log.Info(targetVersion)
		targetVersion = strings.TrimPrefix(targetVersion, "v")
		if strings.Contains(targetVersion, "-") {
			targetVersion = strings.Split(targetVersion, "-")[0]
		}
		if targetVersion == "master" {
			targetVersion = masterMap[comtype]
		}
		curVersion := oltppkg.GetCompVersion(ctx, tc, comtype)
		if strings.Contains(curVersion, targetVersion) {
			log.Info("Seems no need to upgrade.", zap.String("targetVersion", targetVersion), zap.String("curVersion", curVersion))
			continue
		}
		log.Info("Now begin to upgrade ", zap.String("comptype", comtype), zap.String("targetVersion", targetVersion))
		err := tc.Upgrade(
			resource.TiDBClusterUpgradeOptions{
				ComponentUpgradeOption: &resource.TiDBClusterComponentUpgradeOption{
					Target: target,
					//Image:  image,  remove it , since test-infra didn't support upgrade via image any more
				},
				Timeout: timeout})
		if err != nil {
			return err
		}
		log.Info("Upgrade succeeded.", zap.String("comptype", comtype))
		oltppkg.WaitFor(ctx, "Waiting for information update to w/a bug#https://github.com/pingcap/tidb/issues/30907.", 5*time.Minute)
		//Restart workload after upgrade tidb succeeded.
		workloadRunParms.Pause(ctx, 1*time.Minute)
		if strings.Contains(image, ":") {
			targetVersion := strings.Split(image, ":")[1]
			if targetVersion == "master" {
				version := oltppkg.GetCompVersion(ctx, tc, comtype)
				log.Info("Just get version for master:", zap.String("Version", version))
				oltppkg.CheckComponentVersion(ctx, tc, masterMap[comtype], comtype)
			} else {
				oltppkg.CheckComponentVersion(ctx, tc, targetVersion, comtype)
			}
		} else {
			log.Warn("Invalid image.")
		}
	}
	return nil
}