package main

import (
	"context"
	"time"

	"github.com/pingcap/test-infra/caselib/pkg/steps"
	"github.com/pingcap/test-infra/caselib/pkg/utils"
	"github.com/pingcap/test-infra/common/model/resource"
	sdkResource "github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	"go.uber.org/zap"
)

func ImportPlacementMode() steps.Step {
	var (
		tidbRes                     = &resource.ResourceHandle{}
		lightningRes                = &resource.ResourceHandle{}
		opts                        = &steps.LightningImportOpts{}
		dropDatabases               string
		adminCheck                  string
		placementModeIgnore         = []steps.TiDBVariable{steps.TiDBVariable{Name: "TIDB_PLACEMENT_MODE", Value: "\"IGNORE\""}} // supported since 6.0
		placementModeStrict         = []steps.TiDBVariable{steps.TiDBVariable{Name: "TIDB_PLACEMENT_MODE", Value: "\"STRICT\""}} // supported since 6.0
		expectErrorTrue             = true
		expectErrorFalse            = false
		checkPlacementOptionsTables = &steps.Tables{}
		placementOptionsTrue        = true
		placementOptionsFalse       = false
	)

	return steps.
		NewBuilder().
		LoadEnv().
		LoadParamFromEnv(utils.ParamPrefix+"drop_databases", &dropDatabases).
		LoadResourceFromEnv(utils.ParamPrefix+"import_basic_opts", opts).
		LoadResourceFromEnv(utils.ResourcePrefix+"target", tidbRes).
		LoadResourceFromEnv(utils.ResourcePrefix+"lightning", lightningRes).
		LoadParamFromEnv(utils.ParamPrefix+"admin_check", &adminCheck).
		LoadResourceFromEnv(utils.ParamPrefix+"check_placement_options_tables", &checkPlacementOptionsTables).
		DropDatabasesIfExists(tidbRes, &dropDatabases).

		// lighting import data with placement options while placement policy not exist and placement_mode is strict => expect fail
		ConfigGlobalVariable(&placementModeStrict, tidbRes).
		ExecuteSQL(tidbRes, "", "drop PLACEMENT POLICY if exists placement_region;").
		ExpectError(&expectErrorTrue, steps.NewLightningImport(opts, lightningRes, tidbRes)).

		// lighting import data with placement options while placement policy not exist and placement_mode is ignore => expect pass
		ConfigGlobalVariable(&placementModeIgnore, tidbRes).
		ExpectError(&expectErrorFalse, steps.NewLightningImport(opts, lightningRes, tidbRes)).
		CheckTablesPlacementOptions(checkPlacementOptionsTables, &placementOptionsFalse, tidbRes).
		CheckDatabasesCrc32(&opts.Crc32Checksum, &dropDatabases, tidbRes).
		AdminCheck(tidbRes, &adminCheck).
		DropDatabasesIfExists(tidbRes, &dropDatabases).

		// lighting import data with placement options while placement policy exist and placement_mode is ignore => expect Pass and placement options not created.
		ConfigGlobalVariable(&placementModeIgnore, tidbRes).
		ExpectError(&expectErrorFalse, steps.NewLightningImport(opts, lightningRes, tidbRes)).
		CheckTablesPlacementOptions(checkPlacementOptionsTables, &placementOptionsFalse, tidbRes).
		CheckDatabasesCrc32(&opts.Crc32Checksum, &dropDatabases, tidbRes).
		AdminCheck(tidbRes, &adminCheck).
		DropDatabasesIfExists(tidbRes, &dropDatabases).

		// lighting import data with placement options while placement policy exist and placement_mode is strict => expect Pass and placement options created
		// This case (the dumpling output) requires target tikv with label: [server.labels]\nregion = "chengdu"
		// change tikv config to include specified region label
		UpdateTiDBClusterConfig(tidbRes, sdkResource.TiKV, "[server.labels]\nregion = \"chengdu\"", 10*time.Minute).
		ConfigGlobalVariable(&placementModeStrict, tidbRes).
		ExecuteSQL(tidbRes, "", "CREATE PLACEMENT POLICY IF NOT EXISTS `placement_region` PRIMARY_REGION=\"chengdu\" REGIONS=\"chengdu\";").
		ExpectError(&expectErrorFalse, steps.NewLightningImport(opts, lightningRes, tidbRes)).
		CheckTablesPlacementOptions(checkPlacementOptionsTables, &placementOptionsTrue, tidbRes).
		CheckDatabasesCrc32(&opts.Crc32Checksum, &dropDatabases, tidbRes).
		AdminCheck(tidbRes, &adminCheck).
		Build()

}

func main() {
	_, logger := utils.GetCtxAndLogger()
	ctx := steps.New(context.Background())
	if err := ImportPlacementMode().Execute(ctx); err != nil {
		logger.Fatal("ImportPlacementMode failed", zap.Error(err))
	}
	logger.Info("ImportPlacementMode finished")
}