// 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 (
	"database/sql"
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	pdclient "github.com/pingcap/endless/pkg/client/pd"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/workload"
	. "github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/test-infra/sdk/resource"

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

var policyInfos = map[string]struct {
	placement string
}{
	"p1": {placement: "FOLLOWERS=1"},
	"p2": {placement: `PRIMARY_REGION="us-east-1" REGIONS="us-east-1,us-east-2"`},
	"p3": {placement: `LEADER_CONSTRAINTS="[+region=us-east-2]" FOLLOWER_CONSTRAINTS="[+region=us-west-1]"`},
	"p4": {placement: "FOLLOWERS=4"},
}

func buildBRCmd(options []string) string {
	return "export AWS_ACCESS_KEY_ID=minioadmin;export AWS_SECRET_ACCESS_KEY=minioadmin;/br " + strings.Join(options, " ")
}

func clearTiDB(db *sql.DB) {
	MustExec(db, "DROP DATABASE IF EXISTS db2")
	MustExec(db, "DROP DATABASE IF EXISTS test")
	MustExec(db, "DROP PLACEMENT POLICY IF EXISTS p1")
	MustExec(db, "DROP PLACEMENT POLICY IF EXISTS p2")
	MustExec(db, "DROP PLACEMENT POLICY IF EXISTS p3")
	MustExec(db, "DROP PLACEMENT POLICY IF EXISTS p4")
	MustExec(db, "CREATE DATABASE test")
}

func prepareBackupSource(db *sql.DB) {
	clearTiDB(db)
	for name, info := range policyInfos {
		MustExec(db, fmt.Sprintf("CREATE PLACEMENT POLICY %s %s", name, info.placement))
	}
	MustExec(db, "CREATE DATABASE db2")
	MustExec(db, "CREATE TABLE test.t1(a int) PLACEMENT POLICY p1")
	MustExec(db, `CREATE TABLE db2.t1 (a int) PLACEMENT policy p2 PARTITION BY RANGE (a) ( 
			PARTITION p0 VALUES LESS THAN (100) PLACEMENT POLICY p3,
			PARTITION p1 VALUES LESS THAN (1000)
		)`)
	MustExec(db, "CREATE TABLE db2.t2(a int)")
	MustExec(db, "ALTER DATABASE db2 PLACEMENT POLICY p1")
}

func checkPolicies(ctx *SuiteContext, existPoliciesWhenHasPlacement []string, hasPlacement bool) {
	existPolicies := existPoliciesWhenHasPlacement
	if !hasPlacement {
		existPolicies = []string{}
	}

	existPoliciesMap := make(map[string]interface{})
	for _, name := range existPolicies {
		existPoliciesMap[name] = struct{}{}
	}

	for name, info := range policyInfos {
		if _, ok := existPoliciesMap[name]; ok {
			expected := []string{name, fmt.Sprintf("CREATE PLACEMENT POLICY `%s` %s", name, info.placement)}
			Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE PLACEMENT POLICY "+name)))).Should(Equal([][]string{expected}))
		} else {
			Expect(Try(ReadAll(MustQuery(ctx.DB, "SELECT * FROM INFORMATION_SCHEMA.PLACEMENT_POLICIES WHERE POLICY_NAME = '"+name+"'")))).Should(Equal([][]string{}))
		}
	}
}

func checkDatabaseDB2(ctx *SuiteContext, exist bool, hasPlacement bool) {
	if !exist {
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW DATABASES LIKE 'db2'")))).Should(Equal([][]string{}))
		return
	}

	createSQL := "CREATE DATABASE `db2` /*!40100 DEFAULT CHARACTER SET utf8mb4 */"
	if hasPlacement {
		createSQL += " /*T![placement] PLACEMENT POLICY=`p1` */"
	}

	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE DATABASE db2")))).
		Should(Equal([][]string{{"db2", createSQL}}))
}

func checkTableT1InTest(ctx *SuiteContext, exist bool, hasPlacement bool) {
	if !exist {
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA='test' AND TABLE_NAME='t1'")))).Should(Equal([][]string{}))
		return
	}

	getCreateSQL := func(placement string) string {
		return fmt.Sprintf("CREATE TABLE `t1` (\n  "+
			"`a` int(11) DEFAULT NULL\n"+
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin%s", placement)
	}

	if !hasPlacement {
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE test.t1")))).
			Should(Equal([][]string{{"t1", getCreateSQL("")}}))
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW PLACEMENT LIKE 'TABLE test.t1'")))).
			Should(Equal([][]string{}))
		return
	}

	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE test.t1")))).
		Should(Equal([][]string{{"t1", getCreateSQL(" /*T![placement] PLACEMENT POLICY=`p1` */")}}))
	Try(WaitScheduled(ctx.DB, "test.t1", time.Minute))
	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW PLACEMENT LIKE 'TABLE test.t1'")))).Should(Equal([][]string{
		{"TABLE test.t1", policyInfos["p1"].placement, "SCHEDULED"},
	}))

	followerCnt := 1
	if policyInfos["p1"].placement == "FOLLOWERS=3" {
		followerCnt = 3
	}

	CheckRegions(ctx, "test.t1", "", func(r Regions) {
		r.CheckLeaderPlacement("").
			CheckFollowerPlacement("", followerCnt).
			CheckNoPeersLeft()
	})
}

func checkTableT1InDB2(ctx *SuiteContext, exist bool, hasPlacement bool) {
	if !exist {
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA='db2' AND TABLE_NAME='t1'")))).Should(Equal([][]string{}))
		return
	}

	getCreateSQL := func(placement1 string, placement2 string) string {
		return fmt.Sprintf("CREATE TABLE `t1` (\n  "+
			"`a` int(11) DEFAULT NULL\n"+
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin%s\n"+
			"PARTITION BY RANGE (`a`)\n"+
			"(PARTITION `p0` VALUES LESS THAN (100)%s,\n "+
			"PARTITION `p1` VALUES LESS THAN (1000))", placement1, placement2)
	}

	if !hasPlacement {
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE db2.t1")))).
			Should(Equal([][]string{{"t1", getCreateSQL("", "")}}))
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW PLACEMENT LIKE 'TABLE db2.t1'")))).
			Should(Equal([][]string{}))
		return
	}

	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE db2.t1")))).
		Should(Equal([][]string{{"t1", getCreateSQL(" /*T![placement] PLACEMENT POLICY=`p2` */", " /*T![placement] PLACEMENT POLICY=`p3` */")}}))
	Try(WaitScheduled(ctx.DB, "db2.t1", time.Minute))
	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW PLACEMENT LIKE 'TABLE db2.t1%'")))).Should(Equal([][]string{
		{"TABLE db2.t1", policyInfos["p2"].placement, "SCHEDULED"},
		{"TABLE db2.t1 PARTITION p0", policyInfos["p3"].placement, "SCHEDULED"},
		{"TABLE db2.t1 PARTITION p1", policyInfos["p2"].placement, "SCHEDULED"},
	}))

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

	CheckRegions(ctx, "db2.t1", "p1", func(r Regions) {
		r.CheckLeaderPlacement("+region=us-east-1").
			CheckFollowerRegions(2, "us-east-1", "us-east-2").
			CheckNoPeersLeft()
	})
}

func checkTableT2InDB2(ctx *SuiteContext, exist bool) {
	if !exist {
		Expect(Try(ReadAll(MustQuery(ctx.DB, "SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA='db2' AND TABLE_NAME='t2'")))).Should(Equal([][]string{}))
		return
	}

	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW CREATE TABLE db2.t2")))).
		Should(Equal([][]string{{"t2", "CREATE TABLE `t2` (\n  " +
			"`a` int(11) DEFAULT NULL\n" +
			") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin"}}))
	Expect(Try(ReadAll(MustQuery(ctx.DB, "SHOW PLACEMENT LIKE 'TABLE db2.t2'")))).
		Should(Equal([][]string{}))
}

func checkFullRestore(ctx *SuiteContext, hasPlacement bool) {
	checkRestore(ctx, "", "", hasPlacement)
}

func checkRestoreDB(ctx *SuiteContext, db string, hasPlacement bool) {
	checkRestore(ctx, db, "", hasPlacement)
}

func checkRestoreTable(ctx *SuiteContext, tb string, hasPlacement bool) {
	checkRestore(ctx, "", tb, hasPlacement)
}

func checkRestore(ctx *SuiteContext, restoreDB string, restoreTable string, hasPlacement bool) {
	fullRestore := restoreDB == "" && restoreTable == ""
	if restoreDB != "" && restoreTable != "" {
		Fail("only one of restoreDB and restoreTable can be present")
	}

	switch {
	case fullRestore:
		checkPolicies(ctx, []string{"p1", "p2", "p3", "p4"}, hasPlacement)
		checkTableT1InTest(ctx, true, hasPlacement)
		checkDatabaseDB2(ctx, true, hasPlacement)
		checkTableT1InDB2(ctx, true, hasPlacement)
		checkTableT2InDB2(ctx, true)
	case restoreDB == "test":
		checkPolicies(ctx, []string{"p1"}, hasPlacement)
		checkTableT1InTest(ctx, true, hasPlacement)
		checkDatabaseDB2(ctx, false, hasPlacement)
		checkTableT1InDB2(ctx, false, hasPlacement)
		checkTableT2InDB2(ctx, false)
	case restoreDB == "db2":
		checkPolicies(ctx, []string{"p1", "p2", "p3"}, hasPlacement)
		checkTableT1InTest(ctx, false, hasPlacement)
		checkDatabaseDB2(ctx, true, hasPlacement)
		checkTableT1InDB2(ctx, true, hasPlacement)
		checkTableT2InDB2(ctx, true)
	case restoreTable == "test.t1":
		checkPolicies(ctx, []string{"p1"}, hasPlacement)
		checkTableT1InTest(ctx, true, hasPlacement)
		checkDatabaseDB2(ctx, false, hasPlacement)
		checkTableT1InDB2(ctx, false, hasPlacement)
		checkTableT2InDB2(ctx, false)
	case restoreTable == "db2.t1":
		checkPolicies(ctx, []string{"p1", "p2", "p3"}, hasPlacement)
		checkTableT1InTest(ctx, false, hasPlacement)
		checkDatabaseDB2(ctx, true, hasPlacement)
		checkTableT1InDB2(ctx, true, hasPlacement)
		checkTableT2InDB2(ctx, false)
	case restoreTable == "db2.t2":
		checkPolicies(ctx, []string{"p1"}, hasPlacement)
		checkTableT1InTest(ctx, false, hasPlacement)
		checkDatabaseDB2(ctx, true, hasPlacement)
		checkTableT1InDB2(ctx, false, hasPlacement)
		checkTableT2InDB2(ctx, true)
	default:
		Fail(fmt.Sprintf("Unsupported restoreDB: %s, restoreTable: %s", restoreDB, restoreTable))
	}
}

func brBackUp(ctx *SuiteContext, mode string, args ...string) {
	Try(CommandContext(ctx, "rm -rf /data/br-backup&&mkdir -p /data/br-backup").Run(ctx.Resource("minio").(resource.WorkloadNode)))
	Try(CommandContext(ctx, buildBRCmd([]string{
		"backup",
		mode,
		strings.Join(args, " "),
		"--pd", ctx.PD.GetPDAddr(),
		"--s3.endpoint http://minio-peer:9000",
		"-s s3://br-backup/backup",
	})).Run(ctx.Resource("br").(resource.BR)))
}

func brRestore(ctx *SuiteContext, mode string, args ...string) {
	pdClient := Try(pdclient.NewPDClient(ctx.TC)).(*pdclient.Client)
	db := Try(ctx.TC.GetDB("test")).(*sql.DB)
	defer func() {
		Try(db.Close())
	}()

	cmd := []string{
		"restore",
		mode,
		"--pd", pdClient.GetPDAddr(),
		"--s3.endpoint http://minio-peer:9000",
		"-s s3://br-backup/backup",
	}

	if len(args) > 0 {
		cmd = append(cmd, args...)
	}

	Try(CommandContext(ctx, buildBRCmd(cmd)).Run(ctx.Resource("br").(resource.BR)))
}

func BRTest(ctx *SuiteContext) {
	It("br test1, br backup full and restore full restore", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do full backup
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "full")
		checkFullRestore(ctx, true)
	})

	It("br test2, br backup full and restore full but policy p1 exists", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do full backup
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do full restore
		clearTiDB(ctx.DB)
		MustExec(ctx.DB, "create placement policy p1 followers=3")
		brRestore(ctx, "full")
		oldP1 := policyInfos["p1"]
		policyInfos["p1"] = struct{ placement string }{placement: "FOLLOWERS=3"}
		defer func() {
			policyInfos["p1"] = oldP1
		}()
		checkFullRestore(ctx, true)
	})

	It("br test3, br backup database and restore full", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do backup db2
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "db", "--db", "db2")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "full")
		checkRestoreDB(ctx, "db2", false)
	})

	It("br test4, br backup table and restore full", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do backup db2.t1
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "table", "--db", "db2", "-t", "t1")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "full")
		checkRestoreTable(ctx, "db2.t1", false)
	})

	It("br test5, br backup full and restore database db2", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do backup full
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do restore db2
		clearTiDB(ctx.DB)
		brRestore(ctx, "db", "--db", "db2")
		checkRestoreDB(ctx, "db2", true)
	})

	It("br test6, br backup full and restore table test.t1", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do full backup
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "table", "--db", "test", "-t", "t1")
		checkRestoreTable(ctx, "test.t1", true)
	})

	It("br test7, br backup full and restore table db2.t1", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do full backup
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "table", "--db", "db2", "-t", "t1")
		checkRestoreTable(ctx, "db2.t1", true)
	})

	It("br test8, br backup full and restore table db2.t2", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do full backup
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "table", "--db", "db2", "-t", "t2")
		checkRestoreTable(ctx, "db2.t2", true)
	})

	It("br test9, br backup full and restore full with tidb-placement-mode ignore", func() {
		defer func() {
			clearTiDB(ctx.DB)
		}()

		// do full backup
		prepareBackupSource(ctx.DB)
		brBackUp(ctx, "full")

		// do full restore
		clearTiDB(ctx.DB)
		brRestore(ctx, "full", "--with-tidb-placement-mode", "ignore")
		checkFullRestore(ctx, false)
	})
}