package functionality

import (
	"strconv"
	"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"
	"github.com/pingcap/test-infra/sdk/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func constraintsSugarSyntaxTest(ctx *utils.SuiteContext) {
	ginkgo.It("Constraints Sugar Syntax #constraints_sugar_syntax#", func() {
		defer func() {
			workload.MustExec(ctx.DB, "DROP TABLE IF EXISTS t1,t2,t3,tx")
			workload.MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")
		}()
		workload.MustExec(ctx.DB, "DROP TABLE IF EXISTS t1,t2,t3,tx")
		workload.MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")

		workload.MustExec(ctx.DB, "CREATE PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2\" FOLLOWERS=2; ")
		workload.MustExec(ctx.DB, "CREATE TABLE t1 (id INT NOT NULL PRIMARY KEY auto_increment, b VARBINARY(65535)) PLACEMENT POLICY=p1;")
		for n := 1; n < 100; n++ {
			workload.MustExec(ctx.DB, "insert into t1 values (?, ?)", n, strconv.Itoa(n*10))
		}

		// SCHEDULE="EVEN", PRIMARY_REGION should have an even count peers.
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2,us-west-1\" SCHEDULE=\"EVEN\" FOLLOWERS=4;")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerPlacement("+region=us-east-1", 1).
				CheckFollowerRegions(3, "us-east-2", "us-west-1").
				CheckNoPeersLeft()
		})

		// default schedule should be "EVEN"
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2,us-west-1\" FOLLOWERS=4;")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerPlacement("+region=us-east-1", 1).
				CheckFollowerRegions(3, "us-east-2", "us-west-1").
				CheckNoPeersLeft()
		})

		// SCHEDULE="MAJORITY_IN_PRIMARY", PRIMARY_REGION should have a majority count of peers
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2,us-west-1\" FOLLOWERS=4 SCHEDULE=\"MAJORITY_IN_PRIMARY\";")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerPlacement("+region=us-east-1", 2).
				CheckFollowerRegions(2, "us-east-2", "us-west-1").
				CheckNoPeersLeft()
		})

		// SCHEDULE="EVEN" and if peers are less than regions, PRIMARY_REGION should have at least one peer
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2,us-west-1\" FOLLOWERS=1 SCHEDULE=\"MAJORITY_IN_PRIMARY\";")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerRegions(1, "us-east-2", "us-west-1").
				CheckNoPeersLeft()
		})

		// SCHEDULE="MAJORITY_IN_PRIMARY" and if there are only 2 peers (1 follower), only PRIMARY_REGION can be scheduled
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2,us-west-1\" FOLLOWERS=1 SCHEDULE=\"MAJORITY_IN_PRIMARY\";")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerPlacement("+region=us-east-1", 1).
				CheckNoPeersLeft()
		})

		// SCHEDULE="EVEN" and there are only one region
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1\" SCHEDULE=\"EVEN\";")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerPlacement("+region=us-east-1", 2).
				CheckNoPeersLeft()
		})

		// SCHEDULE="MAJORITY_IN_PRIMARY" and there are only one region
		workload.MustExec(ctx.DB, "ALTER PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1\" SCHEDULE=\"MAJORITY_IN_PRIMARY\";")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		utils.CheckRegions(ctx, "t1", "", func(r utils.Regions) {
			r.CheckLeaderPlacement("+region=us-east-1").
				CheckFollowerPlacement("+region=us-east-1", 2).
				CheckNoPeersLeft()
		})
	})

	ginkgo.It("shutdown one region should work", func() {
		defer func() {
			workload.MustExec(ctx.DB, "DROP TABLE IF EXISTS t1")
			workload.MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")
		}()
		workload.MustExec(ctx.DB, "DROP TABLE IF EXISTS t1")
		workload.MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS p1")

		workload.MustExec(ctx.DB, "CREATE PLACEMENT POLICY p1 PRIMARY_REGION=\"us-east-1\" REGIONS=\"us-east-1,us-east-2\" FOLLOWERS=2; ")
		workload.MustExec(ctx.DB, "CREATE TABLE t1 (id INT NOT NULL PRIMARY KEY auto_increment, b VARBINARY(65535)) PLACEMENT POLICY=p1;")
		matcher.Try(utils.WaitScheduled(ctx.DB, "t1", time.Minute))
		for n := 1; n < 100; n++ {
			workload.MustExec(ctx.DB, "insert into t1 values (?, ?)", n, strconv.Itoa(n*10))
		}

		// enable chaos
		tc := ctx.TC
		chaosClient := tc.ChaosClient()
		selector := matcher.Try(tc.ChaosSelect(resource.TiDBClusterSelectorOpts{
			TiDBClusterTiKVSelectorSpec: &resource.TiDBClusterTiKVSelectorSpec{
				StoreLabelSelector: &metav1.LabelSelector{
					MatchLabels: map[string]string{"region": "us-east-2"},
				},
			},
		})).([]resource.ChaosSelector)

		chaosId := matcher.Try(chaosClient.Create(resource.InjectChaosOptions{
			Selector: selector,
			ExperimentSpec: &resource.NetworkLossSpec{
				Loss:        "100",
				Correlation: "100",
			},
		})).(resource.ChaosIdentifier)
		defer func() {
			matcher.Try(tc.ChaosClient().Delete(chaosId))
		}()

		for n := 1; n < 100; n++ {
			rows := utils.ReadAll(utils.MustQuery(ctx.DB, "select id from t1 where id = ?", n))
			Expect(rows).Should(Equal([][]string{{strconv.Itoa(n)}}))
		}
	})
}
