package partition360

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

	. "github.com/onsi/gomega"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/core"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"

	"github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/endless/pkg/workloadnode/cmdstr"
	"github.com/pingcap/endless/testcase/oltp/scenario"
	"github.com/pingcap/endless/testcase/tiflash/pkg"
	"github.com/pingcap/endless/testcase/tiflash/pkg/checker"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload"
)

type Partition360Workload struct {
	ctx                    core.TestContext
	tc                     resource.TiDBCluster
	dbName                 string
	config                 scenario.Config
	subcmd                 cmdstr.GoycsbSubcmd
	workload360            resource.WorkloadNode
	replica                int
	smallTableDateInterval int
	apConnParam            []string
	histograms             map[string]*checker.Histogram
}

func NewWorkload(ctx core.TestContext) workload.Workload {
	dbName := ctx.Param("db_name", "partition360")
	tc := ctx.Resource("tc").(resource.TiDBCluster)
	workload360 := ctx.Resource("workload0").(resource.WorkloadNode)
	replica := Try(strconv.Atoi(ctx.Param("replica", "1"))).(int)
	config := scenario.DefaultConfig()
	// small 表 join 的日期区间
	smallTableDateInterval := Try(strconv.Atoi(ctx.Param("small_table_date_interval", "2"))).(int)
	// 360 负载并发读写数
	config.BettingOpts.Threads = Try(strconv.Atoi(ctx.Param("threads", "100"))).(int)
	config.QHAdOpts.DBName = dbName
	// 删除最老 partition 和新增 partition 相隔时间
	config.QHAdOpts.DropPartitionInterval.Duration = Try(time.ParseDuration(ctx.Param("drop_partition_interval", "20m"))).(time.Duration)
	// 每个 partition 的数据量 约等于 OperationCount * 50
	config.QHAdOpts.OperationCount = Try(strconv.ParseInt(ctx.Param("operation_count", "101"), 10, 64)).(int64)
	// 每个表保留多少个 partition
	config.QHAdOpts.InitialPartitionCount = Try(strconv.Atoi(ctx.Param("initial_partition_count", "10"))).(int)
	config.QHAdOpts.TableNamePrefix = dbName + "_"
	// 创建多少个表
	config.QHAdOpts.TableCount = Try(strconv.Atoi(ctx.Param("table_count", "10"))).(int)
	dbURL := Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
	config.QHAdOpts.MySQL.Host = dbURL.Hostname()
	config.QHAdOpts.MySQL.Port = dbURL.Port()
	// load 360 负载，生成数据
	subcmd := cmdstr.GoycsbLoad
	return &Partition360Workload{
		ctx:                    ctx,
		tc:                     tc,
		dbName:                 dbName,
		config:                 config,
		subcmd:                 subcmd,
		workload360:            workload360,
		replica:                replica,
		smallTableDateInterval: smallTableDateInterval,
		apConnParam:            pkg.OlapConParams,
	}
}

func (w *Partition360Workload) DBName() string {
	return w.dbName
}

func (w *Partition360Workload) SetApThreads(threads int) {
	// 	pass
}

func (w *Partition360Workload) SetApConnParam(apConnParam []string) {
	w.apConnParam = apConnParam
}

func (w *Partition360Workload) GetHistograms() map[string]*checker.Histogram {
	return w.histograms
}

func (w *Partition360Workload) Prepare(ctx context.Context) {
	tidbC := tidb.NewClient(Try(w.tc.GetDB("")).(*sql.DB))
	defer tidbC.Close()
	if Try(tidbC.IsDBExist(ctx, w.dbName)).(bool) {
		log.Info("data is already exists and skip prepare")
		return
	}
	log.Info("start prepare data")
	Try(runWorkload(ctx, w))
	log.Info("finish prepare data")
	db := Try(pkg.OpenDB(w.tc, w.dbName)).(*sql.DB)
	// 在 AP 查询时用到的 small 表
	Try(db.Exec("drop table if exists small"))
	Try(db.Exec("create table small(id int primary key auto_increment,`create_date` date NOT NULL DEFAULT '2000-01-01', `desc` varchar(100))"))
	tidbCli := tidb.NewClient(db)
	defer tidbCli.Close()
	log.Info("set tiflash replica")
	tables := Try(tidbCli.ShowBaseTableNames(w.dbName)).([]string)
	Try(tidbCli.AlterBaseTablesSetTiFlashReplica(w.dbName, tables, w.replica, ""))
	Eventually(func() bool {
		return Try(tidbCli.IsTiFlashReplicaAvailable(w.dbName, tables)).(bool)
	}, 120*time.Minute, 1*time.Minute).Should(BeTrue())
	log.Info("finish set tiflash replica")
	tidbCli2 := tidb.NewClient(Try(pkg.OpenDB(w.tc, w.dbName, "tidb_partition_prune_mode = 'dynamic'")).(*sql.DB))
	defer tidbCli2.Close()
	log.Info("start analyze table in dynamic prune mode")
	Try(tidbCli2.AnalyzeAllTablesInDB(w.dbName))
	log.Info("end analyze table in dynamic prune mode")
}

func (w *Partition360Workload) RunTP(ctx context.Context, errorHandler func(error) error) error {
	// 运行 360 负载，并行运行：
	// 1. 每间隔 droppartitioninterval 时间后，删除最老的 partition，创建新的 partition
	// 2. 从运行开始，自动切换到最新 partition，持续 insert on duplicate key 操作
	w.subcmd = cmdstr.GoycsbRun
	if err := runWorkload(ctx, w); err != nil {
		log.Error("error", zap.Error(err))
		if errorHandler(err) == nil {
			return nil
		}
		return err
	}
	return nil
}

func (w *Partition360Workload) RunAP(ctx context.Context, errorHandler func(err error) error) error {
	db := Try(pkg.OpenDB(w.tc, w.dbName, w.apConnParam...)).(*sql.DB)
	db2 := Try(pkg.OpenDB(w.tc, w.dbName)).(*sql.DB)
	tidbCli := tidb.NewClient(db)
	tables, err := tidbCli.ShowBaseTableNames(w.dbName)
	if err != nil {
		return err
	}
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}
		for _, table := range tables {
			if table == "small" {
				continue
			}
			if err = singleTableAggregation(ctx, db, table); err != nil {
				if err = errorHandler(err); err != nil {
					return err
				}
			}
			if err = selfJoin(ctx, db, table, tables); err != nil {
				if err = errorHandler(err); err != nil {
					return err
				}
			}
			if err = subquery(ctx, db, table, tables); err != nil {
				if err = errorHandler(err); err != nil {
					return err
				}
			}
			// 构造小维度表（表初始化行数 = 大表 partition 数量)
			// 默认值是 10 * 100 = 1000 行
			rowsNum := w.config.QHAdOpts.TableCount * w.config.QHAdOpts.InitialPartitionCount
			if err = smallDimensionTable(ctx, db, db2, w.dbName, table, rowsNum, w.smallTableDateInterval); err != nil {
				if err = errorHandler(err); err != nil {
					return err
				}
			}
		}
	}
}

func runWorkload(ctx context.Context, w *Partition360Workload) error {
	goycsb := cmdstr.NewGoycsb("/go-ycsb")
	if w.subcmd == cmdstr.GoycsbLoad {
		goycsb.Load("mysql")
	} else {
		goycsb.Run("mysql")
	}
	cmdString := goycsb.QHAdOptions(&w.config.QHAdOpts).String()
	cmd := workloadnode.CommandContext(ctx, cmdString)

	// HACK: Kill old go-ycsb process.
	if err := workloadnode.CommandContext(ctx, "pkill -9 go-ycsb || true").Run(w.workload360); err != nil {
		return nil
	}

	// Start client on workloadnode.
	if err := cmd.Start(w.workload360); err != nil {
		return err
	}
	defer cmd.Wait()
	return nil
}

func smallDimensionTable(ctx context.Context, db *sql.DB, db2 *sql.DB, dbName string, table string, rowsNum int, interval int) error {
	beginDate, endDate, err := randomGetTableCreateDateRange(db, table)
	if err != nil {
		return err
	}
	intervalBeginDate, intervalEndDate, err := randomGetTableCreateDateInterval(db, table, interval)
	if err != nil {
		return err
	}
	if err = prepareSmallTable(ctx, db2, dbName, rowsNum, table); err != nil {
		return err
	}
	// 大表 join 小表 on main.create_date = small.create_date
	sqls := []string{
		fmt.Sprintf("select count(*) from %s main join small on main.create_date = small.create_date where main.create_date > '%s' and main.create_date < '%s'", table, beginDate, endDate),
		fmt.Sprintf("select count(*) from %s main join small on main.create_date > small.create_date where main.create_date > '%s' and main.create_date < '%s'", table, intervalBeginDate, intervalEndDate),
		fmt.Sprintf("select count(*) from %s main join small on main.create_date < small.create_date where main.create_date > '%s' and main.create_date < '%s'", table, intervalBeginDate, intervalEndDate),
	}
	for _, sql2 := range sqls {
		log.Info("start query", zap.String("sql", sql2))
		taskStartT := time.Now()
		rows, _, err := pkg.Query(ctx, db, sql2, false)
		taskConsumeT := time.Since(taskStartT)
		if err != nil {
			log.Error("error", zap.String("sql type", "SmallDimensionTable"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows), zap.Error(err), zap.String("sql", sql2))
			return err
		} else {
			log.Info("end query", zap.String("sql type", "SmallDimensionTable"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows))
		}
	}
	return nil
}

func subquery(ctx context.Context, db *sql.DB, table string, subTables []string) error {
	beginDate, endDate, err := randomGetTableCreateDateRange(db, table)
	if err != nil {
		return err
	}
	for _, subTable := range subTables {
		if subTable == "small" {
			continue
		}
		s := fmt.Sprintf("select count(*) from %s m1 where aid in (select aid from %s m2 where m1.id = m2.id and m2.create_date > '%s' and m2.create_date < '%s') and create_date > '%s' and create_date < '%s'", table, subTable, beginDate, endDate, beginDate, endDate)
		log.Info("start query", zap.String("sql", s))
		taskStartT := time.Now()
		rows, _, err := pkg.Query(ctx, db, s, false)
		taskConsumeT := time.Since(taskStartT)
		if err != nil {
			log.Error("error", zap.String("sql type", "SingleTableAggregation"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows), zap.Error(err), zap.String("sql", s))
			return err
		} else {
			log.Info("end query", zap.String("sql type", "SingleTableAggregation"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows))
			return nil
		}
	}
	return nil
}

func selfJoin(ctx context.Context, db *sql.DB, table string, joinTables []string) error {
	beginDate, endDate, err := randomGetTableCreateDateRange(db, table)
	if err != nil {
		return err
	}
	sqls := []string{
		"select count(*) from %s m1 join %s m2 on m1.id = m2.id and m1.create_date = m2.create_date where m1.create_date > '%s' and m1.create_date < '%s'",
		"select count(*) from %s m1 join %s m2 on m1.id = m2.id and m1.create_date = m2.create_date where m2.create_date > '%s' and m2.create_date < '%s'",
		// "select count(*) from %s m1 join %s m2 on m1.id > m2.id and m1.create_date = m2.create_date where m1.create_date > '%s' and m1.create_date < '%s'",
		// "select count(*) from %s m1 join %s m2 on m1.id > m2.id and m1.create_date = m2.create_date where m2.create_date > '%s' and m2.create_date < '%s'",
	}
	for _, joinTable := range joinTables {
		if joinTable == "small" {
			continue
		}
		for _, sql2 := range sqls {
			s := fmt.Sprintf(sql2, table, joinTable, beginDate, endDate)
			log.Info("start query", zap.String("sql", s))
			taskStartT := time.Now()
			rows, _, err := pkg.Query(ctx, db, s, false)
			taskConsumeT := time.Since(taskStartT)
			if err != nil {
				log.Error("error", zap.String("sql type", "SelfJoin"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows), zap.Error(err), zap.String("sql", s))
				return err
			} else {
				log.Info("end query", zap.String("sql type", "SelfJoin"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows))
				return nil
			}
		}
	}
	return nil
}

func singleTableAggregation(ctx context.Context, db *sql.DB, table string) error {
	beginDate, endDate, err := randomGetTableCreateDateRange(db, table)
	if err != nil {
		return err
	}
	beginAid, endAid, err := randomGetTableAidRange(db, table)
	if err != nil {
		return err
	}
	s := fmt.Sprintf("select avg(m_req_all) from %s where create_date > '%s' and create_date < '%s' and aid >= %s and aid <= %s", table, beginDate, endDate, beginAid, endAid)
	log.Info("start query", zap.String("sql", s))
	taskStartT := time.Now()
	rows, _, err := pkg.Query(ctx, db, s, false)
	taskConsumeT := time.Since(taskStartT)
	if err != nil {
		log.Error("error", zap.String("sql type", "SingleTableAggregation"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows), zap.Error(err), zap.String("sql", s))
		return err
	} else {
		log.Info("end query", zap.String("sql type", "SingleTableAggregation"), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Strings("sql result", rows))
		return nil
	}
}

func randomGetTableCreateDateRange(db *sql.DB, table string) (begin, end string, err error) {
	rows, err := db.Query(fmt.Sprintf("select min(create_date), max(create_date) from %s", table))
	if err != nil {
		return begin, end, err
	}
	defer rows.Close()
	for rows.Next() {
		if err := rows.Scan(&begin, &end); err != nil {
			return begin, end, err
		}
	}
	return begin, end, rows.Err()
}

func randomGetTableCreateDateInterval(db *sql.DB, table string, interval int) (begin, end string, err error) {
	rows, err := db.Query(fmt.Sprintf("select distinct create_date from %s limit 1", table))
	if err != nil {
		return begin, end, err
	}
	for rows.Next() {
		if err := rows.Scan(&begin); err != nil {
			return begin, end, err
		}
	}
	rows, err = db.Query(fmt.Sprintf("select date_add('%s', interval %d day)", begin, interval))
	defer rows.Close()
	if err != nil {
		return begin, end, err
	}
	for rows.Next() {
		if err := rows.Scan(&end); err != nil {
			return begin, end, err
		}
	}
	return begin, end, rows.Err()
}

func randomGetTableAidRange(db *sql.DB, table string) (begin, end string, err error) {
	rows, err := db.Query(fmt.Sprintf("select min(aid), max(aid) from %s", table))
	if err != nil {
		return begin, end, err
	}
	defer rows.Close()
	for rows.Next() {
		if err = rows.Scan(&begin, &end); err != nil {
			return begin, end, err
		}
	}
	return begin, end, rows.Err()
}

// small 的 create_date 字段需要包含大表 main 的 create_date 的所有取值，其他 create_date 值可以上下浮动
func prepareSmallTable(ctx context.Context, db *sql.DB, dbName string, rowsNum int, table string) error {
	if _, err := db.Exec("truncate table small"); err != nil {
		return err
	}
	// 获取所有 create_date
	s := fmt.Sprintf("select distinct create_date from %s", table)
	dates, _, err := pkg.Query(ctx, db, s, false)
	if err != nil {
		return err
	}
	dateIndex := 0
	for i := 0; i < rowsNum; i++ {
		s = fmt.Sprintf("insert into small(`create_date`, `desc`) values('%s','test string')", dates[dateIndex])
		dateIndex = (dateIndex + 1) % len(dates)
		if _, err = db.Exec(s); err != nil {
			log.Error("insert data into small error", zap.Error(err))
			return err
		}
	}
	tidbCli := tidb.NewClient(db)
	Eventually(func() bool {
		return Try(tidbCli.IsTiFlashReplicaAvailable(dbName, []string{"small"})).(bool)
	}, 10*time.Minute, 1*time.Minute).Should(BeTrue())
	return nil
}