package auto_increment_upgrade

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

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	_ "github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/test-infra/sdk/resource"
)

var _ = ParameterizedDescribe("Index", 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://benchmark/sysbench-32-1000w-auto", "")
	flagSet.StringVar(&dbname, "db", "sbtest", "")
	flagSet.IntVar(&warehouse, "warehouse", 1000, "")
	flagSet.IntVar(&threads, "threads", 1024, "Threads number.")
	flagSet.DurationVar(&duration, "duration", 3*time.Hour, "Run duration.")

	ctx := suiteTestCtx
	var (
		tc resource.TiDBCluster
		br resource.BR
	)
	Context("tidb upgrade", 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 --check-requirements=false
`
		BeforeEach(func() {
			tc = ctx.Resource("tc").(resource.TiDBCluster)
			br = ctx.Resource("br").(resource.BR)
		})

		Context("auto increment", func() {
			TestCase("prepare", "prepareAutoIncrementDataBeforeUpgrade", func() {
				ctx2, cancel := context.WithTimeout(ctx, duration)
				defer cancel()

				By("get db")
				db := Try(tc.GetDB("test")).(*sql.DB)

				By("Preparing sysbench data")
				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))

				By("auto_increment")
				Try(db.ExecContext(ctx2, fmt.Sprintf("drop table if exists test1")))
				Try(db.ExecContext(ctx2, fmt.Sprintf("drop table if exists test2")))
				Try(db.ExecContext(ctx2, fmt.Sprintf("drop table if exists test3")))
				Try(db.Exec("CREATE TABLE `test1` (`id` int(11) NOT NULL auto_increment,  `k` int(11) NOT NULL DEFAULT '0',  `c` char(120) NOT NULL DEFAULT '',  `pad` char(60) NOT NULL DEFAULT '',  PRIMARY KEY (`id`) NONCLUSTERED,  KEY `k_1` (`k`))"))
				Try(db.Exec("CREATE TABLE `test2` (`id` int(11) NOT NULL auto_increment,  `k` int(11) NOT NULL DEFAULT '0',  `c` char(120) NOT NULL DEFAULT '',  `pad` char(60) NOT NULL DEFAULT '',  PRIMARY KEY (`id`) NONCLUSTERED,  KEY `id_1` (`id`))"))
				Try(db.Exec("CREATE TABLE `test3` (`id` int(11) NOT NULL auto_increment,  `k` int(11) NOT NULL DEFAULT '0',  `c` char(120) NOT NULL DEFAULT '',  `pad` char(60) NOT NULL DEFAULT '',  KEY `id_1` (`id`))"))
				for i := 0; i < 10; i++ {
					Try(db.Exec("insert into test1(k,c,pad) select k,c,pad from sbtest1 limit 1000000"))
					Try(db.Exec("insert into test2(k,c,pad) select k,c,pad from sbtest1 limit 1000000"))
					Try(db.Exec("insert into test3(k,c,pad) select k,c,pad from sbtest1 limit 1000000"))
				}
				Try(db.Exec("alter table test1 SHARD_ROW_ID_BITS = 4"))
				Try(db.Exec("alter table test2 drop primary key"))
				Try(db.Exec("alter table test3 add primary key(id)"))

			})

			TestCase("check", "checkAutoIncrementDataAfterUpgrade", func() {
				By("get db")
				db := Try(tc.GetDB("test")).(*sql.DB)

				By("check data")
				Try(db.Exec("admin check table test1,test2,test3"))
				rowNum := 0
				expectNum := 10000000
				db.QueryRow("select count(*) from test1").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from test2").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from test3").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				//check auto_increment
				Try(db.Exec("insert into test1(k,c,pad) values(2147483647,'a','a')"))

			})
		})
	})
})

var _ = Describe("AutoIncrementUpgradgeTidbCluster", func() {
	It("operator upgrade", func() {
		tc := suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		targetVersion := suiteTestCtx.Param("target_version", "master")
		opt := resource.TiDBClusterUpgradeOptions{
			Version: targetVersion,
			Timeout: 30 * time.Minute,
		}
		err := tc.Upgrade(opt)
		Expect(err).Should(BeNil())
	})
})