package upgrade_test

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

	"github.com/coreos/go-semver/semver"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"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"

	"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/pkg/datagen"
	"github.com/pingcap/endless/testcase/tiflash/upgrade/query"
)

const (
	// when increase resource tiflash/tikv , the error don't raise
	// just record it for triage case in the future
	errRegionUnavailable     = "Region is unavailable"
	errRegionEpochNotMatch   = "Region epoch not match"
	errTaskHasBeenRegistered = "has been registered"
	errMppHangs              = "MPP Task canceled because it seems hangs"
	// 升级过程中, tidb 断开可能引发的错误
	errSocketConnTimeout = "socks connect tcp"
	errIOTimeout         = "i/o timeout"
	errConnRefused       = "connection refused"
	errConnTimeout       = "connection timed out"
	// 升级过程中, TiFlash 断开可能引发的错误
	// 没有完全列举出错误类型，可能还有其他错误
	errTiFlashServerIsTerminating = "TiFlash server is terminating"
	errTiFlashTimeout             = "TiFlash server timeout"
	errFailConnAllAddrs           = "failed to connect to all addresses"
)

const (
	beforeUpgrade = "beforce upgrade"
	upgrading     = "upgrading"
	afterUpgrade  = "afterUpgrade"
)

var errHandler = func(ignores ...string) func(error) error {
	return func(err error) error {
		for _, s := range ignores {
			if err != nil && strings.Contains(err.Error(), s) {
				return nil
			}
		}
		return err
	}
}
var consistencyErrHandlerMap map[string]func(error) error
var dmlErrHandler func(error) error

func initErrHandler() {
	dmlErrHandler = pkg.ErrorHandler(
		pkg.CtxDoneHandler,
		pkg.TiDBInformationSchemaOutDate,
		pkg.InvalidConnectionHandler,
		errHandler(
			errSocketConnTimeout,
			errIOTimeout,
			errConnRefused,
			errConnTimeout,
		))
	// 升级前后，一致性查询预期不报错
	consistencyErrHandlerMap = make(map[string]func(error) error)
	consistencyErrHandlerMap[beforeUpgrade] = pkg.ErrorHandler(
		pkg.TiDBInformationSchemaOutDate,
		errHandler(
			errRegionUnavailable,
			errRegionEpochNotMatch,
			errTaskHasBeenRegistered,
			errMppHangs))

	consistencyErrHandlerMap[afterUpgrade] = consistencyErrHandlerMap[beforeUpgrade]

	// 由于升级过程中，出现节点挂掉，有太多的难以列举全面的网络引发的错误，目前只处理不一致的错误，其余忽略
	consistencyErrHandlerMap[upgrading] = pkg.ErrorHandler(
		pkg.InconsistentHandler,
	)
}

var _ = Describe("TiFlash", func() {
	Context("upgrade", func() {
		ctx := suiteTestCtx
		var (
			tc                   resource.TiDBCluster
			dbName               = "test"
			replicas             int
			duration             time.Duration
			doUpgrade            bool
			prepareRows          int
			upgradeSrcVersion    string
			upgradeTargetVersion string
			upgradeTimeout       time.Duration
			ti                   *datagen.TableInfo
			upgradeStage         atomic.Value
			selectedQueries      []string
		)
		BeforeEach(func() {
			upgradeStage.Store(beforeUpgrade)
			initErrHandler()
			duration = matcher.Try(time.ParseDuration(ctx.Param("duration", "2h"))).(time.Duration)
			upgradeTimeout = matcher.Try(time.ParseDuration(ctx.Param("upgrade_timeout", "90m"))).(time.Duration)
			replicas = matcher.Try(strconv.Atoi(ctx.Param("replicas", "2"))).(int)
			prepareRows = matcher.Try(strconv.Atoi(ctx.Param("prepare_rows", "2000000"))).(int)
			doUpgrade = matcher.Try(strconv.ParseBool(ctx.Param("do_upgrade", "true"))).(bool)
			upgradeSrcVersion = ctx.Param("upgrade_source_version", "v5.4.0")
			upgradeTargetVersion = ctx.Param("upgrade_target_version", "master")
			selectedQueries = matcher.Try(strings.Split(ctx.Param("selected_queries", "q,q1,q2,q3,q4,q6,q12,q17,q18"), ",")).([]string)
			selectedQueries = getUpgradeQuery(upgradeSrcVersion, upgradeTargetVersion, selectedQueries)
			log.Info("after check version", zap.Strings("selected query", selectedQueries))
			tc = ctx.Resource("tc").(resource.TiDBCluster)
			db := matcher.Try(pkg.OpenDB(tc, dbName)).(*sql.DB)
			tidbCli := tidb.NewClient(db)
			matcher.Try(tidbCli.ShowClusterVersionInformation())
			ti = datagen.NewTableInfo(datagen.AutoIncrWideTable, db, false, true)
			matcher.Try(tidbCli.AlterBaseTablesSetTiFlashReplica(dbName, []string{ti.GetTableName()}, replicas, ""))
			gomega.Eventually(func() bool {
				return matcher.Try(tidbCli.IsTiFlashReplicaAvailable(dbName, []string{ti.GetTableName()})).(bool)
			}, 3*time.Minute, 3*time.Second).Should(gomega.BeTrue())
			count := matcher.Try(pkg.Query(suiteTestCtx, db, fmt.Sprintf("select count(*) from %s", ti.GetTableName()), false)).([]string)
			if matcher.Try(strconv.Atoi(count[0])).(int) > 0 {
				log.Info("skip prepare data", zap.String("rows count", count[0]))
				return
			}
			log.Info("prepare data begin", zap.Int("prepare rows", prepareRows))
			matcher.Try(ti.Run(&datagen.Request{
				Ctx: context.Background(),
				PrepareRequest: &datagen.PrepareRequest{
					TotalRows:   prepareRows,
					InsertBatch: 10,
					Threads:     10,
				}}))
			log.Info("prepare data done", zap.Int("prepare rows", prepareRows))
		})
		/*
			From junshen
			升级相关，tiflash 的冒烟测试流程包含下面几步应该会比较好：
			1. 在旧版本上创建表 t0 导入 1000 条数据，建 tiflash 副本，等 available，通过隔离读指定 tiflash 核对查询结果（条数、某个列的 sum）
			2. 升级 tiflash，并观察 tiflash 是否有 crash
			3. 再进行查询 t0，通过隔离读指定 tiflash 核对查询结果
			4. 往 t0 插入或更新一些数据，核对查询结果
			5. 创建另外一个表 t1 导入 1000 条数据，建 tiflash replica 等 available，通过隔离读指定 tiflash 核对查询结果
		*/
		// 期望不出现 invalid connection，socket timeout 等之外的问题
		It("should be consistency with data and available for the write and read #001", func() {
			ctx, cancel := context.WithTimeout(suiteTestCtx, duration)
			g, c := errgroup.WithContext(ctx)
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(ti.Run(&datagen.Request{
					Ctx: c,
					DmlReqs: []*datagen.DMLRequest{
						{
							DmlType:      datagen.Insert,
							Threads:      5,
							InsertBatch:  5,
							ErrorHandler: dmlErrHandler,
						},
						{
							DmlType:      datagen.Update,
							Threads:      1,
							Limit:        1,
							ErrorHandler: dmlErrHandler,
						},
						{
							DmlType:      datagen.Delete,
							Threads:      1,
							Limit:        1,
							ErrorHandler: dmlErrHandler,
						},
					}}))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				db2 := matcher.Try(pkg.OpenDB(tc, dbName)).(*sql.DB)
				for {
					qNo := rand.Intn(len(selectedQueries))
					// special, if query is `count(*)`, log the result
					logResult := false
					if qNo == 0 {
						logResult = true
					}
					log.Info("start query", zap.String("qno", selectedQueries[qNo]))
					err := pkg.DiffTiFlashTiKV(c, db2, fmt.Sprintf(query.Queries[selectedQueries[qNo]], ti.GetTableName()), true, logResult)
					if err != nil {
						log.Error("check consistence error", zap.Error(err), zap.String("sql", fmt.Sprintf(query.Queries[selectedQueries[qNo]], ti.GetTableName())))
						if pkg.CtxDoneHandler(err) == nil {
							return
						}
						matcher.Try(consistencyErrHandlerMap[upgradeStage.Load().(string)](err))
					} else {
						log.Info("end query", zap.String("qno", selectedQueries[qNo]))
					}
				}
			}))
			if doUpgrade {
				g.Go(util.WithGinkgoRecover(func() {
					defer cancel()
					log.Info("stable run 20 minutes before upgrade")
					time.Sleep(10 * time.Minute)
					log.Info("Begin to upgrade", zap.String("targetVersion", upgradeTargetVersion), zap.Duration("timeout", upgradeTimeout))
					upgradeStage.Store(upgrading)
					matcher.Try(tc.Upgrade(
						resource.TiDBClusterUpgradeOptions{
							Timeout: upgradeTimeout,
							Version: upgradeTargetVersion,
						}))
					log.Info("Upgrade succeeded", zap.String("targetVersion", upgradeTargetVersion), zap.Duration("timeout", upgradeTimeout))
					oltppkg.CheckClusterHealth(tc)
					upgradeStage.Store(afterUpgrade)
					// stable run after upgrade
					time.Sleep(10 * time.Minute)
				}))
			}
			matcher.Try(g.Wait())
		})
	})
})

// tiflash 5.3以前版本升级到 5.3 及以后版本的过程中,因为 exchange算子接口的优化导致 5.3 以前版本的 tiflash 的 exchange receiver  不能识别 5.3 及以后版本 的 tiflash 的 exchange sender发送的数据。
// 会出现数据不一致的情况
// 目前有两个解决方案:
// 1. 将5.3  exchange receiver 接口的代码 port 到 5.0/1/2.y 的小版本(或者改一下老 receiver)，然后升级的时候 5.0/1/2.x -> 5.0/1/2.y -> 5.3以后的版本，升级两次。
// 2. 升级过程中set tidb_allow_mpp = 0，等全部升级到 5.3 以后版本的时候，set tidb_allow_mpp = 1。损失是升级过程中原来走 mpp 的查询走不到 mpp，会变慢。
// https://asktug.com/t/topic/662840/23
// 判断如果前后版本可能引起不一致，则只使用 q 查询（只对比 count），否则全 query 跑
func getUpgradeQuery(upgradeSrcVersion, upgradeTargetVersion string, queries []string) []string {
	log.Info("upgrade version info", zap.String("upgrade source version", upgradeSrcVersion), zap.String("upgrade target version", upgradeTargetVersion))
	if upgradeSrcVersion == "master" || upgradeSrcVersion == "nightly" {
		return queries
	}
	srcV := semver.New(upgradeSrcVersion[1:])
	srcComp := srcV.Compare(*semver.New("5.3.0"))
	targetComp := 0
	if upgradeTargetVersion == "master" || upgradeTargetVersion == "nightly" {
		targetComp = 1
	} else {
		targetV := semver.New(upgradeTargetVersion[1:])
		targetComp = targetV.Compare(*semver.New("5.3.0"))
	}
	if srcComp < 0 && targetComp >= 0 {
		return []string{"q"}
	}
	return queries
}