package scale

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

	. "github.com/onsi/gomega"
	pdclient "github.com/pingcap/endless/pkg/client/pd"
	"github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/workloadnode/cluster"
	"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"
	k8swait "k8s.io/apimachinery/pkg/util/wait"
)

func ScaleInTiKVParallel(ctx context.Context, tc resource.TiDBCluster, Percent float32, timeOut time.Duration, ifCheck bool) []string {
	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	storeIdMap := oltppkg.GetStoreIdAddrMap(pdURL, "id", "address")
	var storeIds []string
	for k, _ := range storeIdMap {
		storeIds = append(storeIds, k)
	}
	tobeDel := int(float32(len(storeIds)) * (1 - Percent))

	startTime := time.Now()
	for i := 0; i < tobeDel; i++ {
		log.Info("Delete store ", zap.String("storeId", storeIds[i]))
		oltppkg.ExecPdCtl(tc, fmt.Sprintf("store delete %s", storeIds[i]))
	}
	if ifCheck {
		promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)
		oltppkg.WaitFor(ctx, "waiting for node change to offline", 5*time.Minute)
		WaitForOffline(promURL, 0, 1*time.Minute, timeOut)

		ScaleProgressCheck(ctx, tc, "removing", startTime, time.Now(), ifCheck)

		pd := Try(pdclient.NewPDClient(tc)).(*pdclient.Client)
		Try(pd.DeleteTombstone())
	}
	return storeIds[0:tobeDel]
}
func GetUpStoreNum(promURL *url.URL) float64 {
	query := `sum(pd_cluster_status{type="store_up_count"})`
	curUpStore := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
	log.Info("current curUpStore:", zap.Float64("curUpStore", curUpStore))
	return curUpStore
}
func GetTombstoneNum(promURL *url.URL) float64 {
	query := `sum(pd_cluster_status{type="store_tombstone_count"})`
	curTombstone := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
	log.Info("current curTombstone:", zap.Float64("curTombstone", curTombstone))
	return curTombstone
}
func GetOfflineNum(promURL *url.URL) float64 {
	query := `sum(pd_cluster_status{type="store_offline_count"})`
	offlineNum := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
	log.Info("current offline:", zap.Float64("curoffline", offlineNum))
	return offlineNum
}
func GetPreparingNum(promURL *url.URL) float64 {
	query := `sum(pd_cluster_status{type="store_preparing_count"})`
	offlineNum := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
	log.Info("current offline:", zap.Float64("curoffline", offlineNum))
	return offlineNum
}
func WaitForTombstone(promURL *url.URL, tombestoneNum int, pollingInterval time.Duration, timeOut time.Duration) {
	query := `sum(pd_cluster_status{type="store_tombstone_count"})`
	Eventually(func() float64 {
		curTombstone := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
		log.Info("current curTombstone:", zap.Float64("curTombstone", curTombstone))
		return curTombstone
	}, timeOut, pollingInterval).Should(Equal(float64(tombestoneNum)))
}
func WaitForOffline(promURL *url.URL, tombestoneNum int, pollingInterval time.Duration, timeOut time.Duration) {
	query := `sum(pd_cluster_status{type="store_offline_count"})`
	Eventually(func() float64 {
		curTombstone := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
		log.Info("current offline:", zap.Float64("curoffline", curTombstone))
		return curTombstone
	}, timeOut, pollingInterval).Should(Equal(float64(tombestoneNum)))
}
func WaitForOnline(promURL *url.URL, online int, pollingInterval time.Duration, timeOut time.Duration) {
	query := `sum(pd_cluster_status{type="store_preparing_count"})`
	Eventually(func() float64 {
		curTombstone := Try(oltppkg.GetMetric(promURL, query, time.Now())).(float64)
		log.Info("current curOnline:", zap.Float64("curOnline", curTombstone))
		return curTombstone
	}, timeOut, pollingInterval).Should(Equal(float64(online)))
}
func GetComponentNum(tc resource.MySQLCompatible) map[string]uint32 {
	db := Try(tc.GetDB("")).(*sql.DB)
	defer db.Close()
	tidbClient := tidb.NewClient(db)
	tidbNum := Try(tidbClient.GetInstanceNum("tidb")).(uint32)
	pdNum := Try(tidbClient.GetInstanceNum("pd")).(uint32)
	tikvNum := Try(tidbClient.GetInstanceNum("tikv")).(uint32)

	curNumber := map[string]uint32{
		"tidb": tidbNum,
		"pd":   pdNum,
		"tikv": tikvNum}
	return curNumber
}
func ScaleOutInAndCheck(ctx context.Context, tc resource.TiDBCluster, runTimeParms oltppkg.RunTimePara, scaleOutPercent float64) {
	origNumber := GetComponentNum(tc)
	log.Info(fmt.Sprintf("current component number %v", origNumber))

	oltppkg.WaitFor(ctx, "Waiting for workload running before scale out", runTimeParms.WarmUpTime)

	targetNumber := map[string]uint32{
		"tidb": uint32(math.Floor(float64(origNumber["tidb"]) * (scaleOutPercent))),
		"pd":   uint32(math.Floor(float64(origNumber["pd"]) * (scaleOutPercent))),
		"tikv": uint32(math.Floor(float64(origNumber["tikv"]) * (scaleOutPercent))),
	}

	log.Info(fmt.Sprintf("scale out to target number %v", targetNumber))
	scaleOutStartTime := time.Now()
	for _, value := range []string{"pd", "tidb", "tikv"} {
		log.Info(fmt.Sprintf("Start to scale out %s to %d", value, targetNumber[value]))
		if value == "tikv" {
			scaleOutStartTime = time.Now()
		}
		err := tc.Scale(resource.TiDBClusterScaleOptions{
			Target:   oltppkg.GetResourceTypeByString(value),
			Replicas: targetNumber[value],
			Timeout:  30 * time.Minute,
		})
		Expect(err).Should(BeNil())
		log.Info(fmt.Sprintf("Finish to scale out %s to %d", value, targetNumber[value]))
	}

	log.Info("Waiting for store balance")
	oltppkg.WaitForBalance(ctx, tc, runTimeParms.TotalRunTime/2, 2*time.Minute, 0.1)

	scaleOutEndTime := time.Now()
	ScaleProgressCheck(ctx, tc, "preparing", scaleOutStartTime, scaleOutEndTime, false)

	log.Info(fmt.Sprintf("scale in to original number %v", origNumber))
	scaleInStartTime := time.Now()
	Try(oltppkg.ExecPdCtl(tc, "store limit all 10"))
	for _, value := range []string{"pd", "tidb", "tikv"} {
		log.Info(fmt.Sprintf("Start to scale in %s to %d", value, origNumber[value]))
		timeOut := 30 * time.Minute
		if value == "tikv" {
			scaleInStartTime = time.Now()
			timeOut = runTimeParms.TotalRunTime / 2
		}
		err := tc.Scale(resource.TiDBClusterScaleOptions{
			Target:   oltppkg.GetResourceTypeByString(value),
			Replicas: origNumber[value],
			Timeout:  timeOut,
		})
		Expect(err).Should(BeNil())
		log.Info(fmt.Sprintf("Start to scale in %s to %d", value, origNumber[value]))
	}
	//ScaleInTiKVParallel(ctx, tc, float32((scaleInPercent)))
	scaleInEndTime := time.Now()
	ScaleProgressCheck(ctx, tc, "removing", scaleInStartTime, scaleInEndTime, false)

	log.Info("remove tombstone")
	k8swait.PollImmediate(1*time.Minute, 10*time.Minute, func() (done bool, err error) {
		_, err = oltppkg.ExecPdCtl(tc, "store remove-tombstone")
		if err == nil {
			return true, nil
		}
		return false, nil
	})
}

func ScaleProgressCheck(ctx context.Context, tc resource.TiDBCluster, action string, startTime time.Time, endTime time.Time, ifVerify bool) {
	promURL := Try(tc.ServiceURL(resource.Prometheus)).(*url.URL)
	progressQuery := fmt.Sprintf(`avg(pd_cluster_progress{action="%s"} > 0)`, action)
	leftTimeQuery := `pd_cluster_eta{}`
	cost := endTime.Sub(startTime)
	log.Info("scale time", zap.String("action", action), zap.Time("startTime", startTime), zap.Time("endTime", endTime), zap.Duration("timeCost", cost))
	pctLsit := []float64{0.3, 0.5, 0.8}
	for _, pct := range pctLsit {
		value, err := oltppkg.GetMetric(promURL, progressQuery, startTime.Add(time.Duration((float64(cost) * pct))))

		if err != nil && strings.Contains(err.Error(), "empty") {
			log.Warn("Got empty, which is not expected", zap.String("progressQuery", progressQuery))
		} else {
			deviation := math.Abs(value - pct)
			log.Info("process deviation", zap.Float64("actualProcess", pct), zap.Float64("showProcess", value), zap.Float64("deviation", deviation))
			if ifVerify {
				Expect(deviation).Should(BeNumerically("<", 0.5))
			}
		}
		leftTime, err := oltppkg.GetMetric(promURL, leftTimeQuery, startTime.Add(time.Duration((float64(cost) * pct))))
		if err != nil && strings.Contains(err.Error(), "empty") {
			log.Warn("Got empty, which is not expected", zap.String("progressQuery", progressQuery))
		} else {
			leftTimeAcctual := time.Duration((float64(cost) * (1 - pct)))
			deviation := math.Abs(leftTimeAcctual.Minutes() - leftTime/60)
			log.Info("left time deviation ", zap.Duration("leftTimeAcctual", leftTimeAcctual), zap.Duration("showLeftTime", time.Duration(leftTime*float64(time.Second))), zap.Float64("deviation", deviation))
			log.Info("left time deviation minute", zap.Float64("deviation_minute", deviation))
			if ifVerify {
				Expect(deviation).Should(BeNumerically("<", leftTimeAcctual.Minutes()*0.5))
			}

		}

	}

}

func SetStoreLimitAccordingSize(ctx context.Context, tc cluster.BaseCluster) {
	_, sumSize, err := oltppkg.GetStoreSizeByType(tc, "used_size")
	Expect(err).Should(BeNil())
	//check if > 2T
	if sumSize > 2*1024*1024 {
		Try(oltppkg.BasicClusterExecPdCtl(tc, "store limit all 200"))
	} else {
		Try(oltppkg.BasicClusterExecPdCtl(tc, "store limit all 60"))
	}
}
func InitForScale(ctx context.Context, g *errgroup.Group, tc resource.TiDBCluster, workloadRunParms oltppkg.WorkloadRunParms, caseCmdParms oltppkg.CaseCmdParms, runTimeParms *oltppkg.RunTimePara,
	params ...interface{}) {
	//Start oom and output check before all.
	oltppkg.LogAndBy("Step 1: Start oom and output check before all.")

	qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
		IsCheck:      true,
		Offset:       "2m",
		ToleratePerc: 0.5,
		TolerateDur:  120 * time.Second,
	}
	durJitterCheckParms := oltppkg.DurJitterCheckParameters{
		IsCheck:      true,
		Offset:       "2m",
		ToleratePerc: 5,
		TolerateDur:  600 * time.Second,
	}
	params = append(params, qpsJitterCheckParms, durJitterCheckParms)

	oltppkg.FunCommonCheck(ctx, tc, g, runTimeParms.TotalRunTime, workloadRunParms.Cmdpr, params...)
	oltppkg.WaitFor(ctx, "waiting for workload running", runTimeParms.WarmUpTime)

	//Init parameters and wait cluster to balance
	oltppkg.LogAndBy("Step 2: Set store limit and waiting cluster balance.")
	SetStoreLimitAccordingSize(ctx, tc)
}

func GetScaleWorkloadCheckParms(caseCmdParms oltppkg.CaseCmdParms, ifTidbScale bool) oltppkg.WorkloadCheckParameters {
	var ignoreErr []string
	if ifTidbScale {
		ignoreErr = append(strings.Split(caseCmdParms.SysbenchParms.IgnoreErrors, ","), "8027", "9001", "9002", "invalid connection", "err dial tcp")
	} else {
		ignoreErr = strings.Split(caseCmdParms.SysbenchParms.IgnoreErrors, ",")
	}
	workloadCheckParms := oltppkg.WorkloadCheckParameters{
		IsCheck: true,
		//9001: PD server timeout
		//8027: Information schema is out of date: schema failed to update in 1 lease
		IgnoreErr: ignoreErr,
	}
	return workloadCheckParms
}
func ScaleClusterByNumber(ctx context.Context, tc resource.TiDBCluster, targetNumber map[string]uint32, timeout time.Duration) error {
	for key, value := range targetNumber {
		origNum := oltppkg.GetInstanceNum(tc, key)
		target := oltppkg.GetResourceTypeByString(key)
		if origNum == value {
			log.Info("Set same number with before, so no need to scale.",
				zap.String("key", key), zap.Uint32("number", value))
			continue
		}
		log.Info(fmt.Sprintf("Start scale %s to %d.", key, value))
		err := tc.Scale(resource.TiDBClusterScaleOptions{
			Target:   target,
			Replicas: uint32(value),
			Timeout:  timeout,
		})
		if err != nil && strings.Contains(err.Error(), "timed out") && value > origNum {
			log.Info("scale out time out, cancel the task")
			err2 := tc.Scale(resource.TiDBClusterScaleOptions{
				Target:   target,
				Replicas: uint32(origNum),
				Timeout:  timeout,
			})
			Expect(err2).Should(Succeed())
		}
		Expect(err).Should(Succeed())

		log.Info(fmt.Sprintf("Scale %s to %d finished.", key, value))
	}
	return nil
}
func ScaleClusterByPercent(ctx context.Context, tc resource.TiDBCluster, Percent float32, timeOut time.Duration) error {
	//Get original number of all component in this cluster
	orig_tidb := oltppkg.GetInstanceNum(tc, "tidb")
	orig_pd := oltppkg.GetInstanceNum(tc, "pd")
	orig_tikv := oltppkg.GetInstanceNum(tc, "tikv")
	targetNumber := map[string]uint32{
		"tikv": uint32(math.Ceil(float64(orig_tikv) * float64(Percent))),
		"tidb": uint32(math.Ceil(float64(orig_tidb) * float64(Percent))),
		"pd":   uint32(math.Ceil(float64(orig_pd) * float64(Percent)))}
	return ScaleClusterByNumber(ctx, tc, targetNumber, timeOut)
}