package examples_test

import (
	"context"
	"database/sql"
	"flag"
	"fmt"
	"net/url"
	"time"

	_ "github.com/go-sql-driver/mysql"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/pkg/chaos"
	"github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/matcher/metrics"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/pkg/workload"
	"github.com/pingcap/endless/pkg/workloadnode"
)

var _ = ParameterizedDescribe("failover", func(flagSet *flag.FlagSet) {
	var (
		s3Endpoint        string
		s3AccessKeyID     string
		s3SecretAccessKey string

		storage   string
		dbname    string
		warehouse int
		threads   int
		duration  time.Duration
	)

	flagSet.StringVar(&s3Endpoint, "s3-endpoing", "http://minio.pingcap.net:9000", "")
	flagSet.StringVar(&s3AccessKeyID, "s3-key", "minioadmin", "")
	flagSet.StringVar(&s3SecretAccessKey, "s3-secret", "minioadmin", "")

	flagSet.StringVar(&storage, "storage", "s3://tpcc/br-500m-v4.0.13", "")
	flagSet.StringVar(&dbname, "db", "test", "")
	flagSet.IntVar(&warehouse, "warehouse", 1000, "")
	flagSet.IntVar(&threads, "threads", 1024, "Threads number.")
	flagSet.DurationVar(&duration, "duration", 3*time.Hour, "Run duration.")

	ctx := suiteTestCtx

	// Testbed is testbed/testbed_multizone.yaml.
	Context("with one AZ Down", func() {

		brCmdTpl := `
export AWS_ACCESS_KEY_ID=%s AWS_SECRET_ACCESS_KEY=%s
/br restore full --pd "%s" \
	--storage "%s" \
	--s3.endpoint "%s" \
	--send-credentials-to-tikv=true
`
		goTpcCmdTpl := `/go-tpc tpcc --host %s -P4000 --warehouses %d run -D %s -T %d`

		var (
			tc    resource.TiDBCluster
			br    resource.BR
			goTpc resource.WorkloadNode
			dbURL *url.URL
		)

		BeforeEach(func() {
			tc = ctx.Resource("tc").(resource.TiDBCluster)
			br = ctx.Resource("br").(resource.BR)
			goTpc = ctx.Resource("gotpc").(resource.WorkloadNode)
			dbURL = Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
		})

		// You can run this testcase with a command like the following:
		// ginkgo -focus "#Endless_Sample_Failover#" -- failover -duration=10s
		TestCase("should work (recover the QPS in 1 minutes)", "Endless_Sample_Failover", func() {
			ctx2, cancel := context.WithTimeout(ctx, duration)
			defer cancel()

			By("Preparing TPC-C data")
			// Check if data is already prepared.
			db := Try(tc.GetDB("")).(*sql.DB)
			tidbClient := tidb.NewClient(db)
			exist := Try(tidbClient.IsTableExist(ctx2, dbname, "warehouse")).(bool)
			if !exist {
				pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
				s := fmt.Sprintf(brCmdTpl, s3AccessKeyID, s3SecretAccessKey, pdURL.Host, storage, s3Endpoint)
				cmd := workloadnode.CommandContext(ctx2, s)
				Try(cmd.Run(br))
			}
			log.Info("TPC-C data is ready")

			g, ctx3 := errgroup.WithContext(ctx2)

			By("Running tpcc workload asynchronously")
			g.Go(WithGinkgoRecover(func() {
				s := fmt.Sprintf(goTpcCmdTpl, dbURL.Hostname(), warehouse, dbname, threads)
				cmd := workloadnode.CommandContext(ctx3, s)
				Try(cmd.Start(goTpc))
				defer cmd.Stop()
				if err := cmd.Wait(); err != nil {
					Expect(cmd.ExitDueToCtxDone).To(BeTrue())
				}
			}))

			By("Running metrics checker")
			g.Go(WithGinkgoRecover(func() {
				select {
				case <-time.After(5 * time.Minute): // Time for workload to warm up.
					ConsistentlyWithMetrics(tc).
						ShouldNot(QPSDropMoreThan(0.5).For(1 * time.Minute).WithCtx(ctx3))
				case <-ctx3.Done():
				}
			}))

			By("Running chaos periodically after workload starts")
			chaos := chaos.Chaos{
				Name:           "DownAllTiKVInAZ`zone1`",
				SelectorOpts:   chaos.SelectAllTiKVInOneAZ("zone1"),
				ExperimentSpec: chaos.FailureExperiment(10*time.Second, "@every 5m"),
				Tc:             tc,
			}
			defer chaos.Clean()
			Try(chaos.Run())

			g.Wait()
		})
	})
})

var _ = ParameterizedDescribe("tikvoom", func(flagSet *flag.FlagSet) {
	ctx := suiteTestCtx

	var (
		duration      time.Duration
		threadsNUmber int
		tablesNumber  int
		paddingLength int
	)

	flagSet.DurationVar(&duration, "duration", time.Second*10, "Run duration.")
	flagSet.IntVar(&threadsNUmber, "threads", 1024, "Threads number.")
	flagSet.IntVar(&tablesNumber, "tables", 128, "Tables number.")
	flagSet.IntVar(&paddingLength, "padding", 4000000, "Row padding length.")

	Context("under heavy append workload", func() {
		var (
			tc resource.TiDBCluster
		)

		BeforeEach(func() {
			tc = ctx.Resource("tc").(resource.TiDBCluster)
		})

		TestCase("should not occur", "Endless_Sample_Tikvoom", func() {
			db := Try(tc.GetDB("test")).(*sql.DB)
			ctx2, cancel := context.WithTimeout(ctx, duration)
			defer cancel()

			w := workload.AppendWorkload{
				DB:          db,
				Concurrency: threadsNUmber,
				Tables:      tablesNumber,
				PadLength:   paddingLength,
			}
			Try(w.Prepare())
			go w.Run(ctx2)

			// TODO: use tc.WatchOOM to check if there was an OOM event.
			ConsistentlyWithMetrics(tc).
				ShouldNot(TiKVRestart().WithCtx(ctx2))
		})
	})
})