package pdfun

import (
	"context"
	"database/sql"
	"fmt"
	"net/url"
	"strconv"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/endless/pkg/chaos"
	. "github.com/pingcap/endless/pkg/matcher"
	"go.uber.org/zap"

	pdclient "github.com/pingcap/endless/pkg/client/pd"
	"github.com/pingcap/endless/pkg/client/tidb"
	. "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"
	"golang.org/x/sync/errgroup"
)

func OLTP_EVICT_LEADER_FUN_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimePara := oltppkg.NewRunTimePara(3*time.Hour, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 10*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimePara.TotalRunTime)
	defer cancelAll()
	pd := Try(pdclient.NewPDClient(tc)).(*pdclient.Client)
	Try(pd.SetAllStoreLimit(200))
	defer pd.SetAllStoreLimit(15)

	g, ctxg := errgroup.WithContext(ctxAll)

	promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)

	defer func() {
		Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
		oltppkg.WaitFor(context.Background(), "waiting for leader balance", 3*time.Minute)
	}()

	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	idInstanceMap := oltppkg.GetStoreIdAddrMap(pdURL, "id", "address")
	log.Info(fmt.Sprintf("%v", idInstanceMap))
	var storeID string
	storeAddr := fmt.Sprintf("tc-tikv-%d", oltppkg.GetInstanceNum(tc, "tikv")-1)
	for id, addr := range idInstanceMap {
		if addr == storeAddr {
			storeID = id
			break
		}
	}
	Try(oltppkg.ExecPdCtl(tc, fmt.Sprintf("scheduler add evict-leader-scheduler %s", storeID)))
	Expect(GetEvictLeaderCount(tc)).Should(Equal(1))
	oltppkg.WaitForLeaderToZero(promURL, storeAddr, 10*time.Second, 20*time.Minute)

	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: true,
		TiKVRestartCheck:         false,
		PDRestartCheck:           false,
		TiDBRestartCheck:         false,
	}

	oltppkg.FunCommonCheck(ctxg, tc, g, runTimePara.TotalRunTime, workloadRunParms.Cmdpr, restartCheckParms)

	db := Try(tc.GetDB("")).(*sql.DB)
	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,
		"pd":   initPdNum}
	oltppkg.LogAndBy("Step 2: Scale-in the cluster")
	oltppkg.WaitFor(ctxg, "Warm workload", 10*time.Minute)

	g.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(ctxg, "Wait for workload running", 5*time.Minute)
		log.Info(fmt.Sprintf("scale target %v", targetNumber))
		Try(scale.ScaleClusterByNumber(ctxg, tc, targetNumber, 4*time.Hour))
	}))
	Expect(g.Wait()).Should(BeNil())
}

func OLTP_EVICT_LEADER_FUN_002(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimePara := oltppkg.NewRunTimePara(60*time.Minute, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimePara.TotalRunTime)
	defer cancelAll()

	pd := Try(pdclient.NewPDClient(tc)).(*pdclient.Client)
	Try(pd.SetAllStoreLimit(200))
	defer pd.SetAllStoreLimit(15)

	g, ctxg := errgroup.WithContext(ctxAll)

	promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)

	defer func() {
		Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
		oltppkg.WaitFor(context.Background(), "waiting for leader balance", 3*time.Minute)
	}()

	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	idInstanceMap := oltppkg.GetStoreIdAddrMap(pdURL, "id", "address")
	log.Info(fmt.Sprintf("%v", idInstanceMap))
	var storeID string
	storeAddr := fmt.Sprintf("tc-tikv-%d", oltppkg.GetInstanceNum(tc, "tikv")-1)
	for id, addr := range idInstanceMap {
		if addr == storeAddr {
			storeID = id
			break
		}
	}

	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: true,
		TiKVRestartCheck:         false,
		PDRestartCheck:           false,
		TiDBRestartCheck:         false,
	}

	oltppkg.FunCommonCheck(ctxg, tc, g, runTimePara.TotalRunTime, workloadRunParms.Cmdpr, restartCheckParms)

	db := Try(tc.GetDB("")).(*sql.DB)
	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,
		"pd":   initPdNum}

	oltppkg.LogAndBy("Step 2: Scale-in the cluster")
	oltppkg.WaitFor(ctxg, "Warm workload", 10*time.Minute)

	g.Go(WithGinkgoRecover(func() {
		Try(scale.ScaleClusterByNumber(ctxg, tc, targetNumber, 4*time.Hour))
	}))
	g.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(ctxg, "Wait for workload running", 5*time.Minute)
		Try(oltppkg.ExecPdCtl(tc, fmt.Sprintf("scheduler add evict-leader-scheduler %s", storeID)))
		Expect(GetEvictLeaderCount(tc)).Should(Equal(1))
		oltppkg.WaitForLeaderToZero(promURL, storeAddr, 10*time.Second, 20*time.Minute)
	}))
	Expect(g.Wait()).Should(BeNil())
}

// hit bug https://github.com/tikv/pd/issues/4707
func OLTP_EVICT_LEADER_REL_001(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimePara := oltppkg.NewRunTimePara(60*time.Minute, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimePara.TotalRunTime)
	defer cancelAll()

	promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)

	defer func() {
		Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
		oltppkg.WaitFor(context.Background(), "waiting for leader balance", 3*time.Minute)
	}()
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	idInstanceMap := oltppkg.GetStoreIdAddrMap(pdURL, "id", "address")
	log.Info(fmt.Sprintf("%v", idInstanceMap))

	c := chaos.Chaos{
		Name:           "pd leader",
		Tc:             tc,
		ExperimentSpec: &resource.FailureExperimentSpec{},
		SelectorOpts: resource.TiDBClusterSelectorOpts{
			TiDBClusterPDSelectorSpec: &resource.TiDBClusterPDSelectorSpec{
				IsLeader: &[]bool{true, false}[0],
			},
		},
	}

	g, ctxg := errgroup.WithContext(ctxAll)
	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         true,
		PDRestartCheck:           false,
		TiDBRestartCheck:         true,
	}
	qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
		IsCheck:      false,
		Offset:       "2m",
		ToleratePerc: 0.5,
		TolerateDur:  120 * time.Second,
	}
	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck: true,
		//9001: PD server timeout
		IgnoreErr: append(GetIngoreErrors("fun"), "9001", "8027", "fail"),
	}
	oltppkg.FunCommonCheck(ctxg, tc, g, runTimePara.TotalRunTime, workloadRunParms.Cmdpr, restartCheckParms, qpsJitterCheckParms, workloadCheckParms)

	g.Go(WithGinkgoRecover(func() {
		interval := 1 * time.Minute
	Loop:
		for {
			select {
			case <-time.After(interval):
				log.Info("Remove evict leader scheduler")
				Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
				Eventually(tc, GetLeaderBalanceTime(tc), 10*time.Second).ShouldNot(RegionLeaderDifferentExceeds(0.05).WithCtx(ctxg))
			case <-ctxg.Done():
				break Loop
			}
			reservedIndex := oltppkg.GetRandInt(0, len(idInstanceMap))
			start := 0
			evictNum := 0
			for id, addr := range idInstanceMap {
				if start == reservedIndex {
					start++
					continue
				}
				if evictNum >= 2 {
					break
				}
				evictNum++
				Try(oltppkg.ExecPdCtl(tc, fmt.Sprintf("scheduler add evict-leader-scheduler %s", id)))
				oltppkg.WaitForLeaderToZero(promURL, addr, 10*time.Second, 20*time.Minute)
				start++
			}

			oltppkg.WaitFor(ctxg, "waiting for evict leader works", 5*time.Minute)
			Expect(GetEvictLeaderCount(tc)).Should(Equal(2))
			oltppkg.InjectFaultByChaos(&c, context.Background(), 2*time.Minute)

			time.Sleep(2 * time.Minute)
			//Comment this checker for there is https://github.com/tikv/pd/issues/4707
			Expect(GetEvictLeaderCount(tc)).Should(Equal(2))
		}
	}))
	Expect(g.Wait()).Should(BeNil())
	Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
	oltppkg.WaitForBalance(ctxg, tc, 10*time.Minute, 1*time.Minute, 0.3)
}

func OLTP_EVICT_LEADER_REL_002(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimePara := oltppkg.NewRunTimePara(30*time.Minute, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimePara.TotalRunTime)
	defer cancelAll()

	defer func() {
		Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
		oltppkg.WaitFor(context.Background(), "waiting for leader balance", 3*time.Minute)
	}()
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	idInstanceMap := oltppkg.GetStoreIdAddrMap(pdURL, "id", "address")
	log.Info(fmt.Sprintf("%v", idInstanceMap))

	g, ctxg := errgroup.WithContext(ctxAll)

	promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)

	restartCheckParms := oltppkg.RestartCheckParameters{
		AllComponentRestartCheck: false,
		TiKVRestartCheck:         false,
		PDRestartCheck:           true,
		TiDBRestartCheck:         false,
	}

	oltppkg.FunCommonCheck(ctxg, tc, g, runTimePara.TotalRunTime, workloadRunParms.Cmdpr, restartCheckParms)

	interval := 1 * time.Minute
Loop:
	for {
		select {
		case <-time.After(interval):
			Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
			oltppkg.WaitForBalance(ctxg, tc, 10*time.Minute, 1*time.Minute, 0.2)
		case <-ctxg.Done():
			break Loop
		}
		reservedIndex := oltppkg.GetRandInt(0, len(idInstanceMap))
		start := 0
		var evictAddr string
		for storeId, storeAddr := range idInstanceMap {
			if start != reservedIndex {
				start++
				continue
			}
			evictAddr = storeAddr
			Try(oltppkg.ExecPdCtl(tc, fmt.Sprintf("scheduler add evict-leader-scheduler %s", storeId)))
			oltppkg.WaitForLeaderToZero(promURL, storeAddr, 10*time.Second, 20*time.Minute)
			start++
		}

		oltppkg.WaitFor(ctxg, "waiting for evict leader works", 5*time.Minute)
		Expect(GetEvictLeaderCount(tc)).Should(Equal(1))

		tmpList := strings.Split(evictAddr, "-")
		peerId, err := strconv.Atoi(tmpList[len(tmpList)-1])
		Expect(err).Should(BeNil())

		c := chaos.Chaos{
			Name:           "tikv down",
			Tc:             tc,
			ExperimentSpec: &resource.FailureExperimentSpec{},
			SelectorOpts: resource.TiDBClusterSelectorOpts{
				TiDBClusterTiKVSelectorSpec: &resource.TiDBClusterTiKVSelectorSpec{
					Peer: oltppkg.ToIntPtr(peerId),
				},
			},
		}
		oltppkg.InjectFaultByChaos(&c, context.Background(), 2*time.Minute)

		oltppkg.WaitForLeaderToZero(promURL, evictAddr, 10*time.Second, 20*time.Minute)
		time.Sleep(2 * time.Minute)
		Expect(GetEvictLeaderCount(tc)).Should(Equal(1))
	}
	Expect(g.Wait()).Should(BeNil())
}

func OLTP_EVICT_LEADER_REL_003(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	workloadRunParms oltppkg.WorkloadRunParms,
) {
	runTimePara := oltppkg.NewRunTimePara(60*time.Minute, 10*time.Minute, oltppkg.GetRandDuration(time.Minute, 10), 2*time.Minute)
	ctxAll, cancelAll := context.WithTimeout(ctx, runTimePara.TotalRunTime)
	defer cancelAll()

	defer func() {
		Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
		oltppkg.WaitFor(context.Background(), "waiting for leader balance", 3*time.Minute)
	}()

	promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	idInstanceMap := oltppkg.GetStoreIdAddrMap(pdURL, "id", "address")
	log.Info(fmt.Sprintf("%v", idInstanceMap))

	g, ctxg := errgroup.WithContext(ctxAll)
	qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
		IsCheck:      false,
		Offset:       "2m",
		ToleratePerc: 0.8,
		TolerateDur:  120 * time.Second,
	}
	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck: true,
		//9002: TiKV server timeout
		IgnoreErr: []string{"1213", "1105", "1205", "8022", "8028", "9004", "9007", "1062", "9002", "invalid connection", "fail"},
	}
	oltppkg.FunCommonCheck(ctxg, tc, g, runTimePara.TotalRunTime, workloadRunParms.Cmdpr, qpsJitterCheckParms, workloadCheckParms)

	interval := 10 * time.Minute
Loop:
	for {
		select {
		case <-time.After(interval):
			oltppkg.WaitFor(ctxg, "waiting for recover", runTimePara.RecoveryTime)
			Try(oltppkg.ExecPdCtl(tc, "scheduler remove evict-leader-scheduler"))
			Eventually(tc, GetLeaderBalanceTime(tc), 10*time.Second).ShouldNot(RegionLeaderDifferentExceeds(0.1).WithCtx(ctxg))
			log.Info("Leader balance now.")
		case <-ctxg.Done():
			break Loop
		}

		reservedIndex := oltppkg.GetRandInt(0, len(idInstanceMap))
		start := 0
		var evictAddr string
		var evictId string
		for storeId, storeAddr := range idInstanceMap {
			if start != reservedIndex {
				start++
				continue
			}
			evictAddr = storeAddr
			evictId = storeId
			start++
			break
		}

		log.Info(fmt.Sprintf("idInstanceMap %v", idInstanceMap))
		tmpList := strings.Split(idInstanceMap[evictId], "-")
		log.Info("tmpList", zap.Strings("tmpList", tmpList))
		peerId, err := strconv.Atoi(tmpList[len(tmpList)-1])
		Expect(err).Should(BeNil())

		c := chaos.Chaos{
			Name:           "tikv down",
			Tc:             tc,
			ExperimentSpec: &resource.FailureExperimentSpec{},
			SelectorOpts: resource.TiDBClusterSelectorOpts{
				TiDBClusterTiKVSelectorSpec: &resource.TiDBClusterTiKVSelectorSpec{
					Peer: oltppkg.ToIntPtr(peerId),
				},
			},
		}
		c.Run()

		Try(oltppkg.ExecPdCtl(tc, fmt.Sprintf("scheduler add evict-leader-scheduler %s", evictId)))
		Expect(GetEvictLeaderCount(tc)).Should(Equal(1))

		oltppkg.WaitFor(ctxg, "waiting for evict leader works", 5*time.Minute)
		c.Clean()

		oltppkg.WaitForLeaderToZero(promURL, evictAddr, 10*time.Second, 20*time.Minute)
		Expect(GetEvictLeaderCount(tc)).Should(Equal(1))
	}
	Expect(g.Wait()).Should(BeNil())
}