package replica_manager_test

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

	. "github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"github.com/pingcap/log"
	"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/testcase/tiflash/pkg"
)

var _ = Describe("ClusterManager", func() {
	var (
		dbName                            string
		tiflashReplicas                   int
		tc                                resource.TiDBCluster
		tidbCli                           *tidb.Client
		setAndWaitTiFlashReplicaAvailable func(int)
		selectCountFromRandomTable        func(context.Context, *tidb.Client)
	)
	Context("alter database tiflash replica", func() {
		BeforeEach(func() {
			// data from sql coverage comp, about 30K tables and contains some partition tables
			// alter database coms set tiflash replica 2;
			// Query OK, 0 rows affected (3 hours 15 min 1.26 sec)
			// In total 29920 tables: 29920 succeed, 0 failed, 0 skipped
			brStorage := suiteTestCtx.Param("br_storage", "s3://nfs/tiflash/comp-noflash?access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2fminio.pingcap.net%3a9000&force-path-style=true")
			brStorageWithCollation := suiteTestCtx.Param("br_storage_with_collation", "s3://nfs/tiflash/comp-noflash-collation?access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2fminio.pingcap.net%3a9000&force-path-style=true")
			dbName = suiteTestCtx.Param("db_name", "coms")
			tiflashReplicas = matcher.Try(strconv.Atoi(suiteTestCtx.Param("tiflash_replicas", "2"))).(int)
			// the number of tiflash replica in the data to be recovered
			tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
			tidbCli = tidb.NewClient(matcher.Try(tc.GetDB("")).(*sql.DB))
			setAndWaitTiFlashReplicaAvailable = func(replica int) {
				log.Info("set tiflash replica and wait available", zap.Int("replica", replica))
				matcher.Try(tidbCli.AlterDatabaseSetTiFlashReplica(dbName, replica, ""))
				log.Info("ddl done", zap.Int("replica", replica))
				gomega.Eventually(func() bool {
					return matcher.Try(tidbCli.IsDBTiFlashReplicaAvailable(dbName)).(bool)
				}, 5*time.Hour, 3*time.Minute).Should(gomega.BeTrue())
				log.Info("all table available", zap.Int("replica", replica))
			}
			selectCountFromRandomTable = func(ctx context.Context, cli *tidb.Client) {
				log.Info("do some select count(*) through tiflash engine")
				tableNames := matcher.Try(cli.ShowBaseTableNames(dbName)).([]string)
				for _, n := range tableNames {
					q := fmt.Sprintf("select /*+ read_from_storage(TIFLASH[%s]) */ count(*) from %s ", n, n)
					_, err := cli.ExecContext(ctx, q)
					log.Info("select", zap.String("sql", q))
					if err != nil {
						log.Error("sql error", zap.String("sql", q), zap.Error(err))
						if pkg.CtxDoneHandler(err) == nil {
							return
						}
						gomega.Expect(err).Should(gomega.BeNil())
					}
				}
			}
			if matcher.Try(tidbCli.IsDBExist(suiteTestCtx, dbName)).(bool) {
				return
			}
			br := suiteTestCtx.Resource("br").(resource.BR)
			pkg.MustRestoreWithBR(suiteTestCtx, tc, br, brStorage, brStorageWithCollation)
		})
		It("should success when br restore lots of tables #001#", func() {
			cli := tidb.NewClient(matcher.Try(tc.GetDB(dbName)).(*sql.DB))
			// replica 2
			setAndWaitTiFlashReplicaAvailable(tiflashReplicas)
			c, _ := context.WithTimeout(context.Background(), 5*time.Minute)
			selectCountFromRandomTable(c, cli)
			// do some select
			// replica 1
			setAndWaitTiFlashReplicaAvailable(tiflashReplicas - 1)
			c, _ = context.WithTimeout(context.Background(), 5*time.Minute)
			selectCountFromRandomTable(c, cli)
			// replica 3
			setAndWaitTiFlashReplicaAvailable(tiflashReplicas + 1)
			c, _ = context.WithTimeout(context.Background(), 5*time.Minute)
			selectCountFromRandomTable(c, cli)
			// 	replica 0
			log.Info("set tiflash replica 0")
			matcher.Try(tidbCli.AlterDatabaseSetTiFlashReplica(dbName, 0, ""))
			gomega.Expect(matcher.Try(tidbCli.GetDBAvailableTiFlashReplicasNumber(dbName))).Should(gomega.BeZero())
		})
	})
})