package loadsplit

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"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/query"
	"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"
)

type LoadSplitOpt struct {
	TableSize                 int64
	TableNum                  int64
	PureScan                  bool
	CpuOverloadRatio          string
	CpuOverloadRandom         bool
	CpuBaseEnableRandom       bool
	RegionThreshhold          int64
	DbName                    string
	TableName                 string
	ScanParallel              int64
	ChaosInject               bool
	ScanWithIndex             bool
	CpuBalanceCheckWarmup     time.Duration
	CpuBalanceCheckThreshhold float64
	IfStartUpdateWorkload     bool
}

func WaitForRegionHealth(ctx context.Context, tc resource.TiDBCluster, timeOut time.Duration, interval time.Duration) {
	Eventually(func() bool {
		return oltppkg.CheckRegionHealth(tc)
	}, timeOut, interval).Should(BeTrue())
}

func GetTableRegions(ctx context.Context, tc resource.TiDBCluster, dbName, tableName string) int64 {
	queryTmp := `select count(s.region_id) from INFORMATION_SCHEMA.TIKV_REGION_STATUS s where s.table_name = '%s' and s.db_name='%s'`
	db := Try(tc.GetDB(dbName)).(*sql.DB)
	defer db.Close()
	rows := Try(db.Query(fmt.Sprintf(queryTmp, tableName, dbName))).(*sql.Rows)
	defer rows.Close()
	var region_count float64
	for rows.Next() {
		Try(rows.Scan(&region_count))
	}
	log.Info("Get region count", zap.Int64(tableName, int64(region_count)))
	return int64(region_count)
}

func SetCpuThreshholdRatioRandom(ctxg context.Context, tc resource.TiDBCluster, runTime time.Duration, g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		log.Info("Set cpu threshhold random start.")
		interval := time.Duration(oltppkg.GetRandInt(1, 15)) * time.Minute
		k8swait.PollImmediateWithContext(ctxg, interval, runTime,
			func(context.Context) (done bool, err error) {
				ratio := float64(oltppkg.GetRandInt(5, 7)) * 0.01
				ratioStr := strconv.FormatFloat(ratio, 'f', 2, 64)
				oltppkg.SetAndVerifyConfigOnline(tc, "tikv", "split.region-cpu-overload-threshold-ratio", ratioStr)
				return false, nil
			})
		log.Info("Set cpu threshhold random finished.")
	}))
}

func EnableCpuBaseRandom(ctxg context.Context, tc resource.TiDBCluster, runTime time.Duration, g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		log.Info("Set cpu threshhold random start.")
		interval := time.Duration(oltppkg.GetRandInt(2, 15)) * time.Minute
		k8swait.PollImmediateWithContext(ctxg, interval, runTime,
			func(context.Context) (done bool, err error) {
				randint := oltppkg.GetRandInt(0, 2)
				ratioList := []string{"0", "0.05"}
				oltppkg.SetAndVerifyConfigOnline(tc, "tikv", "split.region-cpu-overload-threshold-ratio", ratioList[randint])
				return false, nil
			})
		log.Info("Set cpu threshhold random finished.")
	}))
}

func RandomRestartInstanceRepeatly(ctxg context.Context,
	tc resource.TiDBCluster,
	runTime time.Duration,
	g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		log.Info("Random restart instance start.")
		interval := time.Duration(oltppkg.GetRandInt(5, 15)) * time.Minute
		k8swait.PollImmediateWithContext(ctxg, interval, runTime,
			func(context.Context) (done bool, err error) {
				oltppkg.RandomlyRestartInstance(tc, resource.TiKV)
				oltppkg.WaitFor(ctxg, "waiting for next restart", 5*time.Minute)
				oltppkg.RandomlyRestartInstance(tc, resource.PD)
				return false, nil
			})
		log.Info("Random restart instance finished.")
	}))
}

func (opt *LoadSplitOpt) RunWorkloadForLoadSplit(ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	benchToolset resource.WorkloadNode,
	runTime time.Duration,
	g *errgroup.Group,
) {
	opt.RunFullTableScanParallel(ctx, tc, runTime, g)
	if !opt.PureScan {
		dbURL, _ := tc.ServiceURL(resource.DBAddr)
		sysbench := oltppkg.NewSysbench("oltp_point_select", "run")
		sysbenchcmd := sysbench.Db(opt.DbName).Host(dbURL.Hostname()).Port(dbURL.Port()).Threads(100).
			Tables(1).TableSize(int(opt.TableSize)).IgnoreError(caseCmdParms.SysbenchParms.IgnoreErrors)
		runcmd := sysbenchcmd.Build()
		Try(oltppkg.ExecWorkloadCmd(ctx, benchToolset, runcmd, false, nil))
	}
	if opt.IfStartUpdateWorkload {
		dbURL, _ := tc.ServiceURL(resource.DBAddr)
		sysbench := oltppkg.NewSysbench("oltp_write_only", "run")
		sysbenchcmd := sysbench.Db(opt.DbName).Host(dbURL.Hostname()).Port(dbURL.Port()).Threads(1).
			Tables(1).TableSize(int(opt.TableSize)).IgnoreError(caseCmdParms.SysbenchParms.IgnoreErrors)
		runcmd := sysbenchcmd.Build()
		Try(oltppkg.ExecWorkloadCmd(ctx, benchToolset, runcmd, false, nil))
	}
}
func (opt *LoadSplitOpt) PrepareData(ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	benchToolset resource.WorkloadNode) {
	log.Info("Prepare data")
	timestamp := strconv.FormatInt(time.Now().UnixNano()/1e6, 10)
	dbName := fmt.Sprintf("%s_%s", "sysbenchprepare", timestamp)
	Try(oltppkg.RecreateDatabase(ctx, tc, dbName))
	dbURL, _ := tc.ServiceURL(resource.DBAddr)

	sysbench := oltppkg.NewSysbench("oltp_insert", "prepare")
	sysbenchcmd := sysbench.Db(dbName).Host(dbURL.Hostname()).Port(dbURL.Port()).Threads(100).
		Tables(1).TableSize(int(opt.TableSize)).IgnoreError(caseCmdParms.SysbenchParms.IgnoreErrors)
	runcmd := sysbenchcmd.Build()
	Try(oltppkg.ExecWorkloadCmd(ctx, benchToolset, runcmd, true, nil))
	opt.DbName = dbName
	opt.TableName = "sbtest1"

	if !opt.ScanWithIndex {
		log.Info("alter index invisible")
		db := Try(tc.GetDB(dbName)).(*sql.DB)
		defer db.Close()
		Try(db.ExecContext(ctx, fmt.Sprintf("alter table %s alter index k_1 invisible", opt.TableName)))
	}
}

func (opt *LoadSplitOpt) ContinueGetRegionCount(ctxg context.Context, tc resource.TiDBCluster, runTime time.Duration, g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		log.Info("Continue get region count start,")
		k8swait.PollImmediateWithContext(ctxg, 5*time.Minute, runTime,
			func(context.Context) (done bool, err error) {
				GetTableRegions(ctxg, tc, opt.DbName, opt.TableName)
				return false, nil
			})
		log.Info("Continue get region count finished.")
	}))
}

func (opt *LoadSplitOpt) RegionCountShouldMoreThan(ctxg context.Context, tc resource.TiDBCluster, runTime time.Duration, g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		log.Info("Region count check start.")
		err := k8swait.PollImmediateWithContext(ctxg, 1*time.Minute, runTime,
			func(context.Context) (done bool, err error) {
				regionCount := GetTableRegions(ctxg, tc, opt.DbName, opt.TableName)
				if regionCount >= opt.RegionThreshhold {
					return true, nil
				}
				return false, nil
			})
		if err != nil {
			log.Info(fmt.Sprintf("err is %v", err))
			Expect(strings.Contains(err.Error(), "timed out waiting for the condition")).ShouldNot(BeTrue())
			Expect(errors.Is(err, k8swait.ErrWaitTimeout)).ShouldNot(BeTrue())
		}

		log.Info("Region count check finished.")
	}))
}

func (opt *LoadSplitOpt) TiKVCpuBalanceCheck(ctxg context.Context, tc resource.TiDBCluster, runTime time.Duration, g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(ctxg, "waiting for check tikv cpu balance", opt.CpuBalanceCheckWarmup)
		log.Info("Tikv CPU check start.")
		ConsistentlyWithMetrics(tc, runTime).ShouldNot(
			TiKVCPUDifferentExceeds(opt.CpuBalanceCheckThreshhold).For(30 * time.Second).WithCtx(ctxg),
		)
		log.Info("Tikv CPU check start.")
	}))
}

func (opt *LoadSplitOpt) RunFullTableScanParallel(ctxg context.Context,
	tc resource.TiDBCluster,
	runTime time.Duration,
	g *errgroup.Group) {
	g.Go(WithGinkgoRecover(func() {
		var wg sync.WaitGroup
		log.Info("full tables scan finished")
		for i := 1; i <= int(opt.ScanParallel); i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				query.FullScanAllTablesRepeatedly(ctxg, tc, runTime, opt.DbName, []string{}, []string{opt.TableName})
			}()
		}
		wg.Wait()
		log.Info("full tables scan finished")
	}))
}

func (opt *LoadSplitOpt) LoadSplitTest(
	ctx context.Context,
	tc resource.TiDBCluster,
	caseCmdParms oltppkg.CaseCmdParms,
	benchToolset resource.WorkloadNode,
) {
	var runTime time.Duration
	var warmUpTime time.Duration
	if caseCmdParms.RunTime.TotalRunTime != 0 {
		runTime = caseCmdParms.RunTime.TotalRunTime
	} else {
		runTime = 1 * time.Hour
	}
	if caseCmdParms.RunTime.WarmUpTime != 0 && !caseCmdParms.SkipRunWorkload {
		warmUpTime = caseCmdParms.RunTime.WarmUpTime
		oltppkg.WaitFor(ctx, "workload warmup", warmUpTime)
	}
	log.Info("Test will run ", zap.Duration("case runTime", runTime))
	ctxAll, cancel := context.WithTimeout(ctx, runTime)
	defer cancel()
	g, ctxg := errgroup.WithContext(ctxAll)

	log.Info("start run workload")
	oltppkg.SetAndVerifyConfigOnline(tc, "tikv", "readpool.unified.max-thread-count", "4")

	corpCacheList := oltppkg.GetConfigByName(tc, "tidb", "tikv-client.copr-cache.capacity-mb")
	coprCacheEnable := false
	for _, coprCache := range corpCacheList {
		if coprCache.Value != "0" {
			coprCacheEnable = true
			break
		}
	}
	if coprCacheEnable {
		opt.IfStartUpdateWorkload = true
	}

	qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
		IsCheck:      true,
		Offset:       "2m",
		ToleratePerc: 0.2,
		TolerateDur:  120 * time.Second,
	}
	if !opt.ChaosInject {
		oltppkg.FunCommonCheck(ctxg, tc, g, runTime, nil, qpsJitterCheckParms)
	} else {
		restartCheckParms := oltppkg.RestartCheckParameters{
			AllComponentRestartCheck: false,
			TiKVRestartCheck:         false,
			PDRestartCheck:           false,
			TiDBRestartCheck:         true,
		}
		qpsJitterCheckParms := oltppkg.QpsJitterCheckParameters{
			IsCheck: false,
		}
		durJitterCheckParms := oltppkg.DurJitterCheckParameters{
			IsCheck: false,
		}
		oltppkg.FunCommonCheck(ctxg, tc, g, runTime, nil, qpsJitterCheckParms, durJitterCheckParms, restartCheckParms)
	}

	opt.PrepareData(ctxg, tc, caseCmdParms, benchToolset)
	opt.RunWorkloadForLoadSplit(ctxg, tc, caseCmdParms, benchToolset, runTime, g)
	opt.ContinueGetRegionCount(ctxg, tc, runTime, g)

	if !opt.CpuOverloadRandom && !opt.ChaosInject {
		oltppkg.SetAndVerifyConfigOnline(tc, "tikv", "split.region-cpu-overload-threshold-ratio", "0.05")
		opt.TiKVCpuBalanceCheck(ctxg, tc, runTime, g)
	}
	if opt.CpuOverloadRandom {
		SetCpuThreshholdRatioRandom(ctxg, tc, runTime, g)
	}
	if opt.ChaosInject {
		oltppkg.WaitFor(ctxg, "Waiting for workload running", 5*time.Minute)
		RandomRestartInstanceRepeatly(ctxg, tc, runTime, g)
	}
	opt.RegionCountShouldMoreThan(ctxg, tc, runTime, g)

	if opt.CpuBaseEnableRandom {
		EnableCpuBaseRandom(ctxg, tc, runTime, g)
	}
	if opt.CpuOverloadRandom {
		SetCpuThreshholdRatioRandom(ctxg, tc, runTime, g)
	}

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