package window_function

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

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/endless/pkg/workloadnode/tiup"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s" // register k8s impl
	"go.uber.org/zap"

	"github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/testcase/tiflash/pkg"
)

var _ = Describe("WindowFunctionTest", func() {
	var (
		tc          resource.TiDBCluster
		tidbClient  *tidb.Client
		db          *sql.DB
		duration    time.Duration
		randomTimes int
	)
	BeforeEach(func() {
		tc = Try(suiteTestCtx.Resource("tc")).(resource.TiDBCluster)
		duration = Try(time.ParseDuration(suiteTestCtx.Param("duration", "1h"))).(time.Duration)
		db = Try(pkg.OpenDB(tc, "")).(*sql.DB)
		randomTimes = Try(strconv.Atoi(suiteTestCtx.Param("randomTimes", "1"))).(int)
		log.Info("create database and table")
		Try(db.Exec("drop database if exists window_function"))
		Try(db.Exec("create database window_function"))
		db = Try(pkg.OpenDB(tc, "window_function", pkg.MultiStatementsParam)).(*sql.DB)
		Try(db.Exec("create table window_function_0 (id int primary key auto_increment, i int, r real, s char(32) DEFAULT 'abc', dt datetime DEFAULT '1000-01-01 00:00:00')"))
		Try(db.Exec("create table window_function_0_part (id int primary key auto_increment, i int, r real, s char(32) DEFAULT 'abc', dt datetime DEFAULT '1000-01-01 00:00:00') partition by hash(id) partitions 5"))
		tidbClient = tidb.NewClient(db)
		log.Info("alter tiflash replica")
		Try(tidbClient.AlterDatabaseSetTiFlashReplica("window_function", 2, ""))
		Eventually(func() bool {
			return Try(tidbClient.IsDBTiFlashReplicaAvailable("window_function")).(bool)
		}, 10*time.Minute, 3*time.Second).Should(BeTrue())
		Try(tidbClient.ShowClusterVersionInformation())
		Try(prepareData(db))
	})
	Context("test window function", func() {
		It("sql will success #window_function_fine_graned_shuffle#", func() {
			ctx, _ := context.WithTimeout(suiteTestCtx, duration)
			for i := 0; i < randomTimes; i++ {
				Try(windowFunctionTest(ctx, db, false, true, true))
			}
		})

		It("sql will success #window_function_test1#", func() {
			ctx, _ := context.WithTimeout(suiteTestCtx, duration)
			for i := 0; i < randomTimes; i++ {
				Try(windowFunctionTest(ctx, db, false, false, false))
			}
		})
	})
})

// new tidb version and mixed tiflash version
var _ = Describe("MixedDeployWindowFunction", func() {
	var (
		db          *sql.DB
		duration    time.Duration
		randomTimes int
	)

	BeforeEach(func() {
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "3h"))).(time.Duration)
		randomTimes = Try(strconv.Atoi(suiteTestCtx.Param("randomTimes", "1"))).(int)
		deployMode := "mixed-version"
		db = deployClusterAndPrepareData(deployMode)
	})

	Context("test window function", func() {
		It("sql will success mixed_deploy_window_function_fine_graned_shuffle", func() {
			ctx, _ := context.WithTimeout(suiteTestCtx, duration)
			for i := 0; i < randomTimes; i++ {
				err := windowFunctionTest(ctx, db, true, true /*enable_fine_graned_shuffle*/, true /*random_batch_size*/)
				Expect(err).ToNot(BeNil())
			}
		})
	})
})

// old tidb version and new tiflash version
var _ = Describe("OldTidbWindowFunction", func() {
	var (
		db          *sql.DB
		duration    time.Duration
		randomTimes int
	)

	BeforeEach(func() {
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "3h"))).(time.Duration)
		randomTimes = Try(strconv.Atoi(suiteTestCtx.Param("randomTimes", "1"))).(int)
		deployMode := "patch-tiflash"
		db = deployClusterAndPrepareData(deployMode)
	})

	Context("test window function", func() {
		It("sql will success old_tidb_window_function_fine_graned_shuffle", func() {
			ctx, _ := context.WithTimeout(suiteTestCtx, duration)
			for i := 0; i < randomTimes; i++ {
				Try(windowFunctionTest(ctx, db, true, true /*enable_fine_graned_shuffle*/, true /*random_batch_size*/))
			}
		})
	})
})

// new tidb version and old tiflash version
var _ = Describe("OldTiflashWindowFunction", func() {
	var (
		db          *sql.DB
		duration    time.Duration
		randomTimes int
	)

	BeforeEach(func() {
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "1h"))).(time.Duration)
		randomTimes = Try(strconv.Atoi(suiteTestCtx.Param("randomTimes", "1"))).(int)
		deployMode := "patch-tidb"
		db = deployClusterAndPrepareData(deployMode)

	})

	Context("test window function", func() {
		It("sql will success old_tiflash_window_function_fine_graned_shuffle", func() {
			ctx, _ := context.WithTimeout(suiteTestCtx, duration)
			for i := 0; i < randomTimes; i++ {
				Try(windowFunctionTest(ctx, db, false, true /*enable_fine_graned_shuffle*/, true /*random_batch_size*/))
			}
		})
	})
})

func deployClusterAndPrepareData(deployMode string) *sql.DB {
	version := suiteTestCtx.Param("version", "v6.1.0")
	arch := suiteTestCtx.Param("arch", "amd64")
	enableTls := suiteTestCtx.Param("enable_tls", "false")
	tiupMirror := suiteTestCtx.Param("mirror", "product")
	tidbNum := matcher.Try(strconv.Atoi(suiteTestCtx.Param("tidb_num", "2"))).(int)
	tikvNum := matcher.Try(strconv.Atoi(suiteTestCtx.Param("tikv_num", "3"))).(int)
	tiflashNum := matcher.Try(strconv.Atoi(suiteTestCtx.Param("tiflash_num", "2"))).(int)

	tiupMirrorSetUp := fmt.Sprintf("tiup mirror set %s", productionMirror)
	tiupMirrorHttp := productionMirror
	if tiupMirror == "staging" {
		tiupMirrorSetUp = fmt.Sprintf("tiup mirror set %s", stagingMirror)
		tiupMirrorHttp = stagingMirror
	}

	controlNodeName := "tiup" // Node 的 hostname = $name-peer
	controlNode := suiteTestCtx.Resource(controlNodeName).(resource.WorkloadNode)

	log.Info(tiupMirrorSetUp)
	setCmd := workloadnode.CommandContext(suiteTestCtx, tiupMirrorSetUp)
	matcher.Try(setCmd.Run(controlNode))
	log.Info("Mirror setup")

	topoStr := renderTopo(tidbNum, tikvNum, tiflashNum)
	log.Info(topoStr)
	deployTopo := fmt.Sprintf(tiupTopoTpl, arch, enableTls, topoStr)
	tc := tiup.NewTiUPCluster(controlNode, controlNodeName+"-peer", deployTopo, clusterName)
	tryDeployAndStartCluster(tc, version)
	switch deployMode {
	case "patch-tidb":
		patchCluster(tc, tiupMirrorHttp, "tikv", "-R tikv")
		patchCluster(tc, tiupMirrorHttp, "tidb", "-R tidb")
	case "patch-tiflash":
		patchCluster(tc, tiupMirrorHttp, "tikv", "-R tikv")
		patchCluster(tc, tiupMirrorHttp, "tiflash", "-R tiflash")
	case "mixed-version":
		patchCluster(tc, tiupMirrorHttp, "tikv", "-R tikv")
		patchCluster(tc, tiupMirrorHttp, "tiflash", "-R tiflash -N 'tiflash-1-peer:9000'")
		patchCluster(tc, tiupMirrorHttp, "tidb", "-R tidb")
	}

	db := Try(pkg.OpenDB(tc, "")).(*sql.DB)
	log.Info("create database and table")
	Try(db.Exec("drop database if exists window_function"))
	Try(db.Exec("create database window_function"))
	db = Try(pkg.OpenDB(tc, "window_function", pkg.MultiStatementsParam)).(*sql.DB)
	Try(db.Exec("create table window_function_0 (id int primary key auto_increment, i int, r real, s char(32) DEFAULT 'abc', dt datetime DEFAULT '1000-01-01 00:00:00')"))
	Try(db.Exec("create table window_function_0_part (id int primary key auto_increment, i int, r real, s char(32) DEFAULT 'abc', dt datetime DEFAULT '1000-01-01 00:00:00') partition by hash(id) partitions 5"))
	tidbClient := tidb.NewClient(db)
	log.Info("alter tiflash replica")
	Try(tidbClient.AlterDatabaseSetTiFlashReplica("window_function", 2, ""))
	Eventually(func() bool {
		return Try(tidbClient.IsDBTiFlashReplicaAvailable("window_function")).(bool)
	}, 10*time.Minute, 3*time.Second).Should(BeTrue())
	Try(tidbClient.ShowClusterVersionInformation())
	Try(prepareData(db))
	return db
}

func patchCluster(tc *tiup.TiUPCluster, tiupMirrorHttp string, role string, flags ...string) {
	nightlyVersion, _ := tc.ControlNodeHost.Exec("tiup list tidb |grep \"nightly ->\" |awk '{print $3}'", 1*time.Minute)
	log.Info(nightlyVersion)
	tidbPatchName := fmt.Sprintf("%s/%s-%s-linux-amd64.tar.gz", tiupMirrorHttp, role, nightlyVersion)
	log.Info(tidbPatchName)
	Try(tc.Patch(tidbPatchName, 10*time.Minute, flags...))
}

func tryDeployAndStartCluster(tc *tiup.TiUPCluster, version string) {
	info, err := tc.Display()
	if err == nil && info.ClusterMeta.ClusterName == clusterName {
		tc.Destroy()
	}
	log.Info("cluster not deploy and start and try to deploy and start", zap.Error(err))
	matcher.Try(tc.Deploy(version, "--format", "json"))
	log.Info("Cluster deployed")
	matcher.Try(tc.Start("--format", "json"))
	log.Info("Cluster started")
}

func prepareData(db *sql.DB) error {
	log.Info("insert data")
	var query string
	for i := 0; i < 5; i++ {
		for _, int_value := range getIntRange() {
			for _, double_value := range getDoubleRange() {
				for _, string_value := range getStringRange() {
					for _, datetime_value := range getDatetimeRange() {
						query = fmt.Sprintf("insert into window_function_0 (i, r, s, dt) VALUES (%s, %s, %s, %s)", int_value, double_value, string_value, datetime_value)
						log.Info("exec sql", zap.String("sql", query))
						if _, err := db.Exec(query); err != nil {
							return err
						}
						query = fmt.Sprintf("insert into window_function_0_part (i, r, s, dt) VALUES (%s, %s, %s, %s)", int_value, double_value, string_value, datetime_value)
						log.Info("exec sql", zap.String("sql", query))
						if _, err := db.Exec(query); err != nil {
							return err
						}
					}
				}
			}
		}
	}
	// Split table, so there will be multiple tiflash nodes to do MPP.
	sql := `split table window_function_0 between (0) and (1000000) regions 5`
	log.Info(sql)
	Try(db.Exec(sql))
	sql = `split table window_function_0_part between (0) and (1000000) regions 5`
	log.Info(sql)
	Try(db.Exec(sql))
	return nil
}

func windowFunctionTest(ctx context.Context, db *sql.DB, isOldTidbVersion bool, enable_fine_graned_shuffle bool, random_batch_size bool) error {
	log.Info("Window Function")
	functions := []string{"row_number()", "rank()", "dense_rank()"}
	colType := []string{"i", "r", "s", "dt"}

	// Will change setting of fine grained shuffle if asked to.
	var streamCountSQL string
	var batchSizeSQL string
	changeFineGrainedShuffleSysVar := func() {
		if !isOldTidbVersion {
			if enable_fine_graned_shuffle {
				maxStreamCnt := 1024
				minStreamCnt := 1
				streamCountSQL = fmt.Sprintf("set @@tiflash_fine_grained_shuffle_stream_count = %d;",
					rand.Intn(maxStreamCnt-minStreamCnt)+minStreamCnt)

				if random_batch_size {
					maxBatchSize := 1000000
					minBatchSize := 1
					batchSizeSQL = fmt.Sprintf("set @@tiflash_fine_grained_shuffle_batch_size = %d;",
						rand.Intn(maxBatchSize-minBatchSize)+minBatchSize)
				}
			} else {
				streamCountSQL = "set @@tiflash_fine_grained_shuffle_stream_count = 0;"
				batchSizeSQL = "set @@tiflash_fine_grained_shuffle_batch_size = 1;"
			}
		}
	}

	tblNames := []string{"window_function_0", "window_function_0_part"}
	mppSet := "set @@session.tidb_allow_mpp=on;set @@session.tidb_enforce_mpp=1;"

	for _, function := range functions {
		for _, partition := range colType {
			for _, order := range colType {
				for _, tblName := range tblNames {
					changeFineGrainedShuffleSysVar()
					query := fmt.Sprintf(`select %s, %s, %s  over w from %s window w as
						(partition by %s order by %s) order by 1, 2, 3`,
						partition, order, function, tblName, partition, order)
					query = mppSet + streamCountSQL + batchSizeSQL + query
					log.Info("exec sql", zap.String("sql", query))
					if err := pkg.DiffTiFlashTiKV(ctx, db, query, true, false); err != nil {
						return err
					}

					for _, partition2 := range colType {
						for _, order2 := range colType {
							changeFineGrainedShuffleSysVar()
							query := fmt.Sprintf(`select %s, %s, %s, %s, %s  over w from %s window w as
								(partition by %s,%s order by %s,%s) order by 1, 2, 3, 4, 5`,
								partition, partition2, order, order2, function, tblName, partition, partition2, order, order2)
							query = mppSet + streamCountSQL + batchSizeSQL + query
							log.Info("exec sql", zap.String("sql", query))
							if err := pkg.DiffTiFlashTiKV(ctx, db, query, true, false); err != nil {
								return err
							}
						}
					}
				}
			}
		}
	}

	getRandomN := func(in []string) (res string, n int) {
		// [1,len(n)]
		n = rand.Intn(len(in)-1) + 1
		for i := 0; i < n; i++ {
			if i != 0 {
				res += ", "
			}
			ele := in[rand.Intn(len(in))]
			res += ele
		}
		return res, n
	}
	getOrderBy := func(n int) (res string) {
		for i := 1; i <= n; i++ {
			if i != 1 {
				res += ", "
			}
			res += fmt.Sprintf("%d", i)
		}
		res = "order by " + res
		return res
	}
	orders := "id"
	for i := 0; i < len(colType); i++ {
		orders += ", "
		orders += colType[i]
	}
	for _, tblName := range tblNames {
		for i := 0; i < 20; i++ {
			func1 := functions[rand.Intn(len(functions))]
			func2 := functions[rand.Intn(len(functions))]
			partitions1, nPart1 := getRandomN(colType)
			partitions2, nPart2 := getRandomN(colType)
			orderBy2 := getOrderBy(nPart1 + nPart2 + 2)
			changeFineGrainedShuffleSysVar()
			// [10, 100]
			limit := fmt.Sprintf("%d", rand.Intn(100-10)+10)
			offset := fmt.Sprintf("%d", rand.Intn(8000))
			// Two function.
			query := fmt.Sprintf(`select %s, %s, %s over w1, %s over w2 from %s window w1 as (partition by %s order by %s), w2 as (partition by %s order by %s) %s limit %s offset %s;`,
				partitions1, partitions2, func1, func2, tblName, partitions1, orders, partitions2, orders, orderBy2, limit, offset)
			query = mppSet + streamCountSQL + batchSizeSQL + query
			log.Info("exec sql", zap.String("sql", query))
			if err := pkg.DiffTiFlashTiKV(ctx, db, query, true, false); err != nil {
				return err
			}

			// GroupBy.
			groupByCol, projGroupByCol := func() (string, string) {
				n := rand.Intn(len(colType))
				c1 := colType[n]
				c2 := colType[(n+1)%len(colType)]
				return c1, c2
			}()
			query = fmt.Sprintf(`select %s over w1, count(%s) from %s group by %s having %s > 10 window w1 as (partition by %s order by %s) order by 1, 2 limit %s offset %s;`,
				func1, projGroupByCol, tblName, groupByCol, groupByCol, partitions1, orders, limit, offset)
			query = mppSet + streamCountSQL + batchSizeSQL + query
			log.Info("exec sql", zap.String("sql", query))
			if err := pkg.DiffTiFlashTiKV(ctx, db, query, true, false); err != nil {
				return err
			}
			// Join
			joinCondCol := colType[rand.Intn(len(colType))]
			query = fmt.Sprintf(`select %s over w1 from %s a join %s b on a.%s = b.%s where a.id < 1000 window w1 as (partition by a.%s order by a.id) order by 1 limit %s offset %s;`,
				func1, tblName, tblName, joinCondCol, joinCondCol, joinCondCol, limit, offset)
			query = mppSet + streamCountSQL + batchSizeSQL + query
			log.Info("exec sql", zap.String("sql", query))
			if err := pkg.DiffTiFlashTiKV(ctx, db, query, true, false); err != nil {
				return err
			}
		}
	}
	return nil
}

func getIntRange() []string {
	return []string{"NULL", "0", "1", "-1", "1.1", "-1.1", "2147483647", "2147483646", "-2147483648", "-2147483647", "1234567890", "-1234567890", "123456.7890", "-123456.7890", "1.0000000000000002"}
}

func getDoubleRange() []string {
	return []string{"NULL", "0.0", "1.1", "-1.1", "2147483647", "2147483646", "-2147483648", "-2147483647", "2.2250738585072014e-308", "-2.2250738585072014e-308", "123456.7890", "-123456.7890", "1.0000000000000002", "4.9406564584124654e-324", "2.2250738585072009e-308", "-2.2250738585072009e-308"}
}

func getStringRange() []string {
	return []string{"NULL", "'aaaaab'", "'basdbga'", "' qwe sdfg sd '", "'@!@%sdaf$@$'", "'%**^@!#^!@#'"}
}

func getDatetimeRange() []string {
	return []string{"NULL", "'1000-01-01 00:00:00'", "'1000-01-02 23:59:59'", "'9999-12-3 00:00:00'", "'9999-12-2 23:59:59'", "'1234-01-01 00:00:00'"}
}