// Copyright 2022 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package compatibility

import (
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/workload"

	. "github.com/pingcap/endless/testcase/geopartition/utils"
)

func TiflashTest(ctx *SuiteContext) {
	It("placement should not add tiflash replica", func() {
		defer func() {
			MustExec(ctx.DB, "DROP TABLE IF EXISTS t1, tp")
			MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")
		}()
		MustExec(ctx.DB, "DROP TABLE IF EXISTS t1, tp")
		MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")

		MustExec(ctx.DB, "CREATE PLACEMENT POLICY p1 PRIMARY_REGION='us-east-1' REGIONS='us-east-1,us-east-2'")
		MustExec(ctx.DB, "CREATE TABLE t1(a int) PLACEMENT POLICY p1")
		Try(WaitScheduled(ctx.DB, "t1", time.Minute))
		MustExec(ctx.DB, `CREATE TABLE tp (id INT) PARTITION BY RANGE (id) ( 
			PARTITION p0 VALUES LESS THAN (100) PLACEMENT POLICY p1,
			PARTITION p1 VALUES LESS THAN (1000)
		)`)
		Try(WaitScheduled(ctx.DB, "tp", time.Minute))

		MustExecError(ctx.DB, 8200, "ALTER TABLE t1 SET TIFLASH REPLICA 1")
		MustExecError(ctx.DB, 8200, "ALTER TABLE tp SET TIFLASH REPLICA 1")

		Expect(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE t1"))[0][1]).Should(Equal("CREATE TABLE `t1` (\n  " +
			"`a` int(11) DEFAULT NULL\n" +
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin /*T![placement] PLACEMENT POLICY=`p1` */"))
		Expect(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE tp"))[0][1]).Should(Equal("CREATE TABLE `tp` (\n" +
			"  `id` int(11) DEFAULT NULL\n" +
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin\n" +
			"PARTITION BY RANGE (`id`)\n" +
			"(PARTITION `p0` VALUES LESS THAN (100) /*T![placement] PLACEMENT POLICY=`p1` */,\n" +
			" PARTITION `p1` VALUES LESS THAN (1000))"))

		CheckRegions(ctx, "t1", "", func(r Regions) {
			r.CheckLeaderPlacement("+region=us-east-1,-engine=tiflash").
				CheckFollowerPlacement("+region=us-east-1,-engine=tiflash", 1).
				CheckFollowerPlacement("+region=us-east-2,-engine=tiflash", 1).
				CheckNoPeersLeft()
		})

		CheckRegions(ctx, "tp", "p0", func(r Regions) {
			r.CheckLeaderPlacement("+region=us-east-1,-engine=tiflash").
				CheckFollowerPlacement("+region=us-east-1,-engine=tiflash", 1).
				CheckFollowerPlacement("+region=us-east-2,-engine=tiflash", 1).
				CheckNoPeersLeft()
		})

		CheckRegions(ctx, "tp", "p1", func(r Regions) {
			r.CheckLeaderPlacement("-engine=tiflash").
				CheckFollowerPlacement("-engine=tiflash", 2).
				CheckNoPeersLeft()
		})
	})

	It("tiflash table should not add placement", func() {
		defer func() {
			MustExec(ctx.DB, "DROP TABLE IF EXISTS t1, tp")
			MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")
		}()
		MustExec(ctx.DB, "DROP TABLE IF EXISTS t1, tp")
		MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")

		MustExec(ctx.DB, "CREATE PLACEMENT POLICY p1 PRIMARY_REGION='us-east-1' REGIONS='us-east-1,us-east-2'")
		MustExec(ctx.DB, "CREATE TABLE t1(a int)")
		Try(WaitScheduled(ctx.DB, "t1", time.Minute))
		MustExec(ctx.DB, `CREATE TABLE tp (id INT) PARTITION BY RANGE (id) ( 
			PARTITION p0 VALUES LESS THAN (100),
			PARTITION p1 VALUES LESS THAN (1000)
		)`)
		MustExec(ctx.DB, "ALTER TABLE t1 SET TIFLASH REPLICA 1")
		MustExec(ctx.DB, "ALTER TABLE tp SET TIFLASH REPLICA 1")

		MustExecError(ctx.DB, 8200, "ALTER TABLE t1 PLACEMENT POLICY p1")
		MustExecError(ctx.DB, 8200, "ALTER TABLE tp PLACEMENT POLICY p1")
		MustExecError(ctx.DB, 8200, "ALTER TABLE tp ADD PARTITION (PARTITION p2 VALUES LESS THAN (10000) PLACEMENT POLICY p1)")

		Expect(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE t1"))[0][1]).Should(Equal("CREATE TABLE `t1` (\n  " +
			"`a` int(11) DEFAULT NULL\n" +
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin"))
		Expect(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE tp"))[0][1]).Should(Equal("CREATE TABLE `tp` (\n" +
			"  `id` int(11) DEFAULT NULL\n" +
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin\n" +
			"PARTITION BY RANGE (`id`)\n" +
			"(PARTITION `p0` VALUES LESS THAN (100),\n" +
			" PARTITION `p1` VALUES LESS THAN (1000))"))

	})

	It("should not schedule placement peer to tiflash store even not enough tikv nodes", func() {
		defer func() {
			MustExec(ctx.DB, "DROP TABLE IF EXISTS tp, t1")
			MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")
			MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p2")
		}()
		MustExec(ctx.DB, "DROP TABLE IF EXISTS tp, t1")
		MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")
		MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")

		MustExec(ctx.DB, "CREATE PLACEMENT POLICY p1 FOLLOWERS=7")
		MustExec(ctx.DB, "CREATE PLACEMENT POLICY p2 FOLLOWERS=8")
		MustExec(ctx.DB, `CREATE TABLE t1 (id INT) PLACEMENT POLICY p1`)
		MustExec(ctx.DB, `CREATE TABLE tp (id INT) PLACEMENT POLICY p1 PARTITION BY RANGE (id) ( 
			PARTITION p0 VALUES LESS THAN (100) PLACEMENT POLICY p2,
			PARTITION p1 VALUES LESS THAN (1000)
		)`)
		Try(WaitScheduledOrPending(ctx.DB, "t1", time.Minute))
		Try(WaitScheduledOrPending(ctx.DB, "tp", time.Minute))
		CheckRegions(ctx, "t1", "", func(r Regions) {
			r.CheckLeaderPlacement("-engine=tiflash").
				CheckFollowerPlacement("-engine=tiflash", 6).
				CheckNoPeersLeft()
		})

		CheckRegions(ctx, "tp", "p0", func(r Regions) {
			r.CheckLeaderPlacement("-engine=tiflash").
				CheckFollowerPlacement("-engine=tiflash", 6).
				CheckNoPeersLeft()
		})

		CheckRegions(ctx, "tp", "p1", func(r Regions) {
			r.CheckLeaderPlacement("-engine=tiflash").
				CheckFollowerPlacement("-engine=tiflash", 6).
				CheckNoPeersLeft()
		})
	})
}