package common_test

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"
	"time"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"github.com/pingcap/endless/pkg/util/schedule"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/pkg/client/tidb"
	"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/tiflash/pkg"
	"github.com/pingcap/endless/testcase/tiflash/pkg/chaos"
	"github.com/pingcap/endless/testcase/tiflash/pkg/checker"
	workload2 "github.com/pingcap/endless/testcase/tiflash/pkg/workload"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/workloads"
)

// use timeout to control stop
const cpuProfile = "timeout %ds bash -c 'while true; do curl -sSL http://127.0.0.1:20292/debug/pprof/profile?seconds=10 -o a; echo start && sleep 10;done'"

var _ = Describe("CommonOLAPTest", func() {
	var (
		wlType             workloads.WorkloadType
		wl                 workload2.Workload
		tc                 resource.TiDBCluster
		duration           time.Duration
		db                 *sql.DB
		tidbClient         *tidb.Client
		chaosIdRec         resource.ChaosIdentifier
		recoveryAfterChaos func(nodeType resource.TiDBCompType, faultType resource.ExperimentKind)
		checkBalance       bool
		mode               tiflashMode
		enableCompact      bool
	)
	BeforeEach(func() {
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		wlType = workloads.WorkloadType(suiteTestCtx.Param("workload", string(workloads.CHWorkloadType)))
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "3h"))).(time.Duration)
		mode = tiflashMode(suiteTestCtx.Param("mode", string(normal)))
		enableCompact = matcher.Try(strconv.ParseBool(suiteTestCtx.Param("compact", "false"))).(bool)
		log.Info("some case params", zap.String("workload", string(wlType)), zap.String("mode", string(mode)),
			zap.Duration("duration", duration), zap.Bool("compact", enableCompact))
		if w, ok := workloads.HtapWorkloads[wlType]; ok {
			wl = w(suiteTestCtx)
		} else {
			Skip("don't have the workload " + string(wlType))
		}
		recoveryAfterChaos = func(nodeType resource.TiDBCompType, faultType resource.ExperimentKind) {
			ctx0, _ := context.WithTimeout(context.Background(), 60*time.Minute)
			g, ctx := errgroup.WithContext(ctx0)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunAP(ctx, pkg.SkipErrorHandler))
			}))
			// 	注入错误
			g.Go(util.WithGinkgoRecover(func() {
				// query run normal before chaos
				time.Sleep(5 * time.Minute)
				chaosIdRecI := matcher.Try(chaos.InjectFault(ctx, tc, nodeType, faultType))
				if chaosIdRecI != nil {
					chaosIdRec = chaosIdRecI.(resource.ChaosIdentifier)
				} else {
					chaosIdRec = nil
				}
				select {
				case <-ctx.Done():
					matcher.Try(chaos.DeleteChaos(tc, chaosIdRec))
				}
			}))
			matcher.Try(g.Wait())
			// 等待 10 分钟集群可以恢复服务
			log.Info("expect stable run ap after sleep 10 minutes")
			time.Sleep(10 * time.Minute)
			log.Info("start stable run 10 minutes without chaos")
			ctx2, _ := context.WithTimeout(context.Background(), 10*time.Minute)
			matcher.Try(wl.RunAP(ctx2, pkg.APCommonErrHandler))
		}
		checkBalance = true
		matcher.Try(pkg.DisableExpressionPushdown(tc, []string{"date_add"}, []string{"tiflash"}))
		wl.Prepare(suiteTestCtx)
		// the database is created after prepare
		db = matcher.Try(tc.GetDB(wl.DBName())).(*sql.DB)
		tidbClient = tidb.NewClient(db)
		matcher.Try(tidbClient.ShowClusterVersionInformation())
		handleTiFlashMode(suiteTestCtx, tc, []string{wl.DBName()}, mode, true)
		// https://github.com/pingcap/tidb/issues/33037
		matcher.Try(tidbClient.Exec("set @@global.tidb_enable_rate_limit_action=0"))
		// auto analyze global stats
		matcher.Try(tidbClient.Exec("set @@global.tidb_partition_prune_mode='dynamic'"))
	})
	// borrow from the oltp cases
	AfterEach(func() {
		By("End One Test")
		chaos.DeleteChaos(tc, chaosIdRec)
		time.Sleep(2 * time.Minute)
		oltppkg.LogAndBy("Post check 1: Test finished, check if all stores and pds are up.")
		oltppkg.CheckClusterHealth(tc)
		if checkBalance {
			// TODO check page store v2 or v3 and uncomment it
			// matcher.Try(pkg.CheckTiFlashTiKVStoreSize(context.Background(), tc))
			oltppkg.LogAndBy("Post check 2: Test finished, check if region size/number, store size if balance in all tiflash instance.")
			ctx2, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
			defer cancel()
			g, ctxg := errgroup.WithContext(ctx2)
			threshhold := 0.5
			checkTime := 2 * time.Minute
			g.Go(util.WithGinkgoRecover(func() {
				matcher.ConsistentlyWithMetrics(tc, checkTime).ShouldNot(metrics.TiFlashStoreRegionSizeDifferentExceeds(threshhold).For(30 * time.Second).WithCtx(ctxg))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.ConsistentlyWithMetrics(tc, checkTime).ShouldNot(metrics.TiFlashStoreRegionCountDifferentExceeds(threshhold).For(30 * time.Second).WithCtx(ctxg))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.ConsistentlyWithMetrics(tc, checkTime).ShouldNot(metrics.TiFlashStoreStoreSizeDifferentExceeds(threshhold).For(30 * time.Second).WithCtx(ctxg))
			}))
			matcher.Try(g.Wait())
		}
	})

	Context("run workload", func() {
		It("should serve stably #001#", func() {
			enableProfile := matcher.Try(strconv.ParseBool(suiteTestCtx.Param("profile", "false"))).(bool)
			// 	长稳运行 3 小时不失败
			ctx0, _ := context.WithTimeout(suiteTestCtx, duration)
			g, ctx := errgroup.WithContext(ctx0)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunTP(ctx, pkg.TPCommonErrHandler))
			}))
			if enableProfile {
				num := matcher.Try(tidbClient.GetInstanceNum("tiflash")).(uint32)
				for i := 0; i < int(num); i++ {
					j := i
					g.Go(util.WithGinkgoRecover(func() {
						// TODO use execAsync to replace it after test-infra tidb cluster support exec async
						// it just use timeout but can't handle the ctx.Done
						pkg.TiFlashExecCmd(tc, fmt.Sprintf(cpuProfile, int(duration.Seconds())), j)
					}))
				}
			}
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunAP(ctx, pkg.APCommonErrHandler))
			}))
			checker.CommonCheck(ctx, tc, g, duration, wl.GetHistograms())
			handleTiFlashMode(ctx, tc, []string{wl.DBName()}, mode, false)
			if enableCompact {
				op := matcher.Try(chaos.NewCompactTiFlashReplica(tc)).(*chaos.CompactTiFlashReplica)
				matcher.Try(schedule.RunPeriodically(ctx, "@every 20m", 10*time.Minute, op, "randomCompactTiFlashReplica"))
			}
			matcher.Try(g.Wait())
		})
		It("should server stably when scale tiflash #002#", func() {
			// 	扩缩容，查询不会失败
			ctx0, _ := context.WithTimeout(suiteTestCtx, duration)
			g, ctx := errgroup.WithContext(ctx0)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunTP(ctx, pkg.TPCommonErrHandler))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				componentNum := matcher.Try(tidbClient.GetAllTypeInstanceNum(suiteTestCtx)).(map[string]uint32)
				log.Info("cluster info", zap.Any("component num", componentNum))
				err := pkg.ScaleCluster(tc, resource.TiFlash, componentNum[resource.TiFlash.String()]+2)
				// 如果 scale-out timeout，最终也会扩容（如果有资源的话），这会影响后续 case（一个 plan 多个 cases 这样的编排）的 balance 检测
				// 如下一个 case 结束了，但是在检测 store balance 的时候，悄悄的正在扩容节点，这时候检测就会不通过
				// 当前可以通过 maxReplicas 字段用于预占资源保证 case 扩容/升级成功, 扩容场景填最大副本数，纯升级场景可填 maxReplicas = replicas
				if err != nil {
					checkBalance = false
					log.Error("scale error", zap.Error(err))
					return
				}
				// 观察 daily run 的扩容速度，一个小时时间太短，需要延长
				if pkg.Sleep(ctx, 3*time.Hour) {
					return
				}
				// TODO check balance after scale-in/out
				componentNum = matcher.Try(tidbClient.GetAllTypeInstanceNum(suiteTestCtx)).(map[string]uint32)
				log.Info("cluster info", zap.Any("component num", componentNum))
				err = pkg.ScaleCluster(tc, resource.TiFlash, componentNum[resource.TiFlash.String()]-2)
				if err != nil {
					checkBalance = false
					log.Error("scale error", zap.Error(err))
					return
				}
				if pkg.Sleep(ctx, 1*time.Hour) {
					return
				}
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunAP(ctx, pkg.APCommonErrHandler))
			}))
			matcher.Try(g.Wait())
		})
		It("should recovery to be normal after chaos #003#", func() {
			// 	kill，网络隔离，网络延迟等 chaos 结束后，查询可以恢复正常
			ctx0, cancel := context.WithTimeout(suiteTestCtx, duration)
			g, _ := errgroup.WithContext(ctx0)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunTP(ctx0, pkg.TPCommonErrHandler))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				defer cancel()
				// 需要设置 maxReplicas 才可以进行反复的重启 pod，保障 pod 重启的时候有足够的资源
				recoveryAfterChaos(resource.TiFlash, resource.ExperimentFailureKind)
				recoveryAfterChaos(resource.TiFlash, chaos.ExperimentFullDiskKind)
				recoveryAfterChaos(resource.TiFlash, resource.ExperimentNetworkPartitionKind)
				recoveryAfterChaos(resource.TiFlash, resource.ExperimentNetworkDelayKind)
				recoveryAfterChaos(resource.TiFlash, resource.ExperimentNetworkLossKind)

				// recoveryAfterChaos(resource.TiKV, resource.ExperimentFailureKind)
				// recoveryAfterChaos(resource.TiKV, resource.ExperimentNetworkPartitionKind)
				// recoveryAfterChaos(resource.TiKV, resource.ExperimentNetworkDelayKind)
				// recoveryAfterChaos(resource.TiKV, resource.ExperimentNetworkLossKind)
				//
				// recoveryAfterChaos(resource.PD, resource.ExperimentFailureKind)
				// recoveryAfterChaos(resource.PD, resource.ExperimentNetworkPartitionKind)
				// recoveryAfterChaos(resource.PD, resource.ExperimentNetworkDelayKind)
				// recoveryAfterChaos(resource.PD, resource.ExperimentNetworkLossKind)
			}))
			g.Go(util.WithGinkgoRecover(func() {
				checker.TiFlashCrashChecker(tc, duration*2)
			}))
			matcher.Try(g.Wait())
		})
		It("should serve stably when ddl #004#", func() {
			// ddl 执行时间较长，使用自定义的 db
			db2 := matcher.Try(pkg.OpenDB(tc, wl.DBName())).(*sql.DB)
			ctx0, _ := context.WithTimeout(suiteTestCtx, duration)
			g, ctx := errgroup.WithContext(ctx0)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunTP(ctx, pkg.TPCommonErrHandler))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(wl.RunAP(ctx, pkg.APCommonErrHandler))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(pkg.AddDropColumn(ctx, db2, wl.DBName(), pkg.ErrorHandler(
					pkg.TiFlashStaleSchemaHandler, pkg.TiDBUnsupportedModifyColumnHandler, pkg.DDLWriteConflict, pkg.DDLDuplicateColumn, pkg.ColumnExists)))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				// change column type 消耗大量的时间，为了不影响其他 ddl，等待一段时间再执行.
				if pkg.Sleep(ctx, 3*time.Hour) {
					return
				}
				matcher.Try(pkg.ChangeColumnType(ctx, db2, wl.DBName(), pkg.ErrorHandler(
					pkg.TiFlashStaleSchemaHandler, pkg.TiDBUnsupportedModifyColumnHandler, pkg.DDLWriteConflict, pkg.DDLDuplicateColumn, pkg.ColumnExists)))
			}))
			// 为了不破坏其他负载的业务逻辑（如 ch），只在 360 partition 负载 truncate partition
			if wlType == workloads.Partition360WorkloadType {
				g.Go(util.WithGinkgoRecover(func() {
					matcher.Try(pkg.RandomTruncateTable(ctx, tidbClient, db2, wl.DBName(), pkg.ErrorHandler(
						pkg.TiFlashStaleSchemaHandler, pkg.TiDBUnsupportedModifyColumnHandler, pkg.DDLWriteConflict, pkg.DDLDuplicateColumn)))
				}))
			}
			checker.CommonCheck(ctx, tc, g, duration, wl.GetHistograms())
			handleTiFlashMode(ctx, tc, []string{wl.DBName()}, mode, false)
			if enableCompact {
				op := matcher.Try(chaos.NewCompactTiFlashReplica(tc)).(*chaos.CompactTiFlashReplica)
				matcher.Try(schedule.RunPeriodically(ctx, "@every 20m", 10*time.Minute, op, "randomCompactTiFlashReplica"))
			}
			matcher.Try(g.Wait())
		})

		// 目前一个 segment 默认最大是 1 G，假设平均 500 MB，默认 gc 时间间隔是 10 分钟，一次 gc 15 个 segment，所以可以按 10 min gc 7G 的进度来算
		It("should gc completely when set tiflash replica 0 #005#", func() {
			defer func() {
				checkBalance = false
			}()
			tbls := matcher.Try(tidbClient.ShowTableNames(wl.DBName())).([]string)
			totalSizeOnDisk := float64(0)
			for _, tbl := range tbls {
				size := matcher.Try(tidbClient.GetTiFlashSizeFromSystemTables(wl.DBName(), tbl, "total_stable_size_on_disk")).(uint64)
				totalSizeOnDisk += float64(size)
			}
			minutes := (int(totalSizeOnDisk/1024/1024/1024/7) + 1) * 10
			log.Info("wait gc", zap.Float64("tiflash stable size(GB) on disk", totalSizeOnDisk/1024/1024/1024),
				zap.Int("sleep minutes", minutes))
			matcher.Try(tidbClient.AlterBaseTablesSetTiFlashReplica(wl.DBName(), tbls, 0, ""))
			gomega.Eventually(func() uint64 {
				var total uint64
				for _, tbl := range tbls {
					size := matcher.Try(tidbClient.GetTiFlashSizeFromSystemTables(wl.DBName(), tbl, "total_size")).(uint64)
					total += size
				}
				return total
			}, time.Duration(minutes)*time.Minute, 1*time.Minute).Should(gomega.BeZero())
		})
	})
})

type tiflashMode string

// 1. normal. alter table xx set tiflash replica normal
// 2. fast. alter table xx set tiflash replica fast
// 3. conver. normal->fast->mormal->……
// 4. mix. some table fast and some table normal
// when set mode conver, mix or fast, we should set consistency param in workload false
// and we check consistency in the htap consistency cases
const (
	normal tiflashMode = "normal"
	fast   tiflashMode = "fast"
	conver tiflashMode = "conver"
	mix    tiflashMode = "mix"
)

// handleTiFlashMode is called in stable and ddl cases
func handleTiFlashMode(ctx context.Context, tc resource.TiDBCluster, dbs []string, mode tiflashMode, prepareStage bool) {
	tidbCli := tidb.NewClient(matcher.Try(tc.GetDB("")).(*sql.DB))
	switch mode {
	case normal, fast:
		for _, dbName := range dbs {
			tbls := matcher.Try(tidbCli.ShowTableNames(dbName)).([]string)
			log.Info("convert tiflash mode", zap.String("db name", dbName), zap.Strings("table names", tbls), zap.String("mode", string(mode)))
			matcher.Try(tidbCli.AlterBaseTablesSetTiFlashMode(dbName, tbls, string(mode)))
		}
	case conver:
		if !prepareStage {
			op := matcher.Try(chaos.NewConverTiFlashMode(tc, dbs)).(*chaos.ConverTiFlashMode)
			matcher.Try(schedule.RunPeriodically(ctx, "@every 20m", 10*time.Minute, op, "randomConverTiFlashMode"))
		}
	case mix:
		for _, dbName := range dbs {
			tbls := matcher.Try(tidbCli.ShowTableNames(dbName)).([]string)
			tbls2 := tbls[:len(tbls)/2]
			log.Info("convert tiflash mode", zap.String("db name", dbName), zap.Strings("table names", tbls), zap.String("mode", string(fast)))
			matcher.Try(tidbCli.AlterBaseTablesSetTiFlashMode(dbName, tbls2, string(fast)))
			tbls3 := tbls[len(tbls)/2:]
			log.Info("convert tiflash mode", zap.String("db name", dbName), zap.Strings("table names", tbls), zap.String("mode", string(normal)))
			matcher.Try(tidbCli.AlterBaseTablesSetTiFlashMode(dbName, tbls3, string(normal)))
		}
	}
}