package tpch

import (
	"context"
	"database/sql"
	"strconv"
	"strings"
	"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"
	"golang.org/x/sync/errgroup"

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

type TpchWorkload struct {
	ctx                    core.TestContext
	tc                     resource.TiDBCluster
	br                     resource.BR
	factor                 int
	brStorage              string
	brStorageWithCollation string
	dbName                 string
	tiflashReplicas        int
	olapThreads            int
	apConnParam            []string
	histograms             map[string]*checker.Histogram
}

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

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

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

func (w *TpchWorkload) 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/tpch1-noflash2?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/tpch1-noflash-collation?access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2fminio.pingcap.net%3a9000&force-path-style=true")
	factor := matcher.Try(strconv.Atoi(ctx.Param("factor", "1"))).(int)
	dbName := ctx.Param("db_name", "tpch")
	tiflashReplicas := matcher.Try(strconv.Atoi(ctx.Param("tiflash_replicas", "2"))).(int)
	olapThreads := matcher.Try(strconv.Atoi(ctx.Param("olap_threads", "1"))).(int)
	tc := ctx.Resource("tc").(resource.TiDBCluster)
	br := ctx.Resource("br").(resource.BR)

	histograms := make(map[string]*checker.Histogram)
	for _, q := range Queries {
		histograms[q.Name] = checker.NewHistogram(q.Name)
	}
	return &TpchWorkload{
		ctx:                    ctx,
		tc:                     tc,
		br:                     br,
		brStorage:              brStorage,
		brStorageWithCollation: brStorageWithCollation,
		factor:                 factor,
		dbName:                 dbName,
		tiflashReplicas:        tiflashReplicas,
		olapThreads:            olapThreads,
		histograms:             histograms,
		apConnParam:            pkg.OlapConParams,
	}
}

func (w *TpchWorkload) Prepare(ctx context.Context) {
	tidbCli := tidb.NewClient(matcher.Try(w.tc.GetDB("")).(*sql.DB))
	if matcher.Try(tidbCli.IsDBExist(ctx, w.dbName)).(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)
	}, 60*time.Minute, 3*time.Second).Should(gomega.BeTrue())
	// it need sleep 5 minutes to wait for 1 factor tpch after br
	sleepT := time.Duration(w.factor) * 5 * 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")

	db := matcher.Try(w.tc.GetDB(w.dbName)).(*sql.DB)
	defer db.Close()
	_, err := db.Exec(createViewRevenue0)
	gomega.Expect(err != nil && !strings.Contains(err.Error(), "already exists")).Should(gomega.BeFalse())
}

func (w *TpchWorkload) RunTP(ctx context.Context, errorHandler func(error) error) error {
	return nil
}

func (w *TpchWorkload) RunAP(ctx context.Context, errorHandler func(err error) error) error {
	g, ctx2 := errgroup.WithContext(ctx)
	db := matcher.Try(pkg.OpenDB(w.tc, w.dbName, w.apConnParam...)).(*sql.DB)
	for i := 0; i < w.olapThreads; i++ {
		g.Go(func() error {
			for {
				for _, q := range Queries {
					log.Info("start Query", zap.String("Query Name", q.Name))
					taskStartT := time.Now()
					_, rows, err := pkg.Query(ctx2, db, q.Query, true)
					taskConsumeT := time.Since(taskStartT)
					w.histograms[q.Name].Update(taskConsumeT.Milliseconds())
					if err != nil {
						log.Error("end Query", zap.String("Query Name", q.Name), zap.String("consume time", taskConsumeT.Round(time.Millisecond).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("Query Name", q.Name), zap.String("consume time", taskConsumeT.Round(time.Millisecond).String()), zap.Int("result count", rows))
					}
				}
			}
		})
	}
	return g.Wait()
}