package ch

import (
	"context"
	"database/sql"
	"fmt"
	"math"
	"math/rand"
	"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"
)

/*
some known issues needed to workaround
1. https: // github.com/pingcap/tics/issues/1745
TiFlash can not be used immediately after br restored data
for 1k warehouses, wait 30 minutes
2. some function can't not be pushed to tiflash
Scalar function 'ascii'(signature: ASCII) can not be pushed to tiflash
Scalar function 'mod'(signature: ModInt) can not be pushed to tiflash
Error 1105: DB::TiFlashException: ModIntSignedSigned is not supported.
3. DECIMAL value is out of range in '(5, 2)' -- q19
q5,q7,q8,q9,q10,q11,q15,q16,q19,q20,q21
q1,q2,q3,q4,q6,q12,q13,q14,q17,q18,q22
4. q15 view skip
*/

type ChWorkload struct {
	ctx                    core.TestContext
	tc                     resource.TiDBCluster
	br                     resource.BR
	goTpc                  resource.WorkloadNode
	dbURL                  *url.URL
	warehouses             int
	brStorage              string
	brStorageWithCollation string
	dbName                 string
	tiflashReplicas        int
	olapThreads            int
	selectedQueries        string
	enableConsistencyCheck bool
	apConnParam            []string
	histograms             map[string]*checker.Histogram
}

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

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

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

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

func NewWorkload(ctx core.TestContext) workload.Workload {
	// dbName warehouses tiflashReplicas depends on the br storage
	brStorage := ctx.Param("br_storage", "s3://nfs/tiflash/ch-100-flash0?access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2fminio.pingcap.net%3a9000&force-path-style=true")
	brStorageWithCollation := ctx.Param("br_storage_with_collation", "s3://nfs/tiflash/ch-100-flash0-collation?access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2fminio.pingcap.net%3a9000&force-path-style=true")
	dbName := ctx.Param("db_name", "tpcc")
	warehouses := matcher.Try(strconv.Atoi(ctx.Param("warehouse", "100"))).(int)
	enableConsistencyCheck := matcher.Try(strconv.ParseBool(ctx.Param("consistency", "true"))).(bool)
	tiflashReplicas := matcher.Try(strconv.Atoi(ctx.Param("tiflash_replicas", "2"))).(int)
	olapThreads := matcher.Try(strconv.Atoi(ctx.Param("olap_threads", "1"))).(int)
	selectedQueries := ctx.Param("selected_queries", "q1,q3,q4,q6,q12,q13,q17,q18,q22")
	tc := ctx.Resource("tc").(resource.TiDBCluster)
	br := ctx.Resource("br").(resource.BR)
	goTpc := ctx.Resource("gotpc").(resource.WorkloadNode)
	dbURL := matcher.Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)

	return &ChWorkload{
		ctx:                    ctx,
		tc:                     tc,
		br:                     br,
		goTpc:                  goTpc,
		dbURL:                  dbURL,
		warehouses:             warehouses,
		brStorage:              brStorage,
		brStorageWithCollation: brStorageWithCollation,
		dbName:                 dbName,
		tiflashReplicas:        tiflashReplicas,
		olapThreads:            olapThreads,
		selectedQueries:        selectedQueries,
		enableConsistencyCheck: enableConsistencyCheck,
		apConnParam:            pkg.OlapConParams,
	}

}
func (w *ChWorkload) Prepare(ctx context.Context) {
	tidbCli := tidb.NewClient(matcher.Try(w.tc.GetDB("")).(*sql.DB))
	defer tidbCli.Close()
	if matcher.Try(tidbCli.IsTableExist(ctx, w.dbName, "warehouse")).(bool) {
		log.Info("data is already exists and skip prepare")
		return
	}
	pkg.MustRestoreWithBR(ctx, w.tc, w.br, w.brStorage, w.brStorageWithCollation)

	tbls := matcher.Try(tidbCli.ShowTableNames(w.dbName)).([]string)
	matcher.Try(tidbCli.AlterBaseTablesSetTiFlashReplica(w.dbName, tbls, w.tiflashReplicas, ""))
	gomega.Eventually(func() bool {
		return matcher.Try(tidbCli.IsTiFlashReplicaAvailable(w.dbName, tbls)).(bool)
	}, 3*time.Hour, 3*time.Second).Should(gomega.BeTrue())
	// it need sleep 10 minutes to wait for 100 tpcc after br
	sleepT := time.Duration(math.Sqrt(float64(w.warehouses))) * time.Minute
	log.Info("set tiflash replica and wait after br restore", zap.Duration("wait time", sleepT))
	time.Sleep(sleepT)
	log.Info("start analyze table in default prune mode")
	matcher.Try(tidbCli.AnalyzeAllTablesInDB(w.dbName))
	log.Info("end analyze table in default prune mode")

	tidbCli2 := tidb.NewClient(matcher.Try(pkg.OpenDB(w.tc, "", pkg.DynamicPartitionPruneMode)).(*sql.DB))
	defer tidbCli2.Close()
	log.Info("start analyze table in dynamic prune mode")
	matcher.Try(tidbCli2.AnalyzeAllTablesInDB(w.dbName))
	log.Info("end analyze table in dynamic prune mode")
}

func (w *ChWorkload) RunTP(ctx context.Context, errorHandler func(error) error) error {
	g, ctx2 := errgroup.WithContext(ctx)
	chTpCmd := fmt.Sprintf(`/go-tpc ch run --host %s -P %s --warehouses %d -T 32 -D %s  --acThreads 0 --ignore-error`, w.dbURL.Hostname(), w.dbURL.Port(), w.warehouses, w.dbName)
	cmd := workloadnode.CommandContext(ctx2, chTpCmd)
	pr := pkg.AssignCmdStdout(cmd)
	g.Go(func() error {
		defer cmd.Stop()
		return cmd.Run(w.goTpc)
	})
	g.Go(func() error {
		return pkg.InfoCmdStdout(ctx2, pr, errorHandler)
	})
	if w.enableConsistencyCheck {
		g.Go(func() error {
			return w.checkConsistence(ctx2)
		})
	}
	return g.Wait()
}

func (w *ChWorkload) RunAP(ctx context.Context, errorHandler func(err error) error) error {
	g, ctx2 := errgroup.WithContext(ctx)
	for i := 0; i < w.olapThreads; i++ {
		g.Go(func() error {
			return w.RunAPWithComment(ctx2, errorHandler, "")
		})
	}
	return g.Wait()
}

// 忽略检查中的非一致性错误
func (w *ChWorkload) checkConsistence(ctx context.Context) error {
	db := matcher.Try(w.tc.GetDB(w.dbName)).(*sql.DB)
	for {
		err := w.check(ctx, db)
		if err != nil {
			log.Error(err.Error())
			if pkg.CtxDoneHandler(err) == nil {
				return nil
			}
			if pkg.InconsistentHandler(err) != nil {
				return err
			}
		}
	}
}

/*
the mysql comment grammar that identify the query
*/
func (w *ChWorkload) RunAPWithComment(ctx context.Context, errorHandler func(error) error, comment string) error {
	// the longest duration is q18 which is about 1m32，the sort operator is not push down and it will consume some tidb memory
	qs := strings.Split(w.selectedQueries, ",")
	db := matcher.Try(pkg.OpenDB(w.tc, w.dbName, w.apConnParam...)).(*sql.DB)
	for {
		qNo := rand.Intn(len(qs))
		log.Info("start query", zap.String("qno", qs[qNo]))
		taskStartT := time.Now()
		_, rows, err := pkg.Query(ctx, db, comment+Queries[qs[qNo]], true)
		taskConsumeT := time.Since(taskStartT)
		if err != nil {
			log.Error("end query", zap.String("qno", qs[qNo]), zap.String("sql comment", comment), 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("qno", qs[qNo]), zap.String("sql comment", comment), zap.String("consume time", taskConsumeT.Round(time.Second).String()), zap.Int("result count", rows))
		}
	}
}