package main

import (
	"context"
	"flag"
	"fmt"
	"github.com/pingcap/log"
	pd_component "github.com/pingcap/schrodinger-test/once/components/pd"
	"github.com/pingcap/schrodinger-test/once/drivers/pd/pd_ctl"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"go.uber.org/zap"
	"os"
	"os/signal"
	"syscall"
	"time"
)

/*
issue: https://internal.pingcap.net/jira/browse/TIKV-574
pr: https://github.com/pingcap/pd/pull/1073
desc: evict-leader should not transfer leader to stores which reject leaders
note: This case does not support 2.0
steps:
 1. 1 PD, 2 TiKVs
 2. Get regions' distribution
 3. Add scheduler evict-leader to TiKV1
 4. Add scheduler evict-leader to TiKV2
 5. Check region distribution for 10 rounds. No regions transferred from TiKV2 to TiKV1
*/

func main() {
	flag.Parse()

	ctx, cancel := context.WithCancel(context.Background())

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		sig := <-sc
		log.Info("git signal to exit", zap.String("signal", sig.String()))
		cancel()
		os.Exit(0)
	}()

	if !local {
		cli := client.New(ctx, client.ModeOnline)
		pdEndpoint = cli.PdService()
	}
	setup(ctx, fmt.Sprintf("http://%s", pdEndpoint))

	pd := pd_component.PD{}
	stores, err := pd.ListStores()
	if err != nil {
		log.Fatal("List stores failed", zap.Error(err))
	}

	if len(stores) < 2 {
		log.Fatal("No enough stores", zap.Int("current store count", len(stores)))
	}
	rejectStore := stores[0]
	evictStore := stores[1]
	log.Info(fmt.Sprintf("reject store: %v, evict store: %v", rejectStore, evictStore))

	slMap, err := pd.GetStoreRegionLeaderMap()
	if err != nil {
		log.Fatal("Failed to get store leader map", zap.Error(err))
	}
	evictLeaders, ok := slMap[evictStore.Store.ID]
	if !ok {
		log.Fatal("No leaders found on evict store", zap.Uint64("store ID", evictStore.Store.ID))
	}

	if err = rejectStore.AddScheduler("evict-leader-scheduler"); err != nil {
		log.Fatal("Failed to add scheduler evict-leader-scheduler for store", zap.Uint64("store ID", rejectStore.Store.ID), zap.Error(err))
	}
	log.Info("Add scheduler evict-leader-scheduler", zap.Uint64("store ID", rejectStore.Store.ID))

	if err = evictStore.AddScheduler("evict-leader-scheduler"); err != nil {
		log.Fatal("Failed to add scheduler evict-leader-scheduler for store ", zap.Uint64("store ID", evictStore.Store.ID), zap.Error(err))
	}
	log.Info("Add scheduler evict-leader-scheduler", zap.Uint64("store ID", evictStore.Store.ID))

	log.Info("Checking leader location")
	time.Sleep(5 * time.Second)
	for i := 0; i < 10; i++ {
		time.Sleep(1 * time.Second)
		currentSLMap, err := pd.GetStoreRegionLeaderMap()
		if err != nil {
			log.Fatal("Failed to get store leader map", zap.Error(err))
		}
		leadersOnReject, ok := currentSLMap[rejectStore.Store.ID]
		if !ok {
			continue
		}
		for _, leader := range leadersOnReject {
			for _, evict := range evictLeaders {
				if leader.Leader.StoreID == evict.ID {
					log.Fatal(
						"Test Failed",
						zap.Uint64("Region ID", leader.ID),
						zap.Uint64("Original Store ID", evictStore.Store.ID),
						zap.Uint64("Current Store ID", rejectStore.Store.ID),
					)
				}
			}
		}
	}

	log.Info("TearDown: remove schedulers")
	if err = rejectStore.RemoveScheduler("evict-leader-scheduler"); err != nil {
		log.Fatal("Remove store evict-leader-scheduler failed", zap.Uint64("Store ID", rejectStore.Store.ID))
	}
	if err = evictStore.RemoveScheduler("evict-leader-scheduler"); err != nil {
		log.Fatal("Remove store evict-leader-scheduler failed", zap.Uint64("Store ID", evictStore.Store.ID))
	}

	log.Info("Test Passed")
}

var (
	pdTag      string
	pdEndpoint string
	local      bool
)

func init() {
	flag.StringVar(&pdTag, "pd-tag", "master", "PD build tag")
	flag.StringVar(&pdEndpoint, "pd-endpoint", "", "PD Endpoint")
	flag.BoolVar(&local, "local", false, "Run locally")
}

func setup(ctx context.Context, endpoint string) {
	bin := "/tmp/master/bin/pd-ctl"
	if !local {
		dest := "/tmp/" + pdTag
		log.Info("Installing pd-ctl", zap.String("tag", pdTag), zap.String("path", dest))
		binS, err := pd_ctl.InstallPDCtl(ctx, pdTag, dest)
		if err != nil {
			log.Fatal(
				"Install pd ctl-failed",
				zap.String("tag", pdTag),
				zap.String("dest", dest),
				zap.Error(err),
			)
		}
		bin = binS
	}

	pd_ctl.Setup(ctx, bin, endpoint)
}