package li_auto_test

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

	. "github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"github.com/pingcap/endless/testcase/tiflash/pkg"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	"github.com/pingcap/tidb/parser"
	"github.com/pingcap/tidb/parser/ast"
	"github.com/pingcap/tidb/parser/format"
	_ "github.com/pingcap/tidb/types/parser_driver"
	"github.com/rcrowley/go-metrics"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/matcher"
)

var tidbParser = parser.New()

const queriesDirName = "query"

func parseDbNames(s string) (res []string, err error) {
	stmtNodes, _, err := tidbParser.Parse(s, "", "")
	if err != nil {
		return
	}
	for _, stmt := range stmtNodes {
		switch stmt.(type) {
		case *ast.CreateDatabaseStmt:
			createDBStmt := stmt.(*ast.CreateDatabaseStmt)
			res = append(res, createDBStmt.Name)
		default:
		}
	}
	return
}

func splitSqls(s string) (res []string, err error) {
	stmtNodes, _, err := tidbParser.Parse(s, "", "")
	if err != nil {
		return
	}
	var sb strings.Builder
	for _, stmt := range stmtNodes {
		ctx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
		if err := stmt.Restore(ctx); err != nil {
			panic(err)
		}
		res = append(res, sb.String())
		sb.Reset()
	}
	return
}

type queriesExecuteRes []*queryExecuteInfo

func newQueriesExecuteRes(queryNames, queries []string) queriesExecuteRes {
	res := make([]*queryExecuteInfo, len(queries))
	for i, q := range queries {
		h := metrics.NewHistogram(metrics.NewExpDecaySample(1028, 0.015))
		res[i] = &queryExecuteInfo{
			queryId:   queryNames[i],
			query:     q,
			errors:    make(map[string]bool),
			plans:     make(map[string]string),
			histogram: h}
		metrics.MustRegister(q, h)
	}
	return res
}

func (o queriesExecuteRes) execute(ctx context.Context, cli *sql.DB, threads int, rounds int) error {
	g, _ := errgroup.WithContext(ctx)
	for i := 0; i < threads; i++ {
		g.Go(func() error {
			round := 0
			for round < rounds {
				for _, q := range o {
					err := q.executeQueryAndRecordMetrics(ctx, cli)
					if err != nil {
						if pkg.CtxDoneHandler(err) == nil {
							return nil
						}
						q.errors[err.Error()] = true
					}
				}
				round++
			}
			return nil
		})
	}
	return g.Wait()
}

type queryExecuteInfo struct {
	queryId   string
	query     string
	errors    map[string]bool
	plans     map[string]string
	histogram metrics.Histogram
}

func (o *queryExecuteInfo) getPlans(cli *sql.DB) {
	q := fmt.Sprintf(`select PLAN_DIGEST,first_value(plan) over (partition by plan_digest) from information_schema.statements_summary_history where query_sample_text="%s")`, o.query)
	var planDigest, plan string
	rows, err := cli.Query(q)
	if err != nil {
		return
	}
	defer rows.Close()
	for rows.Next() {
		if err := rows.Scan(&planDigest, &plan); err != nil {
			log.Error("Sql error",
				zap.String("sql", q),
				zap.Error(err))
			continue
		}
		o.plans[planDigest] = plan
	}
}

func (o *queryExecuteInfo) executeQueryAndRecordMetrics(ctx context.Context, cli *sql.DB) error {
	start := time.Now()
	rows, err := cli.QueryContext(ctx, o.query)
	defer func() {
		if rows != nil {
			rows.Close()
		}
		o.histogram.Update(int64(time.Since(start) / time.Microsecond))
	}()
	if err != nil {
		return err
	}
	return nil
}

func executeSingleStmt(ctx context.Context, cli *tidb.Client, s string, errHandler func(error) error) error {
	res, err := splitSqls(s)
	if err != nil {
		return err
	}
	log.Info("execute sql numbers", zap.Int("num", len(res)))
	for i, r := range res {
		_, err := cli.ExecContext(ctx, r)
		log.Info("execute progress", zap.String("current", fmt.Sprintf("%d/%d", i+1, len(res))))
		if errHandler != nil && errHandler(err) != nil {
			return err
		}
	}
	return nil
}

var _ = Describe("LiAuto", func() {
	var duration time.Duration
	var rounds int
	var tc resource.TiDBCluster
	var db *sql.DB
	var tidbClient *tidb.Client
	BeforeEach(func() {
		duration = matcher.Try(time.ParseDuration(suiteTestCtx.Param("duration", "1h"))).(time.Duration)
		rounds = matcher.Try(strconv.Atoi(suiteTestCtx.Param("rounds", "3"))).(int)
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		db = matcher.Try(tc.GetDB("test")).(*sql.DB)
		tidbClient = tidb.NewClient(db)
	})
	Context("when execute ap query with no data in schema", func() {
		It("shouldn't fail", func() {
			dbNames := matcher.Try(parseDbNames(createDBSqls)).([]string)
			queryNames, queries, err := pkg.ReadFileNamesAndContentsFromEmbedFs(queryFS, queriesDirName)
			// readSqlsFromDir()
			gomega.Expect(err).Should(gomega.BeNil())
			dbsExists := true
			for _, dbName := range dbNames {
				dbsExists = dbsExists && matcher.Try(tidbClient.IsDBExist(suiteTestCtx, dbName)).(bool)
				if !dbsExists {
					break
				}
			}
			if !dbsExists {
				log.Info("dbs are not exists,need create")
				matcher.Try(executeSingleStmt(suiteTestCtx, tidbClient, createDBSqls, nil))
				matcher.Try(executeSingleStmt(suiteTestCtx, tidbClient, createTableSqls, nil))
				matcher.Try(executeSingleStmt(suiteTestCtx, tidbClient, createViewSqls, nil))
				for _, dbName := range dbNames {
					tbls := matcher.Try(tidbClient.ShowTableNames(dbName)).([]string)
					matcher.Try(tidbClient.AlterBaseTablesSetTiFlashReplica(dbName, tbls, 2, ""))
					gomega.Eventually(func() bool {
						return matcher.Try(tidbClient.IsTiFlashReplicaAvailable(dbName, tbls)).(bool)
					}, 20*time.Minute, 3*time.Second).Should(gomega.BeTrue())
				}
			}
			log.Info("start execute query")
			executeRes := newQueriesExecuteRes(queryNames, queries)
			conParam := append(pkg.OlapConParams, "sql_mode=''")
			db2 := matcher.Try(pkg.OpenDB(tc, "test", conParam...)).(*sql.DB)
			ctx, _ := context.WithTimeout(context.Background(), duration)
			matcher.Try(executeRes.execute(ctx, db2, 1, rounds))
			for _, r := range executeRes {
				log.Info("execute info",
					zap.String("query id", r.queryId),
					zap.Any("errors", r.errors))
			}
			for _, r := range executeRes {
				gomega.Expect(len(r.errors)).Should(gomega.BeZero())
			}
		})
	})
})
