package index_data_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/ch-1k-v5", "")
	flagSet.StringVar(&dbname, "db", "tpcc", "")
	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("clustered index", func() {
			TestCase("prepare", "prepareDataBeforeUpgrade", func() {
				ctx2, cancel := context.WithTimeout(ctx, duration)
				defer cancel()

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

				By("set clustered index")
				Try(db.ExecContext(ctx2, "set global tidb_enable_clustered_index =on "))

				By("Preparing TPC-C data")
				Try(db.ExecContext(ctx2, fmt.Sprintf("drop database if exists %s", dbname)))
				Try(db.ExecContext(ctx2, fmt.Sprintf("create database %s", dbname)))
				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("create expression index")
				db = Try(tc.GetDB(dbname)).(*sql.DB)
				Try(db.Exec("create index expression_idx1 on customer ((lower(c_city)))"))
				Try(db.Exec("create index expression_idx2 on customer ((upper(c_city)))"))
				Try(db.Exec("create index expression_idx3 on customer ((reverse(c_city)))"))
				Try(db.Exec("create index expression_idx4 on customer ((md5(c_city)))"))
				Try(db.Exec("create index expression_idx5 on customer ((vitess_hash(c_w_id)))"))

				By("create secondary index")
				Try(db.Exec("create index secondary_idx on customer (c_phone)"))

				By("create prefix index")
				Try(db.Exec("create index prefix_idx on customer (c_city(3))"))

				By("create invisible index")
				Try(db.Exec("create index invisible_idx on customer (c_street_1)"))
				Try(db.Exec("alter table customer alter index invisible_idx invisible"))

			})

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

				By("check index data")
				db.Exec("admin check customer")
				rowNum := 0
				expectNum := 30000000
				db.QueryRow("select count(*) from customer").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(idx_customer)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(expression_idx1)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(expression_idx2)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(expression_idx3)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(expression_idx4)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(expression_idx5)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(secondary_idx)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(prefix_idx)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
				db.QueryRow("select count(*) from customer use index(invisible_idx)").Scan(&rowNum)
				Expect(rowNum).To(Equal(expectNum))
			})
		})
	})
})

var _ = Describe("upgradeCluster", 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())
	})
})