package zto

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

	_ "github.com/go-sql-driver/mysql"
	"github.com/onsi/gomega"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/core"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	"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/workloadnode"
	"github.com/pingcap/endless/testcase/tiflash/pkg"
	"github.com/pingcap/endless/testcase/tiflash/pkg/checker"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload"
)

type ZtoWorkload struct {
	ctx                    core.TestContext
	tc                     resource.TiDBCluster
	olapWorkloads          []resource.WorkloadNode
	dbName                 string
	dbURL                  *url.URL
	prepareRows            int
	startDate              string
	olapThreads            int
	queriesUrl             string
	enableConsistencyCheck bool
	apConnParam            []string
	histograms             map[string]*checker.Histogram
}

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

func (w *ZtoWorkload) SetApThreads(threads int) {
	w.olapThreads = threads
}

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

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

const beforeDay = 2

func NewWorkload(ctx core.TestContext) workload.Workload {
	prepareRows := matcher.Try(strconv.Atoi(ctx.Param("prepare_rows", "100000000"))).(int)
	startDate := ctx.Param("start_date", time.Now().Add(-beforeDay*24*time.Hour).Format("2006-01-02"))
	olapThreads := matcher.Try(strconv.Atoi(ctx.Param("olap_threads", "1"))).(int)
	queriesUrl := ctx.Param("query_url", "http://fileserver.pingcap.net/download/pingcap/qa/tiflash/queries.tar.gz")
	enableConsistencyCheck := matcher.Try(strconv.ParseBool(ctx.Param("consistency", "true"))).(bool)

	tc := ctx.Resource("tc").(resource.TiDBCluster)
	// 4 个客户端高并发写入
	olapWorkloads := []resource.WorkloadNode{ctx.Resource("olapworkload1").(resource.WorkloadNode),
		ctx.Resource("olapworkload2").(resource.WorkloadNode),
		ctx.Resource("olapworkload3").(resource.WorkloadNode),
		ctx.Resource("olapworkload4").(resource.WorkloadNode)}
	dbURL := matcher.Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
	return &ZtoWorkload{
		ctx:                    ctx,
		tc:                     tc,
		olapWorkloads:          olapWorkloads,
		dbURL:                  dbURL,
		dbName:                 "rtdb",
		prepareRows:            prepareRows,
		startDate:              startDate,
		olapThreads:            olapThreads,
		queriesUrl:             queriesUrl,
		enableConsistencyCheck: enableConsistencyCheck,
		apConnParam:            pkg.OlapConParams,
	}
}

func (w *ZtoWorkload) Prepare(ctx context.Context) {
	tidbCli := tidb.NewClient(matcher.Try(w.tc.GetDB("")).(*sql.DB))
	if matcher.Try(tidbCli.IsTableExist(ctx, w.dbName, "zto_ssmx_bill_detail")).(bool) {
		log.Info("data is already exists and skip prepare")
		return
	}
	prepareCmd := fmt.Sprintf(`./zto -address %s -ratio 0 -update-thread 24 -insert-count %d -batch-insert 10 -trace-back-days 15 -trace-back-unsign-days 15 >> debug.log 2>&1`,
		w.dbURL.Host, w.prepareRows/4)
	g, ctx2 := errgroup.WithContext(ctx)
	for _, wn := range w.olapWorkloads {
		gwn := wn
		g.Go(func() error {
			cmd := workloadnode.CommandContext(ctx2, prepareCmd)
			defer cmd.Stop()
			return cmd.Run(gwn)
		})
	}
	g.Go(func() error {
		cmd := workloadnode.CommandContext(ctx2, fmt.Sprintf(`./zto -address %s -ratio 0 -update-thread 2 -insert-count %d -batch-insert 10 -trace-back-days 15 -trace-back-unsign-days 15 -prepare-baseorg true >> debug.log 2>&1`,
			w.dbURL.Host, w.prepareRows/4))
		defer cmd.Stop()
		return cmd.Run(w.olapWorkloads[0])
	})
	matcher.Try(g.Wait())
}

func (w *ZtoWorkload) RunTP(ctx context.Context, errorHandler func(error) error) error {
	runCmd := fmt.Sprintf(`./zto -address %s -ratio 0.8 -update-thread 24 -insert-count %d -batch-insert 10 -trace-back-days 15 -trace-back-unsign-days 15`,
		w.dbURL.Host, math.MaxUint32)
	g, ctx2 := errgroup.WithContext(ctx)
	for _, wn := range w.olapWorkloads {
		gwn := wn
		g.Go(func() error {
			cmd := workloadnode.CommandContext(ctx2, runCmd)
			defer cmd.Stop()
			err := cmd.Run(gwn)
			if err != nil {
				// kill 2
				if *(cmd.ExitCode) == 130 {
					return nil
				}
				return err
			}
			return nil
		})
	}
	if w.enableConsistencyCheck {
		g.Go(func() error {
			return w.checkConsistence(ctx2)
		})
	}
	return g.Wait()
}

func (w *ZtoWorkload) RunAP(ctx context.Context, errorHandler func(err error) error) error {
	matcher.Try(pkg.InstallArchive(ctx, w.queriesUrl, "./"))
	fileNames, queries, err := pkg.ReadFileNamesAndContentsFromDir("./queries")
	gomega.Expect(err).Should(gomega.BeNil())
	// TODO 偶尔会出现 Expression #1 of SELECT list is not in GROUP BY clause and contains nonaggregated column
	// 待确认
	conParam := append(w.apConnParam, "sql_mode=''")
	db := matcher.Try(pkg.OpenDB(w.tc, w.dbName, conParam...)).(*sql.DB)
	g, ctx2 := errgroup.WithContext(ctx)
	log.Info("test ap", zap.Int("query nums", len(queries)))
	for i := 0; i < w.olapThreads; i++ {
		g.Go(func() error {
			for {
				for i, v := range queries {
					log.Info("start query", zap.String("qname", fileNames[i]))
					taskStartT := time.Now()
					// 每次查询重新生成过滤条件，避免数据扫描越来越多
					startDate := time.Now().Add(-beforeDay * 24 * time.Hour).Format("2006-01-02")
					_, rows, err := pkg.Query(ctx2, db, strings.Replace(v, "$startDate", startDate, -1), true)
					taskConsumeT := time.Since(taskStartT)
					if err != nil {
						log.Error("end query", zap.String("qname", fileNames[i]), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Int("result count", rows), zap.Error(err))
						if pkg.CtxDoneHandler(err) == nil {
							return nil
						}
						if errorHandler != nil {
							err = errorHandler(err)
							if err != nil {
								log.Error("ap query exit for the error", zap.Error(err))
								return err
							}
						} else {
							return err
						}
					} else {
						log.Info("end query", zap.String("qname", fileNames[i]), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Int("result count", rows))
					}
				}
			}
		})
	}
	return g.Wait()
}

// 忽略检查中的非一致性错误
func (w *ZtoWorkload) checkConsistence(ctx context.Context) error {
	db := matcher.Try(w.tc.GetDB(w.dbName)).(*sql.DB)
	q := "select count(*) from zto_ssmx_bill_detail"
	for {
		err := pkg.DiffTiFlashTiKV(ctx, db, q, false, true)
		if err != nil {
			log.Error("execute query failed", zap.String("sql", q), zap.Error(err))
			if pkg.CtxDoneHandler(err) == nil {
				return nil
			}
			if pkg.InconsistentHandler(err) != nil {
				log.Error("check consistence exit for the error", zap.Error(err))
				return err
			}
		}
	}
}