package high_availablity_test

import (
	"context"
	"database/sql"
	"fmt"
	"math/rand"
	"sync/atomic"
	"time"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/testcase/tiflash/pkg"
	"github.com/pingcap/endless/testcase/tiflash/pkg/chaos"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/ch"
)

/*
>=release-5.1
*/
const times = 2

func runChaos(ctx context.Context, tc resource.TiDBCluster, tidbClient *tidb.Client, stopF func(), startF func()) {
	opt := resource.TiDBClusterSelectorOpts{
		TiDBClusterTiFlashSelectorSpec: &resource.TiDBClusterTiFlashSelectorSpec{
			All: true,
		},
	}
	selector := matcher.Try(tc.ChaosSelect(opt)).([]resource.ChaosSelector)
	var cID resource.ChaosIdentifier
	defer chaos.DeleteChaos(tc, cID)
	for i := 0; i < times; i++ {
		s := rand.Intn(len(selector))
		log.Info("begin create chaos to stop tiflash node", zap.Int("round", i+1),
			zap.String("node name", selector[s].HumanReadableName()))
		cID = matcher.Try(tc.ChaosClient().Create(resource.InjectChaosOptions{
			Selector:       []resource.ChaosSelector{selector[s]},
			ExperimentSpec: &resource.FailureExperimentSpec{},
		})).(resource.ChaosIdentifier)
		// create chaos 是一个异步操作，不是马上就生效的，此时设置标记位后，tiflash 还没有 stop
		// 基于 tiflash 不 stop 的时候查询没问题的前提（预先跑了几分钟 query），这种情况应该问题不大
		stopF()
		gomega.Eventually(func() []bool {
			return matcher.Try(tidbClient.IsTiFlashStoreStatusUp(context.Background(), false, "", selector[s].HumanReadableName())).([]bool)
		}, 10*time.Minute, 3*time.Second).ShouldNot(gomega.ContainElement(true))

		log.Info("end stop tiflash node", zap.Int("round", i+1),
			zap.String("node name", selector[s].HumanReadableName()))
		if pkg.Sleep(ctx, 5*time.Minute) {
			matcher.Try(chaos.DeleteChaos(tc, cID))
			return
		}
		log.Info("begin delete chaos to start tiflash node", zap.Int("round", i+1),
			zap.String("node name", selector[s].HumanReadableName()))
		matcher.Try(chaos.DeleteChaos(tc, cID))
		gomega.Eventually(func() []bool {
			return matcher.Try(tidbClient.IsTiFlashStoreStatusUp(context.Background(), false, "", selector[s].HumanReadableName())).([]bool)
		}, 10*time.Minute, 3*time.Second).ShouldNot(gomega.ContainElement(false))
		log.Info("end start tiflash node", zap.Int("round", i+1),
			zap.String("node name", selector[s].HumanReadableName()))
		startF()
		if pkg.Sleep(ctx, 5*time.Minute) {
			return
		}
	}
}

/*
case1
	单 tiflash replica
	stop 节点，查询失败，start 节点查询恢复
*/
var _ = Describe("SingleTiFlashReplicaRestartNode", func() {
	var tc resource.TiDBCluster
	var chWorkload *ch.ChWorkload
	var isStop atomic.Value
	var errCount uint32
	var tidbClient *tidb.Client
	var duration time.Duration
	var skipErrHandler = func(err error) error {
		// 节点 stop; 节点 start 且失败次数小于 1(tiflash stop -> query start -> tiflash start -> query end)
		if isStop.Load().(bool) || (!isStop.Load().(bool) && atomic.LoadUint32(&errCount) < 1) {
			if !isStop.Load().(bool) && atomic.LoadUint32(&errCount) < 1 {
				atomic.AddUint32(&errCount, 1)
			}
			return nil
		}
		log.Info("error not skip", zap.Bool("isstop", isStop.Load().(bool)), zap.Uint32("errcount", atomic.LoadUint32(&errCount)))
		return err
	}
	var stopF = func() {
		isStop.Store(true)
		atomic.StoreUint32(&errCount, 0)
	}
	var startF = func() {
		isStop.Store(false)
		atomic.StoreUint32(&errCount, 0)
	}
	BeforeEach(func() {
		isStop.Store(false)
		atomic.StoreUint32(&errCount, 0)
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "5h"))).(time.Duration)
		chWorkload = ch.NewWorkload(suiteTestCtx).(*ch.ChWorkload)
		chWorkload.Prepare(suiteTestCtx)
		db := matcher.Try(tc.GetDB(chWorkload.DBName())).(*sql.DB)
		tidbClient = tidb.NewClient(db)
	})
	Context("when stop & start one tiflash node", func() {

		It("should have high availability", func() {
			log.Info("SingleTiFlashReplicaHighAvailability")
			tbls := matcher.Try(tidbClient.ShowTableNames(chWorkload.DBName())).([]string)
			matcher.Try(tidbClient.AlterBaseTablesSetTiFlashReplica(chWorkload.DBName(), tbls, 1, ""))
			gomega.Eventually(func() bool {
				return matcher.Try(tidbClient.IsTiFlashReplicaAvailable(chWorkload.DBName(), tbls)).(bool)
			}, 60*time.Minute, 3*time.Second).Should(gomega.BeTrue())
			log.Info("set tiflash replica and wait 20 minutes after replica change")
			time.Sleep(20 * time.Minute)
			ctx, cancel := context.WithTimeout(suiteTestCtx, duration)
			g, _ := errgroup.WithContext(ctx)
			g.Go(util.WithGinkgoRecover(func() {
				defer cancel()
				matcher.Try(chWorkload.RunAP(ctx, skipErrHandler))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				defer cancel()
				time.Sleep(5 * time.Minute)
				runChaos(ctx, tc, tidbClient, stopF, startF)
			}))
			matcher.Try(g.Wait())
		})
	})
})

/*
case2
	2 tiflash replica，stop 一个节点，单并发查询最多挂一次，之后查询稳定可用
	等待一段时间后，start 上述节点，查询可用性不受影响
*/
var _ = Describe("MultipleTiFlashReplicaRestartNode", func() {
	var tc resource.TiDBCluster
	var chWorkload *ch.ChWorkload
	var tidbClient *tidb.Client
	var duration time.Duration
	var isStop atomic.Value
	var errCount uint32
	var skipErrHandler = func(err error) error {
		// 如果节点 stop , 且 query 失败次数小于 1
		if isStop.Load().(bool) && atomic.LoadUint32(&errCount) < 1 {
			atomic.AddUint32(&errCount, 1)
			return nil
		}
		log.Info("error not skip", zap.Bool("isstop", isStop.Load().(bool)), zap.Uint32("errcount", atomic.LoadUint32(&errCount)))
		return err
	}
	var stopF = func() {
		isStop.Store(true)
		atomic.StoreUint32(&errCount, 0)
	}
	var startF = func() {
		isStop.Store(false)
		atomic.StoreUint32(&errCount, 0)
	}
	BeforeEach(func() {
		isStop.Store(false)
		atomic.StoreUint32(&errCount, 0)
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "5h"))).(time.Duration)
		chWorkload = ch.NewWorkload(suiteTestCtx).(*ch.ChWorkload)
		chWorkload.Prepare(suiteTestCtx)
		db := matcher.Try(tc.GetDB(chWorkload.DBName())).(*sql.DB)
		tidbClient = tidb.NewClient(db)
	})
	Context("when stop & start one tiflash node", func() {
		It("should have high availability", func() {
			log.Info("MultipleTiFlashReplicaHighAvailability")
			// As the interval time is much larger than query time and no write , there are 2 situations that may be happened
			// 1. query start -> tiflash stop -> query end -> query start and end -> query start -> query end -> tiflash start
			// 																	     query start -> tiflash start -> query end
			// 2. tiflash stop -> query start -> query end -> query start and end -> query start -> query end -> tiflash start
			// 																	     query start -> tiflash start -> query end
			// the query will be failed in situation1 and not be failed in situation2. we simple the checker that the query can only be failed once at most.
			ctx, cancel := context.WithTimeout(suiteTestCtx, duration)
			g, _ := errgroup.WithContext(ctx)
			g.Go(util.WithGinkgoRecover(func() {
				defer cancel()
				matcher.Try(chWorkload.RunAP(ctx, skipErrHandler))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				defer cancel()
				time.Sleep(5 * time.Minute)
				runChaos(ctx, tc, tidbClient, stopF, startF)
			}))
			matcher.Try(g.Wait())
		})
	})
})

/*
case3
	2 tiflash replica，stop 一个 az，查询会失败，之后稳定可用
	start az，查询稳定可用
*/
var _ = Describe("MultipleTiFlashReplicaRestartAZ", func() {
	var duration time.Duration
	var chWorkload *ch.ChWorkload
	var tc resource.TiDBCluster
	var cID resource.ChaosIdentifier
	var tidbClient *tidb.Client
	var db *sql.DB
	BeforeEach(func() {
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "5h"))).(time.Duration)
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		chWorkload = ch.NewWorkload(suiteTestCtx).(*ch.ChWorkload)
		chWorkload.Prepare(suiteTestCtx)
		db = matcher.Try(tc.GetDB(chWorkload.DBName())).(*sql.DB)
		tidbClient = tidb.NewClient(db)
	})
	Context("when stop & start one az", func() {
		It("should have high availability", func() {
			log.Info("MultipleAZMultipleTiFlashReplicaHighAvailability")
			ctx2, cancel := context.WithTimeout(suiteTestCtx, duration)
			defer cancel()
			zoneLabels := matcher.Try(tidbClient.GetTiFlashZoneStoreIdsMap(ctx2)).(map[string][]string)
			if len(zoneLabels) < 3 {
				Skip(fmt.Sprintf("TiFlash zone num(%d) is less than 3", len(zoneLabels)))
			}
			tbls := matcher.Try(tidbClient.ShowTableNames(chWorkload.DBName())).([]string)
			matcher.Try(tidbClient.AlterBaseTablesSetTiFlashReplica(chWorkload.DBName(), tbls, 2, "'zone'"))
			gomega.Eventually(func() bool {
				return matcher.Try(tidbClient.IsTiFlashReplicaAvailable(chWorkload.DBName(), tbls)).(bool)
			}, 60*time.Minute, 3*time.Second).Should(gomega.BeTrue())
			var zones [3]string
			var i int
			for k, _ := range zoneLabels {
				zones[i] = k
				i++
				if i >= 3 {
					break
				}
			}
			gomega.Eventually(func() bool {
				return matcher.Try(tidbClient.DoTablesHaveMutipleTiFlashReplicaInEachZone(ctx2, chWorkload.DBName(), tbls)).(bool)
			}, 30*time.Minute, 3*time.Second).ShouldNot(gomega.BeTrue())
			// query is normal before start chaos
			ctx1, _ := context.WithTimeout(context.Background(), 5*time.Minute)
			matcher.Try(chWorkload.RunAP(ctx1, nil))
			for i := 0; i < times; i++ {
				zone := zones[rand.Intn(len(zones))]
				// zone = "zone1"
				opt := resource.TiDBClusterSelectorOpts{
					TiDBClusterTiFlashSelectorSpec: &resource.TiDBClusterTiFlashSelectorSpec{
						StoreLabelSelector: &metav1.LabelSelector{
							MatchLabels: map[string]string{
								"zone": zone,
							},
						},
					},
				}
				selector := matcher.Try(tc.ChaosSelect(opt)).([]resource.ChaosSelector)
				log.Info("begin stop tiflash az node",
					zap.String("selector", zone), zap.Int("round", i+1))
				cID = matcher.Try(tc.ChaosClient().Create(resource.InjectChaosOptions{
					Selector:       selector,
					ExperimentSpec: &resource.FailureExperimentSpec{},
				})).(resource.ChaosIdentifier)
				// 等待该 zone tiflash 真正 stop 后, 集群可以稳定
				gomega.Eventually(func() []bool {
					return matcher.Try(tidbClient.IsTiFlashStoreStatusUp(context.Background(), false, zone, "")).([]bool)
				}, 10*time.Minute, 3*time.Second).ShouldNot(gomega.ContainElement(true))
				log.Info("end stop tiflash az node",
					zap.String("selector", zone), zap.Int("round", i+1))
				time.Sleep(10 * time.Minute)
				ctx3, _ := context.WithTimeout(context.Background(), 10*time.Minute)
				matcher.Try(chWorkload.RunAP(ctx3, nil))
				log.Info("start tiflash az node",
					zap.String("selector", zone), zap.Int("round", i+1))
				matcher.Try(chaos.DeleteChaos(tc, cID))
				// 启动两个 tiflash 节点的时间可能较长，此处设置 20 minutes
				ctx4, _ := context.WithTimeout(context.Background(), 20*time.Minute)
				matcher.Try(chWorkload.RunAP(ctx4, nil))
			}
		})
	})
	AfterEach(func() {
		chaos.DeleteChaos(tc, cID)
		log.Info("after each delete cid")
	})
})