package global_stats

import (
	"database/sql"
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/testcase/upgrade/utils"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type analyzeType string

var (
	sqls = []string{
		"create table rp%s (a bigint primary key, b varchar(20)) partition by range(a) (partition p0 values less than (10), partition p1 values less than (20), partition p2 values less than (30));",
		"create table hp%s (a bigint primary key, b varchar(20)) partition by hash(a) partitions 3;",
		"insert into rp%s values (1, 'a'), (2, 'b'), (3, 'c'), (14, 'd'), (16, 'e'), (18, 'f'), (20, 'h'), (25, 'i'), (27, 'j'), (29, 'k');",
		"insert into hp%s values (1, 'a'), (2, 'b'), (3, 'c'), (14, 'd'), (16, 'e'), (18, 'f'), (20, 'h'), (25, 'i'), (27, 'j'), (29, 'k');",
		"drop stats rp%s;",
		"drop stats hp%s;",
	}

	listPartition = []string{
		"create table lp%s (a bigint primary key, b varchar(20)) partition by list(a) (partition p0 values in (1,2,3), partition p1 values in (14,16,18), partition p3 values in (20,25,27,29));",
		"insert into lp%s values (1, 'a'), (2, 'b'), (3, 'c'), (14, 'd'), (16, 'e'), (18, 'f'), (20, 'h'), (25, 'i'), (27, 'j'), (29, 'k');",
		"drop stats lp%s;",
	}
	analyzeColumns   analyzeType = " columns a"
	analyzeOption    analyzeType = " with 10 buckets"
	analyzePartition analyzeType = " partition p0"

	AnalyzeTableWithoutConfig     analyzeType = ""
	AnalyzeTableWithConfig        analyzeType = analyzeOption
	AnalyzeTableWithColumns       analyzeType = analyzeColumns
	AnalyzeTableWithColumnsConfig analyzeType = analyzeColumns + analyzeOption

	AnalyzePartitionWithoutConfig     analyzeType = analyzePartition
	AnalyzePartitionWithConfig        analyzeType = analyzePartition + analyzeOption
	AnalyzePartitionWithColumns       analyzeType = analyzePartition + analyzeColumns
	AnalyzePartitionWithColumnsConfig analyzeType = analyzePartition + analyzeColumns + analyzeOption
)

func AnalyzeTable(db *sql.DB, suffixName string, analyzeOption analyzeType) {
	tables := []string{"rp", "hp"}
	var count string
	if err := db.QueryRow("show tables like 'lp%'").Scan(&count); err != nil {
		if err != sql.ErrNoRows {
			Expect(err).Should(BeNil(), "show tables failed.")
		}
	}
	if count != "" {
		tables = append(tables, "lp")
	}
	for _, tb := range tables {
		var version, mode string
		db.QueryRow("select @@tidb_analyze_version, @@tidb_partition_prune_mode").Scan(&version, &mode)
		sql := fmt.Sprintf("analyze table %s%s %s", tb, suffixName, analyzeOption)
		log.Info("execute sql", zap.String("sql", sql), zap.String("analyze_version", version), zap.String("partition_prune_mode", mode))
		for i := 0; i < 5; i++ {
			_, err := db.Exec(sql)
			if err != nil {
				if strings.Contains(err.Error(), "[stats]: some partition level statistics") {
					log.Info("ignore the error", zap.Error(err))
					time.Sleep(3 * time.Second)
					continue
				}
				Expect(err).Should(BeNil(), "analyze table failed with sql", sql)
			}
			break
		}
	}
}

func PrepareData(db *sql.DB, suffixName string) {
	var initSql []string
	initSql = append(initSql, sqls...)
	if CheckVersion(db, ">= 5.0.0") {
		initSql = append(initSql, listPartition...)
		matcher.Try(db.Exec("set session tidb_enable_list_partition = 1;"))
	}
	for _, sql := range initSql {
		sql = fmt.Sprintf(sql, suffixName)
		log.Info("execute sql", zap.String("sql", sql))
		_, err := db.Exec(sql)
		Expect(err).Should(BeNil(), "prepare data failed.")
	}
}

func PostCheck(db *sql.DB, suffixName string) {
	if CheckVersion(db, "< 6.1.0") {
		return
	}
	matcher.Try(db.Exec("set tidb_analyze_version = 2;"))
	matcher.Try(db.Exec("set tidb_partition_prune_mode = 'dynamic';"))
	var version, pruneMode string
	db.QueryRow("select @@tidb_analyze_version").Scan(&version)
	db.QueryRow("select @@tidb_partition_prune_mode").Scan(&pruneMode)
	log.Info("[PostCheck]", zap.String("analyze_version", version), zap.String("partition_prune_mode", pruneMode))
	AnalyzeTable(db, suffixName, AnalyzePartitionWithColumnsConfig)
	AnalyzeTable(db, suffixName, AnalyzeTableWithoutConfig)
	AnalyzeTable(db, suffixName, AnalyzeTableWithColumnsConfig)
}